feat: Enhance reliability by fetching full state when we're missing a lot of auth events

src/service/rooms/event_handler/fetch_and_handle_outliers.rs

@@ -338,8 +338,10 @@ pub(super) async fn fetch_and_handle_missing_events<'a, Pdu>(
 			.candidate_backfill_servers(room_id)
 			.await;
 		candidates.insert(origin.to_owned());
+		candidates.retain(|sn| self.services.globals.server_name() != sn);
+		assert_ne!(candidates.len(), 0, "no candidates to fetch missing events from");
 		let mut seeded_events =
-			HashMap::with_capacity(events.len() + (events.len().saturating_mul(3)));
+			HashMap::with_capacity(events.len().saturating_add(events.len().saturating_mul(3)));
 		trace!(
 			"Fetching {} unknown PDUs on demand from {} candidates",
 			events.len(),
@@ -392,8 +394,32 @@ pub(super) async fn fetch_and_handle_missing_events<'a, Pdu>(
 							continue;
 						},
 					};
-				todo.extend(auth_events);
-				seeded_events.insert(event_id, value);
+				let mut have_all_auth = true;
+				for auth_event_id in auth_events {
+					if let Ok(local_pdu) = self.services.timeline.get_pdu(&next_id).await {
+						trace!("Found auth event {next_id} in db");
+						seeded_events.insert(id.clone(), local_pdu.into_canonical_object());
+						continue;
+					}
+					if seeded_events.contains_key(&auth_event_id) {
+						trace!(%auth_event_id, "Already found auth event");
+						continue;
+					}
+					todo.push_back(auth_event_id);
+					have_all_auth = false;
+				}
+				// Insert this PDU back at the end of the queue so that it will be resolved once
+				// all of its auth events have been fetched.
+				// TODO: This may result in infinite looping, needs a breaker
+				if have_all_auth {
+					debug!(%event_id, "Have all auth events");
+					seeded_events.insert(event_id, value);
+				} else {
+					debug_warn!(
+						"Fetched {event_id} but missing some auth events, will have to re-fetch."
+					);
+					todo.push_back(event_id);
+				}
 			}
 		}
 
@@ -408,7 +434,7 @@ pub(super) async fn fetch_and_handle_missing_events<'a, Pdu>(
 		let mut pdus = HashMap::with_capacity(seeded_ordered.len());
 		for id in seeded_ordered {
 			let pdu_json = seeded_events.remove(&id).unwrap();
-			trace!("Handling missing event {id}");
+			debug_info!("Handling missing event {id} as outlier");
 			match Box::pin(self.handle_outlier_pdu(
 				origin,
 				create_event,
@@ -424,6 +450,10 @@ pub(super) async fn fetch_and_handle_missing_events<'a, Pdu>(
 				},
 				| Err(e) => warn!("Authentication of event {id} failed: {e:?}"),
 			}
+
+			// TODO: should this try to promote to timeline?
+			// If we got here, we probably weren't able to promote it before
+			// now.
 		}
 
 		trace!("Fetched and handled {} missing PDUs", pdus.len());

src/service/rooms/event_handler/fetch_state.rs

@@ -1,111 +1,227 @@
 use std::collections::{HashMap, hash_map};
 
-use conduwuit::{
-	Err, Event, PduEvent, Result, debug, debug_warn, err, implement, utils::IterStream,
-};
+use conduwuit::{Err, Event, PduEvent, Result, debug, debug_warn, err, utils::IterStream, warn};
 use futures::StreamExt;
 use ruma::{
-	EventId, OwnedEventId, RoomId, ServerName, api::federation::event::get_room_state_ids,
+	EventId, OwnedEventId, RoomId, ServerName,
+	api::federation::event::{get_room_state, get_room_state_ids},
 	events::StateEventType,
 };
 
 use crate::{conduwuit::utils::stream::BroadbandExt, rooms::short::ShortStateKey};
 
