perf: Improve gap filling, handle missing auth events better

src/admin/debug/commands.rs

@@ -1173,16 +1173,16 @@ impl crate::Context<'_> {
 		earliest: OwnedEventId,
 		via: String,
 	) -> Result {
-		let earliest_pdu = self.services.rooms.timeline.get_pdu(&earliest).await?;
+		let latest_pdu = self.services.rooms.timeline.get_pdu(&latest).await?;
 		let events = self
 			.services
 			.rooms
 			.event_handler
-			.backfill_missing_events(
-				earliest_pdu.room_id_or_hash(),
-				HashSet::from_iter(vec![latest]),
+			.get_missing_events(
+				&latest_pdu.room_id_or_hash(),
+				&latest_pdu,
 				vec![earliest],
-				ServerName::parse(via)?,
+				&ServerName::parse(via)?,
 			)
 			.await?;
 		self.write_str(&format!("Found {} events:\n\n```\n", events.len()))

src/service/rooms/event_handler/fetch_and_handle_outliers.rs

@@ -3,11 +3,16 @@ use std::{
 	time::Instant,
 };
 
+use assign::assign;
 use conduwuit::{
-	Err, Event, PduEvent, debug, debug_info, debug_warn, err,
-	matrix::event::gen_event_id_canonical_json, state_res::lexicographical_topological_sort,
-	trace, utils::continue_exponential_backoff_secs, warn,
+	Err, Event, PduEvent, debug, debug_info, debug_warn, err, error,
+	matrix::event::gen_event_id_canonical_json,
+	state_res::lexicographical_topological_sort,
+	trace,
+	utils::{IterStream, continue_exponential_backoff_secs, stream::BroadbandExt},
+	warn,
 };
+use futures::StreamExt;
 use ruma::{
 	CanonicalJsonObject, CanonicalJsonValue, EventId, MilliSecondsSinceUnixEpoch, OwnedEventId,
 	OwnedRoomId, OwnedServerName, RoomId, ServerName, UInt,
@@ -437,4 +442,128 @@ impl super::Service {
 		trace!("Missing_events: {backfilled_events:?}");
 		Ok(backfilled_events)
 	}
+
+	/// Uses `/_matrix/federation/v1/get_missing_events` to fill gaps in the
+	/// DAG.
+	///
+	/// When this function is called, the "earliest events" (current forward
+	/// extremities) will be collected, and the function will loop with an
+	/// exponentially incrementing limit (up to 100 per request) until it has
+	/// filled the gap, i.e. when the remote says there's no more events.
+	///
+	/// This function will iterate until the remote returns no more events,
+	/// increasing the limit by a factor of 10. If 20 iterations are reached or
+	/// 200 events are backfilled, the function will give up and return what it
+	/// has, to avoid pulling in too much data (for example, absurdly large
+	/// gaps).
+	///
+	/// This function does not persist the events. The caller is responsible for
+	/// passing them through handle_incoming_pdu.
+	///
+	/// ## Parameters
+	///
+	/// - `room_id`: The room's ID.
+	/// - `head`: The event we are potentially missing prev_events for.
+	/// - `tail`: The most recently known events in the graph (typically forward
+	///   extremities).
+	/// - `via`: The server to ask for missing events.
+	pub async fn get_missing_events(
+		&self,
+		room_id: &RoomId,
+		head: &PduEvent,
+		tail: Vec<OwnedEventId>,
+		via: &ServerName,
+	) -> conduwuit::Result<HashMap<OwnedEventId, PduEvent>> {
+		#[cfg(debug_assertions)]
+		{
+			let missing_count = head
+				.prev_events()
+				.stream()
+				.broad_filter_map(|event_id| async move {
+					match self
+						.services
+						.timeline
+						.get_non_outlier_pdu_json(event_id)
+						.await
+						.inspect(|_| debug!("Found prev_event {event_id} locally."))
+						.inspect_err(
+							|e| debug!(%e, "Could not find prev_event {event_id} locally."),
+						) {
+						| Ok(_) => None,
+						| Err(_) => Some(event_id),
+					}
+				})
+				.count()
+				.await;
+			debug_assert_ne!(
+				missing_count, 0,
+				"event passed to get_missing_events is not missing any events (wasteful call)"
+			);
+		};
+
+		let mut discovered = HashMap::with_capacity(20);
+		let mut latest_events = vec![head.event_id().to_owned()];
+		let mut iterations = 0_u8;
+		loop {
+			iterations = iterations.saturating_add(1);
+			let limit = iterations.saturating_mul(10);
+			debug_info!(%limit, %via, %iterations, "Attempting to gap fill missing events");
+			let response: get_missing_events::v1::Response = self
+				.services
+				.sending
+				.send_federation_request(
+					via,
+					assign!(
+						get_missing_events::v1::Request::new(
+							room_id.to_owned(),
+							tail.clone(),
+							latest_events.clone()
+						),
+						{limit: limit.into()}
+					),
+				)
+				.await?;
+
+			if response.events.is_empty() {
+				debug_info!(%via, "Finished gap filling missing events (remote returned no more events).");
+				break;
+			}
+			debug_info!("Got {} events back from remote", response.events.len());
+
+			latest_events.clear();
+			for raw_event in response.events {
+				let (_, event_id, pdu_json) = self.parse_incoming_pdu(&raw_event).await?;
+				let pdu = PduEvent::from_id_val(&event_id, pdu_json).map_err(|e| {
+					err!(Request(BadJson("Failed to parse backfilled event {event_id}: {e}")))
+				})?;
+
+				for prev_event_id in pdu.prev_events() {
+					if discovered.contains_key(prev_event_id) {
+						continue;
+					}
+					if self
+						.services
+						.timeline
+						.non_outlier_pdu_exists(prev_event_id)
+						.await
+					{
+						continue;
+					}
+					latest_events.push(prev_event_id.to_owned());
+					break;
+				}
+
+				discovered.insert(event_id.clone(), pdu);
+			}
+
+			if latest_events.is_empty() {
+				break;
+			} else if discovered.len() >= 200 || iterations >= 20 {
+				error!("Gap too large, giving up");
+				break;
+			}
+		}
+
+		Ok(discovered)
+	}
 }

src/service/rooms/event_handler/fetch_prev.rs

@@ -1,9 +1,9 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 
 use conduwuit::{
 	Event, PduEvent, debug, debug_info,
-	result::DebugInspect,
 	utils::{BoolExt, IterStream, stream::BroadbandExt},
+	warn,
 };
 use futures::StreamExt;
 use ruma::{RoomId, ServerName};
@@ -28,8 +28,6 @@ impl super::Service {
 					.timeline
 					.get_non_outlier_pdu_json(event_id)
 					.await
-					.inspect(|_| debug!("Found prev_event {event_id} locally."))
-					.inspect_err(|e| debug!(%e, "Could not find prev_event {event_id} locally."))
 					.is_ok()
 					.or(|| event_id.to_owned())
 			})
@@ -48,12 +46,7 @@ impl super::Service {
 			.await;
 
 		let backfilled = self
-			.backfill_missing_events(
-				room_id.to_owned(),
-				HashSet::from_iter(vec![incoming_pdu.event_id().to_owned()]),
-				tail,
-				origin.to_owned(),
-			)
+			.get_missing_events(room_id, incoming_pdu, tail, origin)
 			.await?;
 		debug_info!("Fetched {} missing events", backfilled.len());
 
@@ -75,13 +68,19 @@ impl super::Service {
 
 		for event_id in to_persist {
 			debug_info!("Persisting fetched prev event {event_id}");
-			let pdu_event = backfilled.get(&event_id).cloned().unwrap_or_else(|| {
-				panic!("Event {event_id} was in backfill response but not in map")
-			});
-			let obj = pdu_event.to_canonical_object();
-			self.upgrade_outlier_to_timeline_pdu(pdu_event, obj, create_event, origin, room_id)
+			let obj = mapped.get(&event_id).cloned().unwrap();
+			match self
+				.handle_outlier_pdu(origin, create_event, &event_id, room_id, obj, false)
 				.await
-				.debug_inspect(|_| debug_info!("Persisted fetched prev event {event_id}"))?;
+			{
+				| Ok((pdu, val)) =>
+					self.upgrade_outlier_to_timeline_pdu(pdu, val, create_event, origin, room_id)
+						.await,
+				| Err(e) => {
+					warn!("Failed to persist prev_event {event_id}: {e}");
+					continue;
+				},
+			}?;
 		}
 
 		// NOTE because i keep forgetting: the caller persists incoming_pdu.

src/service/rooms/event_handler/handle_outlier_pdu.rs

@@ -1,15 +1,16 @@
-use std::collections::{BTreeMap, HashMap, HashSet, hash_map};
+use std::collections::{BTreeMap, HashMap, hash_map};
 
 use conduwuit::{
-	Err, Event, PduEvent, Result, debug, debug_info, debug_warn, err, state_res, trace, warn,
+	Err, Event, PduEvent, Result, debug, debug_info, debug_warn, err, info, state_res, trace,
+	warn,
 };
-use futures::{StreamExt, future::ready};
+use futures::future::ready;
 use ruma::{
 	CanonicalJsonObject, CanonicalJsonValue, EventId, OwnedEventId, RoomId, ServerName,
-	events::StateEventType,
+	api::federation::authorization::get_event_authorization, events::StateEventType,
 };
 
-use super::{build_local_dag, check_room_id, get_room_version_rules};
+use super::{check_room_id, get_room_version_rules};
 use crate::rooms::timeline::pdu_fits;
 
 impl super::Service {
@@ -112,79 +113,52 @@ impl super::Service {
 		}
 
 		// Fetch any missing ones & reject invalid ones
-		let mut missing_auth_events: HashSet<OwnedEventId> = pdu_event
-			.auth_events()
-			.filter(|id| !auth_events.contains_key(*id))
-			.map(ToOwned::to_owned)
-			.collect();
-
-		if !missing_auth_events.is_empty() {
-			debug_info!(
-				"Fetching {} missing auth events for outlier event {event_id}",
-				missing_auth_events.len()
-			);
-			let tail = self
+		if auth_events.len() != pdu_event.auth_events().count() {
+			info!("Missing some auth events, asking remote for auth chain");
+			let response: get_event_authorization::v1::Response = self
 				.services
-				.state
-				.get_forward_extremities(room_id)
-				.collect::<Vec<_>>()
-				.await;
-			let backfilled = self
-				.backfill_missing_events(
-					room_id.to_owned(),
-					HashSet::from_iter(vec![event_id.to_owned()]),
-					tail,
-					origin.to_owned(),
-				)
-				.await?;
-			debug_info!("Fetched {} missing auth events for {event_id}", backfilled.len());
-			let mapped = backfilled
-				.iter()
-				.map(|(eid, evt)| {
-					let mut obj = evt.to_canonical_object();
-					obj.remove("event_id"); // event_id is inserted by backfill_missing_events
-					(eid.clone(), obj)
-				})
-				.collect::<HashMap<_, _>>();
-			let local_dag = if mapped.len() == 1 {
-				mapped.keys().map(ToOwned::to_owned).collect()
-			} else {
-				build_local_dag(&mapped).await?
-			};
-			debug_info!("Preparing to handle {} missing auth events", backfilled.len());
-			for prev_event_id in local_dag {
-				let obj = mapped
-					.get(&prev_event_id)
-					.expect("We should have this event in memory");
-				debug_info!("Handling prev {prev_event_id}");
-				let (prev, _) = Box::pin(self.handle_outlier_pdu(
+				.sending
+				.send_federation_request(
 					origin,
-					create_event,
-					&prev_event_id,
-					room_id,
-					obj.clone(),
-					false,
-				))
-				.await?;
-				if missing_auth_events.contains(&*prev_event_id) {
-					missing_auth_events.remove(&prev_event_id);
-					auth_events.insert(prev_event_id, prev);
+					get_event_authorization::v1::Request::new(
+						room_id.to_owned(),
+						event_id.to_owned(),
+					),
+				)
+				.await
+				.map_err(|e| {
+					err!(Request(Forbidden(
+						"Remote server is not divulging incoming event's auth chain: {e}"
+					)))
+				})?;
+			let mut auth_chain_map = HashMap::with_capacity(response.auth_chain.len());
+			for auth_pdu_json in response.auth_chain {
+				let (auth_event_room_id, auth_event_id, auth_pdu_json) =
+					self.parse_incoming_pdu(&auth_pdu_json).await?;
+				if auth_event_room_id != room_id {
+					return Err!(Request(BadJson(
+						"Auth event {auth_event_id} is in {auth_event_room_id}, not {room_id}."
+					)));
 				}
-				debug_info!("Finished handling prev auth event");
+				let auth_pdu = PduEvent::from_id_val(&auth_event_id, auth_pdu_json)
+					.map_err(|e| err!(Request(BadJson("Invalid PDU {auth_event_id}: {e}"))))?;
+				auth_chain_map.insert(auth_event_id, auth_pdu);
 			}
-		} else {
-			debug!("No missing auth events for outlier event {event_id}");
-		}
-		// reject if we are still missing some auth events.
-		// If we're still missing prev events, we will fetch them individually later,
-		// but there's no reason for us to be missing auth events now we've gapfilled
-		// the DAG.
-		if !missing_auth_events.is_empty() {
-			// Don't reject: this could be a temporary condition
-			return Err!(Request(InvalidParam(
-				"Could not fetch all auth events for outlier event {event_id}, still missing: \
-				 {missing_auth_events:?}"
-			)));
+			for aid in pdu_event.auth_events() {
+				if auth_events.contains_key(aid) {
+					continue;
+				}
+				if let Some(auth_event) = auth_chain_map.get(aid) {
+					auth_events.insert(aid.to_owned(), auth_event.clone());
+				} else {
+					return Err!(Request(Forbidden(
+						"Remote server is not divulging incoming event's auth events (missing: \
+						 {aid})"
+					)));
+				}
+			}
+			// TODO: do events received from auth chain need persisting? that
+			// sounds awfully slow
 		}
 
 		// 6. Reject "due to auth events" if the event doesn't pass auth based on the