fix(federation): Queue PDUs during remote joins

changelog.d/2063.bugfix.md

@@ -0,0 +1 @@
+Fixed federation PDUs received during remote joins being dropped before their room state was available. Contributed by @eleboucher.

conduwuit-example.toml

@@ -449,6 +449,10 @@
 #
 #max_concurrent_inbound_transactions = 150
 
+# Maximum number of federation PDUs queued while completing a remote join.
+#
+#remote_join_pdu_queue_capacity = 50
+
 # Maximum age (in seconds) for cached federation transaction responses.
 # Entries older than this will be removed during cleanup.
 #

src/api/server/send.rs

@@ -9,7 +9,7 @@ use conduwuit::{
 	result::LogErr,
 	trace,
 	utils::{
-		IterStream, ReadyExt, millis_since_unix_epoch,
+		IterStream, millis_since_unix_epoch,
 		stream::{BroadbandExt, TryBroadbandExt, automatic_width},
 	},
 	warn,
@@ -142,13 +142,26 @@ async fn process_inbound_transaction(
 	sender: Sender<WrappedTransactionResponse>,
 ) {
 	let txn_start_time = Instant::now();
-	let pdus = body
-		.pdus
-		.iter()
-		.stream()
-		.broad_then(|pdu| services.rooms.event_handler.parse_incoming_pdu(pdu, None))
-		.inspect_err(|e| warn!("Could not parse incoming PDU: {e}"))
-		.ready_filter_map(Result::ok);
+	let pdus = body.pdus.iter().stream().broad_filter_map(|pdu| async {
+		match services
+			.rooms
+			.event_handler
+			.parse_incoming_pdu(pdu, None)
+			.await
+		{
+			| Ok(pdu) => Some(pdu),
+			| Err(error) => {
+				if !services
+					.rooms
+					.event_handler
+					.queue_pending_join_pdu(&body.identity, pdu.clone())
+				{
+					warn!("Could not parse incoming PDU: {error}");
+				}
+				None
+			},
+		}
+	});
 
 	let edus = body
 		.edus

src/core/config/check.rs

@@ -115,6 +115,13 @@ pub fn check(config: &Config) -> Result {
 		));
 	}
 
