fix: Add startup migration to re-process existing extremities on-disk

src/service/migrations.rs

@@ -8,12 +8,12 @@ use std::{
 };
 
 use conduwuit::{
-	Err, Error, Event, Pdu, Result, debug, debug_info, debug_warn, err, error, info,
+	Err, Error, Event, Pdu, PduCount, Result, debug, debug_info, debug_warn, err, error, info,
 	result::NotFound,
 	trace,
 	utils::{
 		IterStream, ReadyExt,
-		stream::{TryExpect, TryIgnore},
+		stream::{TryExpect, TryIgnore, automatic_width},
 		to_canonical_object,
 	},
 	warn,
@@ -167,6 +167,11 @@ migration!(stringify!($name) => $name)
 		// names, preventing users from signing in.
 		VersionBump(22),
 		migration!(fix_oauth_scope_names),
+		// Removes any extremities in roomid_pduleaves that we do not know about (and
+		// consequently should not reference).
+		migration!(remove_unknown_extremities),
+		// Re-calculates extremities for any rooms that actually lack any extremities at all.
+		migration!(consistency_check_extremity_tracker),
 	]
 };
 
@@ -979,3 +984,69 @@ async fn fix_oauth_scope_names(services: &Services) -> Result {
 
 	Ok(())
 }
+
+async fn remove_unknown_extremities(services: &Services) -> Result {
+	services.db["roomid_pduleaves"]
+		.keys()
+		.ignore_err()
+		.for_each_concurrent(
+			automatic_width(),
+			|(room_id, event_id): (OwnedRoomId, OwnedEventId)| async move {
+				if !services.rooms.timeline.pdu_exists(&event_id).await {
+					debug!("Removing unknown extremity in {room_id}: {event_id}");
+					services
+						.rooms
+						.state
+						.remove_forward_extremity(&room_id, &event_id);
+				}
+			},
+		)
+		.await;
+	Ok(())
+}
+
+async fn consistency_check_extremity_tracker(services: &Services) -> Result {
+	let mut rooms = services.rooms.metadata.iter_ids();
+
+	// For every room the server is a member of,
+	while let Some(room_id) = rooms.next().await {
+		// If this room has no extremities, re-calculate
+		if services
+			.rooms
+			.state
+			.get_forward_extremities(&room_id)
+			.count()
+			.await
+			>= 1
+		{
+			continue;
+		}
+
+		warn!(
+			"Room {room_id} is missing extremity information, recalculating. This may take a \
+			 long time if this room is old."
+		);
+		services
+			.rooms
+			.timeline
+			.pdus(&room_id, Some(PduCount::Normal(0)))
+			.ignore_err()
+			.for_each(|(_, pdu)| {
+				let room_id = &room_id;
+				async move {
+					services
+						.rooms
+						.timeline
+						.progress_extremities(&pdu)
+						.await
+						.inspect_err(|e| {
+							error!("Failed to progress extremities in {room_id}: {e:?}");
+						})
+						.ok();
+				}
+			})
+			.await;
+	}
+
+	Ok(())
+}