-/// Call /state_ids to find out what the state at this pdu is. We trust the
-/// server's response to some extend (sic), but we still do a lot of checks
-/// on the events
-#[implement(super::Service)]
-#[tracing::instrument(
-	level = "debug",
-	skip_all,
-	fields(%origin),
-)]
-pub(super) async fn fetch_state<Pdu>(
-	&self,
-	origin: &ServerName,
-	create_event: &Pdu,
-	room_id: &RoomId,
-	event_id: &EventId,
-) -> Result<Option<HashMap<u64, OwnedEventId>>>
-where
-	Pdu: Event + Send + Sync,
-{
-	let res: get_room_state_ids::v1::Response = self
-		.services
-		.sending
-		.send_federation_request(
-			origin,
-			get_room_state_ids::v1::Request::new(event_id.to_owned(), room_id.to_owned()),
-		)
-		.await
-		.inspect_err(|e| debug_warn!("Fetching state for event failed: {e}"))?;
+impl super::Service {
+	/// Asks a remote server what the state at this event is.
+	/// It first attempts to call `GET /_matrix/federation/v1/state_ids` (fast).
+	/// If any events are missing, they are fetched from the remote, and
+	/// persisted as outliers, before being returned back to this function. If
+	/// we are missing a lot of events locally (>=50), this function falls back
+	/// to requesting the full state in PDU format from the remote (`GET
+	/// /_matrix/federation/v1/state, very slow in large rooms), and persists
+	/// them directly.
+	///
+	/// The end result is a result containing a map of shortstatekeys to event
+	/// IDs. The underlying `Option` is always `Some`.
+	#[tracing::instrument(skip(self, create_event))]
+	pub(super) async fn fetch_state(
+		&self,
+		origin: &ServerName,
+		create_event: &PduEvent,
+		room_id: &RoomId,
+		event_id: &EventId,
+	) -> Result<Option<HashMap<u64, OwnedEventId>>> {
+		let res: get_room_state_ids::v1::Response = self
+			.services
+			.sending
+			.send_federation_request(
+				origin,
+				get_room_state_ids::v1::Request::new(event_id.to_owned(), room_id.to_owned()),
+			)
+			.await
+			.inspect_err(|e| debug_warn!("Fetching state for event failed: {e}"))?;
 
-	debug!("Fetching state events");
-	let mut state_events: HashMap<OwnedEventId, PduEvent> =
-		HashMap::with_capacity(res.pdu_ids.len());
-	let to_fetch: Vec<OwnedEventId> = res
-		.pdu_ids
-		.clone()
-		.into_iter()
-		.stream()
-		.broad_filter_map(|event_id| async move {
-			if self.services.timeline.pdu_exists(&event_id).await {
-				None
+		debug!("Fetching state events");
+		let mut state_events: HashMap<OwnedEventId, PduEvent> =
+			HashMap::with_capacity(res.pdu_ids.len());
+		let to_fetch: Vec<OwnedEventId> = res
+			.pdu_ids
+			.clone()
+			.into_iter()
+			.stream()
+			.broad_filter_map(|event_id| async move {
+				if self.services.timeline.pdu_exists(&event_id).await {
+					None
+				} else {
+					Some(event_id)
+				}
+			})
+			.collect()
+			.await;
+		if !to_fetch.is_empty() {
+			if to_fetch.len() >= 100 {
+				// That's a lot of events to fetch, just ask for the full state
+				// at that point.
+				debug_warn!(
+					to_fetch = to_fetch.len(),
+					"Fetching full state from remote server for event"
+				);
+				state_events.extend(
+					self.fetch_full_state(origin, create_event, room_id, event_id)
+						.await?,
+				);
 			} else {
-				Some(event_id)
+				debug!(
+					to_fetch = to_fetch.len(),
+					"Fetching missing events for state from remote"
+				);
+				state_events.extend(
+					self.fetch_and_handle_missing_events(origin, to_fetch, create_event, room_id)
+						.await,
+				);
 			}
-		})
-		.collect()
-		.await;
-	if !to_fetch.is_empty() {
-		if to_fetch.len() >= 100 {
-			// That's a lot of events to fetch, just ask for the full state at
-			// that point. TODO: fetch /state
 		}
-		state_events.extend(
-			self.fetch_and_handle_missing_events(origin, to_fetch, create_event, room_id)
-				.await,
-		);
-	}
 
-	let mut state: HashMap<ShortStateKey, OwnedEventId> =
-		HashMap::with_capacity(state_events.len());
-	for (event_id, pdu) in state_events {
-		let state_key = pdu
-			.state_key()
-			.ok_or_else(|| err!(Database("Found non-state pdu in state events: {event_id}")))?;
+		let mut state: HashMap<ShortStateKey, OwnedEventId> =
+			HashMap::with_capacity(state_events.len());
+		for (event_id, pdu) in state_events {
+			let state_key = pdu.state_key().ok_or_else(|| {
+				err!(Database("Found non-state pdu in state events: {event_id}"))
+			})?;
 
-		let shortstatekey = self
+			let shortstatekey = self
+				.services
+				.short
+				.get_or_create_shortstatekey(&pdu.kind().to_string().into(), state_key)
+				.await;
+
+			match state.entry(shortstatekey) {
+				| hash_map::Entry::Vacant(v) => {
+					v.insert(pdu.event_id().to_owned());
+				},
+				| hash_map::Entry::Occupied(_) => {
+					return Err!(Database(
+						"State event's type and state_key combination exists multiple times \
+						 ({event_id}): {}, {}",
+						pdu.kind(),
+						state_key
+					));
+				},
+			}
+		}
+
+		// The original create event must still be in the state
+		let create_shortstatekey = self
 			.services
 			.short
-			.get_or_create_shortstatekey(&pdu.kind().to_string().into(), state_key)
-			.await;
+			.get_shortstatekey(&StateEventType::RoomCreate, "")
+			.await?;
 
-		match state.entry(shortstatekey) {
-			| hash_map::Entry::Vacant(v) => {
-				v.insert(pdu.event_id().to_owned());
-			},
-			| hash_map::Entry::Occupied(_) => {
-				return Err!(Database(
-					"State event's type and state_key combination exists multiple times \
-					 ({event_id}): {}, {}",
-					pdu.kind(),
-					state_key
-				));
-			},
+		if state.get(&create_shortstatekey).map(AsRef::as_ref) != Some(create_event.event_id()) {
+			return Err!(Database("Incoming event refers to wrong create event."));
 		}
+
+		Ok(Some(state))
 	}
 
-	// The original create event must still be in the state
-	let create_shortstatekey = self
-		.services
-		.short
-		.get_shortstatekey(&StateEventType::RoomCreate, "")
-		.await?;
-
-	if state.get(&create_shortstatekey).map(AsRef::as_ref) != Some(create_event.event_id()) {
-		return Err!(Database("Incoming event refers to wrong create event."));
+	/// Fetches the full state via `GET /_matrix/federation/v1/state` from a
+	/// remote server, and persists all the incoming auth chain events and
+	/// state events as outliers, for use later.
+	///
+	/// Any events that cannot be persisted are dropped with a warning.
+	/// TODO: make it noisy?
+	pub(super) async fn fetch_full_state(
+		&self,
+		origin: &ServerName,
+		create_event: &PduEvent,
+		room_id: &RoomId,
+		event_id: &EventId,
+	) -> Result<HashMap<OwnedEventId, PduEvent>> {
+		let res: get_room_state::v1::Response = self
+			.services
+			.sending
+			.send_federation_request(
+				origin,
+				get_room_state::v1::Request::new(event_id.to_owned(), room_id.to_owned()),
+			)
+			.await
+			.inspect_err(|e| debug_warn!("Fetching state for event failed: {e}"))?;
+		debug!(count = res.auth_chain.len(), "Handling incoming auth chain...");
+		res.auth_chain
+			.iter()
+			.stream()
+			.broad_filter_map(|raw_event_json| async {
+				if let Some(parsed) = self.parse_incoming_pdu(raw_event_json).await.ok()
+					&& parsed.0 == room_id
+				{
+					Some(parsed)
+				} else {
+					None
+				}
+			})
+			.for_each_concurrent(
+				None,
+				|(incoming_room_id, incoming_event_id, incoming_event_json)| async move {
+					self.handle_outlier_pdu(
+						origin,
+						create_event,
+						&incoming_event_id,
+						&incoming_room_id,
+						incoming_event_json,
+						true,
+					)
+					.await
+					.inspect_err(|e| {
+						warn!(
+							%incoming_room_id,
+							%incoming_event_id,
+							?e,
+							"Failed to handle auth chain event from state fetch"
+						);
+					})
+					.ok();
+				},
+			)
+			.await;
+		debug!(count = res.pdus.len(), "Handling incoming state PDUs...");
+		Ok(res
+			.pdus
+			.iter()
+			.stream()
+			.broad_filter_map(|raw_event_json| async {
+				if let Some(parsed) = self.parse_incoming_pdu(raw_event_json).await.ok()
+					&& parsed.0 == room_id
+				{
+					Some(parsed)
+				} else {
+					None
+				}
+			})
+			.broad_filter_map(
+				|(incoming_room_id, incoming_event_id, incoming_event_json)| async move {
+					self.handle_outlier_pdu(
+						origin,
+						create_event,
+						&incoming_event_id,
+						&incoming_room_id,
+						incoming_event_json,
+						true,
+					)
+					.await
+					.inspect_err(|e| {
+						warn!(
+							%incoming_room_id,
+							%incoming_event_id,
+							?e,
+							"Failed to handle state event from state fetch"
+						);
+					})
+					.ok()
+				},
+			)
+			.fold(HashMap::new(), |mut acc, (event, _)| async move {
+				acc.insert(event.event_id().to_owned(), event);
+				acc
+			})
+			.await)
 	}
-
-	Ok(Some(state))
 }

src/service/rooms/event_handler/upgrade_outlier_pdu.rs

@@ -23,17 +23,14 @@
 };
 
 #[implement(super::Service)]
-pub(super) async fn upgrade_outlier_to_timeline_pdu<Pdu>(
+pub(super) async fn upgrade_outlier_to_timeline_pdu(
 	&self,
 	incoming_pdu: PduEvent,
 	mut val: CanonicalJsonObject,
-	create_event: &Pdu,
+	create_event: &PduEvent,
 	origin: &ServerName,
 	room_id: &RoomId,
-) -> Result<Option<RawPduId>>
-where
-	Pdu: Event + Send + Sync,
-{
+) -> Result<Option<RawPduId>> {
 	// Skip the PDU if we already have it as a timeline event
 	if let Ok(pduid) = self
 		.services