+	if config.remote_join_pdu_queue_capacity == 0 {
+		return Err!(Config(
+			"remote_join_pdu_queue_capacity",
+			"The remote join PDU queue capacity must be greater than zero"
+		));
+	}
+
 	// yeah, unless the user built a debug build hopefully for local testing only
 	if cfg!(not(debug_assertions)) && config.server_name == "your.server.name" {
 		return Err!(Config(

src/core/config/mod.rs

@@ -555,6 +555,12 @@ pub struct Config {
 	#[serde(default = "default_max_concurrent_inbound_transactions")]
 	pub max_concurrent_inbound_transactions: usize,
 
+	/// Maximum number of federation PDUs queued while completing a remote join.
+	///
+	/// default: 50
+	#[serde(default = "default_remote_join_pdu_queue_capacity")]
+	pub remote_join_pdu_queue_capacity: usize,
+
 	/// Maximum age (in seconds) for cached federation transaction responses.
 	/// Entries older than this will be removed during cleanup.
 	///
@@ -3006,6 +3012,8 @@ fn default_max_fetch_prev_events() -> u16 { 1024 }
 
 fn default_max_concurrent_inbound_transactions() -> usize { 150 }
 
+fn default_remote_join_pdu_queue_capacity() -> usize { 50 }
+
 fn default_transaction_id_cache_max_age_secs() -> u64 { 60 * 60 * 2 }
 
 fn default_transaction_id_cache_max_entries() -> usize { 8192 }

src/service/rooms/event_handler/mod.rs

@@ -7,6 +7,7 @@ mod handle_incoming_pdu;
 mod handle_outlier_pdu;
 mod parse_incoming_pdu;
 pub mod pdu_checks;
+mod pending_join_pdus;
 mod policy_server;
 mod resolve_state;
 mod state_at_incoming;
@@ -20,9 +21,10 @@ pub use fetch_and_handle_outliers::{
 	DagBuilderTree, GET_MISSING_EVENTS_MAX_BATCH_SIZE, build_local_dag,
 };
 use ruma::{
-	OwnedEventId, OwnedRoomId, events::room::create::RoomCreateEventContent,
+	OwnedEventId, OwnedRoomId, OwnedServerName, events::room::create::RoomCreateEventContent,
 	room_version_rules::RoomVersionRules,
 };
+use serde_json::value::RawValue as RawJsonValue;
 use tokio::sync::{Notify, mpsc};
 
 use crate::{Dep, globals, rooms, sending, server_keys};
@@ -30,6 +32,7 @@ pub struct Service {
 	pub mutex_federation: RoomMutexMap,
 	pub federation_handletime: SyncRwLock<HandleTimeMap>,
 	pub extremity_squashers: SyncRwLock<HashMap<OwnedRoomId, mpsc::Sender<(usize, bool)>>>,
+	joining_rooms: SyncRwLock<HashMap<OwnedRoomId, mpsc::Sender<PendingJoinPdu>>>,
 	services: Services,
 	server_shutdown: Notify,
 	me: std::sync::Weak<Self>,
@@ -55,6 +58,11 @@ struct Services {
 type RoomMutexMap = MutexMap<OwnedRoomId, ()>;
 type HandleTimeMap = HashMap<OwnedRoomId, (OwnedEventId, Instant)>;
 
+enum PendingJoinPdu {
+	Pdu(OwnedServerName, Box<RawJsonValue>),
+	Complete,
+}
+
 #[async_trait]
 impl crate::Service for Service {
 	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
@@ -63,6 +71,7 @@ impl crate::Service for Service {
 			mutex_federation: RoomMutexMap::new(),
 			federation_handletime: HandleTimeMap::new().into(),
 			extremity_squashers: SyncRwLock::new(HashMap::new()),
+			joining_rooms: HashMap::new().into(),
 			services: Services {
 				globals: args.depend::<globals::Service>("globals"),
 				sending: args.depend::<sending::Service>("sending"),

src/service/rooms/event_handler/pending_join_pdus.rs

@@ -0,0 +1,82 @@
+use conduwuit::{debug, debug_warn};
+use ruma::{CanonicalJsonObject, RoomId, ServerName};
+use serde_json::value::RawValue as RawJsonValue;
+
+impl super::Service {
+	pub fn begin_remote_join(&self, room_id: &RoomId) {
+		let (tx, mut rx) = tokio::sync::mpsc::channel(
+			self.services.server.config.remote_join_pdu_queue_capacity,
+		);
+		self.joining_rooms.write().insert(room_id.to_owned(), tx);
+		let Some(service) = self.me.upgrade() else {
+			return;
+		};
+		let room_id = room_id.to_owned();
+
+		self.services.server.runtime().spawn(async move {
+			let mut pdus = Vec::new();
+			let mut complete = false;
+			while let Some(pdu) = rx.recv().await {
+				match pdu {
+					| super::PendingJoinPdu::Pdu(origin, pdu) => pdus.push((origin, pdu)),
+					| super::PendingJoinPdu::Complete => {
+						complete = true;
+						break;
+					},
+				}
+			}
+			if !complete {
+				return;
+			}
+			let _room_lock = service.mutex_federation.lock(room_id.as_str()).await;
+
+			for (origin, pdu) in pdus {
+				let Ok((room_id, event_id, value)) = service.parse_incoming_pdu(&pdu, None).await
+				else {
+					debug_warn!("Failed to parse PDU queued during remote join");
+					continue;
+				};
+				if let Err(error) = service
+					.handle_incoming_pdu(&origin, &room_id, &event_id, value, false)
+					.await
+				{
+					debug_warn!(%room_id, %event_id, "Failed to process PDU queued during remote join: {error}");
+				}
+			}
+		});
+	}
+
+	pub fn cancel_remote_join(&self, room_id: &RoomId) {
+		self.joining_rooms.write().remove(room_id);
+	}
+
+	pub fn queue_pending_join_pdu(&self, origin: &ServerName, pdu: Box<RawJsonValue>) -> bool {
+		let Ok(value) = serde_json::from_str::<CanonicalJsonObject>(pdu.get()) else {
+			return false;
+		};
+		let Some(room_id) = value.get("room_id").and_then(|value| value.as_str()) else {
+			return false;
+		};
+		let Ok(room_id) = RoomId::parse(room_id) else {
+			return false;
+		};
+		let joining_rooms = self.joining_rooms.read();
+		let Some(tx) = joining_rooms.get(&room_id) else {
+			return false;
+		};
+		let pdu = super::PendingJoinPdu::Pdu(origin.to_owned(), pdu);
+		if tx.try_send(pdu).is_err() {
+			debug_warn!(%room_id, "Dropping PDU received during remote join because the queue is full");
+		} else {
+			debug!(%room_id, "Queued PDU received during remote join");
+		}
+		true
+	}
+
+	pub async fn process_pending_join_pdus(&self, room_id: &RoomId) {
+		let Some(tx) = self.joining_rooms.write().remove(room_id) else {
+			return;
+		};
+		let _ = tx.send(super::PendingJoinPdu::Complete).await;
+	}
+}

src/service/rooms/membership/mod.rs

@@ -1,7 +1,8 @@
 use std::{collections::HashMap, sync::Arc};
 
 use conduwuit::{
-	Err, Event, Pdu, Result, Server, debug, debug_info, debug_warn, err, error, info, is_true,
+	Err, Event, Pdu, Result, Server, debug, debug_info, debug_warn, defer, err, error, info,
+	is_true,
 	matrix::{
 		StateKey,
 		event::{gen_event_id, gen_event_id_canonical_json},
@@ -404,6 +405,8 @@ impl Service {
 			.mutex_federation
 			.lock(room_id.as_str())
 			.await;
+		self.services.event_handler.begin_remote_join(room_id);
+		defer! {{ self.services.event_handler.cancel_remote_join(room_id); }}
 		info!("Asking {remote_server} for send_join in room {room_id}");
 		let send_join_response = match self
 			.services
@@ -675,6 +678,10 @@ impl Service {
 		drop(cork);
 
 		self.services.sync.wake_all_joined(room_id).await;
+		self.services
+			.event_handler
+			.process_pending_join_pdus(room_id)
+			.await;
 
 		Ok(())
 	}