shard sender into multiple task workers by destination hash

rename Destination::Normal variant

tracing instruments

Signed-off-by: Jason Volk <jason@zemos.net>

src/service/sending/sender.rs

@@ -1,7 +1,10 @@
 use std::{
 	collections::{BTreeMap, HashMap, HashSet},
 	fmt::Debug,
-	sync::atomic::{AtomicU64, AtomicUsize, Ordering},
+	sync::{
+		atomic::{AtomicU64, AtomicUsize, Ordering},
+		Arc,
+	},
 	time::{Duration, Instant},
 };
 
@@ -66,29 +69,56 @@ pub const PDU_LIMIT: usize = 50;
 pub const EDU_LIMIT: usize = 100;
 
 impl Service {
-	#[tracing::instrument(skip_all, level = "debug")]
-	pub(super) async fn sender(&self) -> Result<()> {
+	#[tracing::instrument(skip(self), level = "debug")]
+	pub(super) async fn sender(self: Arc<Self>, id: usize) -> Result {
 		let mut statuses: CurTransactionStatus = CurTransactionStatus::new();
 		let mut futures: SendingFutures<'_> = FuturesUnordered::new();
-		let receiver = self.receiver.lock().await;
 
-		self.initial_requests(&mut futures, &mut statuses).await;
-		while !receiver.is_closed() {
-			tokio::select! {
-				request = receiver.recv_async() => match request {
-					Ok(request) => self.handle_request(request, &mut futures, &mut statuses).await,
-					Err(_) => break,
-				},
-				Some(response) = futures.next() => {
-					self.handle_response(response, &mut futures, &mut statuses).await;
-				},
-			}
-		}
-		self.finish_responses(&mut futures).await;
+		self.startup_netburst(id, &mut futures, &mut statuses)
+			.boxed()
+			.await;
+
+		self.work_loop(id, &mut futures, &mut statuses).await;
+
+		self.finish_responses(&mut futures).boxed().await;
 
 		Ok(())
 	}
 
+	#[tracing::instrument(
+		name = "work",
+		level = "trace"
+		skip_all,
+		fields(
+			futures = %futures.len(),
+			statuses = %statuses.len(),
+		),
+	)]
+	async fn work_loop<'a>(
+		&'a self,
+		id: usize,
+		futures: &mut SendingFutures<'a>,
+		statuses: &mut CurTransactionStatus,
+	) {
+		let receiver = self
+			.channels
+			.get(id)
+			.map(|(_, receiver)| receiver.clone())
+			.expect("Missing channel for sender worker");
+		loop {
+			tokio::select! {
+				Some(response) = futures.next() => {
+					self.handle_response(response, futures, statuses).await;
+				},
+				request = receiver.recv_async() => match request {
+					Ok(request) => self.handle_request(request, futures, statuses).await,
+					Err(_) => return,
+				},
+			}
+		}
+	}
+
+	#[tracing::instrument(name = "response", level = "debug", skip_all)]
 	async fn handle_response<'a>(
 		&'a self,
 		response: SendingResult,
@@ -138,13 +168,14 @@ impl Service {
 			self.db.mark_as_active(new_events.iter());
 
 			let new_events_vec = new_events.into_iter().map(|(_, event)| event).collect();
-			futures.push(self.send_events(dest.clone(), new_events_vec).boxed());
+			futures.push(self.send_events(dest.clone(), new_events_vec));
 		} else {
 			statuses.remove(dest);
 		}
 	}
 
 	#[allow(clippy::needless_pass_by_ref_mut)]
+	#[tracing::instrument(name = "request", level = "debug", skip_all)]
 	async fn handle_request<'a>(
 		&'a self,
 		msg: Msg,
@@ -154,13 +185,19 @@ impl Service {
 		let iv = vec![(msg.queue_id, msg.event)];
 		if let Ok(Some(events)) = self.select_events(&msg.dest, iv, statuses).await {
 			if !events.is_empty() {
-				futures.push(self.send_events(msg.dest, events).boxed());
+				futures.push(self.send_events(msg.dest, events));
 			} else {
 				statuses.remove(&msg.dest);
 			}
 		}
 	}
 
+	#[tracing::instrument(
+		name = "finish",
+		level = "info",
+		skip_all,
+		fields(futures = %futures.len()),
+	)]
 	async fn finish_responses<'a>(&'a self, futures: &mut SendingFutures<'a>) {
 		use tokio::{
 			select,
@@ -183,9 +220,16 @@ impl Service {
 		}
 	}
 
+	#[tracing::instrument(
+		name = "netburst",
+		level = "debug",
+		skip_all,
+		fields(futures = %futures.len()),
+	)]
 	#[allow(clippy::needless_pass_by_ref_mut)]
-	async fn initial_requests<'a>(
+	async fn startup_netburst<'a>(
 		&'a self,
+		id: usize,
 		futures: &mut SendingFutures<'a>,
 		statuses: &mut CurTransactionStatus,
 	) {
@@ -195,6 +239,10 @@ impl Service {
 		let mut active = self.db.active_requests().boxed();
 
 		while let Some((key, event, dest)) = active.next().await {
+			if self.shard_id(&dest) != id {
+				continue;
+			}
+
 			let entry = txns.entry(dest.clone()).or_default();
 			if self.server.config.startup_netburst_keep >= 0 && entry.len() >= keep {
 				warn!("Dropping unsent event {dest:?} {:?}", String::from_utf8_lossy(&key));
@@ -207,19 +255,27 @@ impl Service {
 		for (dest, events) in txns {
 			if self.server.config.startup_netburst && !events.is_empty() {
 				statuses.insert(dest.clone(), TransactionStatus::Running);
-				futures.push(self.send_events(dest.clone(), events).boxed());
+				futures.push(self.send_events(dest.clone(), events));
 			}
 		}
 	}
 
-	#[tracing::instrument(skip_all, level = "debug")]
+	#[tracing::instrument(
+		name = "select",,
+		level = "debug",
+		skip_all,
+		fields(
+			?dest,
+			new_events = %new_events.len(),
+		)
+	)]
 	async fn select_events(
 		&self,
 		dest: &Destination,
 		new_events: Vec<QueueItem>, // Events we want to send: event and full key
 		statuses: &mut CurTransactionStatus,
 	) -> Result<Option<Vec<SendingEvent>>> {
-		let (allow, retry) = self.select_events_current(dest.clone(), statuses)?;
+		let (allow, retry) = self.select_events_current(dest, statuses)?;
 
 		// Nothing can be done for this remote, bail out.
 		if !allow {
@@ -249,7 +305,7 @@ impl Service {
 		}
 
 		// Add EDU's into the transaction
-		if let Destination::Normal(server_name) = dest {
+		if let Destination::Federation(server_name) = dest {
 			if let Ok((select_edus, last_count)) = self.select_edus(server_name).await {
 				debug_assert!(select_edus.len() <= EDU_LIMIT, "exceeded edus limit");
 				events.extend(select_edus.into_iter().map(SendingEvent::Edu));
@@ -260,10 +316,9 @@ impl Service {
 		Ok(Some(events))
 	}
 
-	#[tracing::instrument(skip_all, level = "debug")]
 	fn select_events_current(
 		&self,
-		dest: Destination,
+		dest: &Destination,
 		statuses: &mut CurTransactionStatus,
 	) -> Result<(bool, bool)> {
 		let (mut allow, mut retry) = (true, false);
@@ -292,7 +347,11 @@ impl Service {
 		Ok((allow, retry))
 	}
 
-	#[tracing::instrument(skip_all, level = "debug")]
+	#[tracing::instrument(
+		name = "edus",,
+		level = "debug",
+		skip_all,
+	)]
 	async fn select_edus(&self, server_name: &ServerName) -> Result<(Vec<Vec<u8>>, u64)> {
 		// selection window
 		let since = self.db.get_latest_educount(server_name).await;
@@ -329,7 +388,12 @@ impl Service {
 		Ok((events, max_edu_count.load(Ordering::Acquire)))
 	}
 
-	/// Look for presence
+	/// Look for device changes
+	#[tracing::instrument(
+		name = "device_changes",
+		level = "trace",
+		skip(self, server_name, max_edu_count)
+	)]
 	async fn select_edus_device_changes(
 		&self,
 		server_name: &ServerName,
@@ -386,6 +450,11 @@ impl Service {
 	}
 
 	/// Look for read receipts in this room
+	#[tracing::instrument(
+		name = "receipts",
+		level = "trace",
+		skip(self, server_name, max_edu_count)
+	)]
 	async fn select_edus_receipts(
 		&self,
 		server_name: &ServerName,
@@ -420,6 +489,7 @@ impl Service {
 	}
 
 	/// Look for read receipts in this room
+	#[tracing::instrument(name = "receipts", level = "trace", skip(self, since, max_edu_count))]
 	async fn select_edus_receipts_room(
 		&self,
 		room_id: &RoomId,
@@ -484,6 +554,11 @@ impl Service {
 	}
 
 	/// Look for presence
+	#[tracing::instrument(
+		name = "presence",
+		level = "trace",
+		skip(self, server_name, max_edu_count)
+	)]
 	async fn select_edus_presence(
 		&self,
 		server_name: &ServerName,
@@ -554,29 +629,33 @@ impl Service {
 		Some(presence_content)
 	}
 
-	async fn send_events(&self, dest: Destination, events: Vec<SendingEvent>) -> SendingResult {
+	fn send_events(&self, dest: Destination, events: Vec<SendingEvent>) -> SendingFuture<'_> {
 		//debug_assert!(!events.is_empty(), "sending empty transaction");
 		match dest {
-			| Destination::Normal(ref server) =>
-				self.send_events_dest_normal(&dest, server, events).await,
-			| Destination::Appservice(ref id) =>
-				self.send_events_dest_appservice(&dest, id, events).await,
-			| Destination::Push(ref userid, ref pushkey) =>
-				self.send_events_dest_push(&dest, userid, pushkey, events)
-					.await,
+			| Destination::Federation(server) =>
+				self.send_events_dest_federation(server, events).boxed(),
+			| Destination::Appservice(id) => self.send_events_dest_appservice(id, events).boxed(),
+			| Destination::Push(user_id, pushkey) =>
+				self.send_events_dest_push(user_id, pushkey, events).boxed(),
 		}
 	}
 
-	#[tracing::instrument(skip(self, dest, events), name = "appservice")]
+	#[tracing::instrument(
+		name = "appservice",
+		level = "debug",
+		skip(self, events),
+		fields(
+			events = %events.len(),
+		),
+	)]
 	async fn send_events_dest_appservice(
 		&self,
-		dest: &Destination,
-		id: &str,
+		id: String,
 		events: Vec<SendingEvent>,
 	) -> SendingResult {
-		let Some(appservice) = self.services.appservice.get_registration(id).await else {
+		let Some(appservice) = self.services.appservice.get_registration(&id).await else {
 			return Err((
-				dest.clone(),
+				Destination::Appservice(id.clone()),
 				err!(Database(warn!(?id, "Missing appservice registration"))),
 			));
 		};
@@ -633,23 +712,29 @@ impl Service {
 		)
 		.await
 		{
-			| Ok(_) => Ok(dest.clone()),
-			| Err(e) => Err((dest.clone(), e)),
+			| Ok(_) => Ok(Destination::Appservice(id)),
+			| Err(e) => Err((Destination::Appservice(id), e)),
 		}
 	}
 
-	#[tracing::instrument(skip(self, dest, events), name = "push")]
+	#[tracing::instrument(
+		name = "push",
+		level = "info",
+		skip(self, events),
+		fields(
+			events = %events.len(),
+		),
+	)]
 	async fn send_events_dest_push(
 		&self,
-		dest: &Destination,
-		userid: &OwnedUserId,
-		pushkey: &str,
+		user_id: OwnedUserId,
+		pushkey: String,
 		events: Vec<SendingEvent>,
 	) -> SendingResult {
-		let Ok(pusher) = self.services.pusher.get_pusher(userid, pushkey).await else {
+		let Ok(pusher) = self.services.pusher.get_pusher(&user_id, &pushkey).await else {
 			return Err((
-				dest.clone(),
-				err!(Database(error!(?userid, ?pushkey, "Missing pusher"))),
+				Destination::Push(user_id.clone(), pushkey.clone()),
+				err!(Database(error!(?user_id, ?pushkey, "Missing pusher"))),
 			));
 		};
 
@@ -677,17 +762,17 @@ impl Service {
 			let rules_for_user = self
 				.services
 				.account_data
-				.get_global(userid, GlobalAccountDataEventType::PushRules)
+				.get_global(&user_id, GlobalAccountDataEventType::PushRules)
 				.await
 				.map_or_else(
-					|_| push::Ruleset::server_default(userid),
+					|_| push::Ruleset::server_default(&user_id),
 					|ev: PushRulesEvent| ev.content.global,
 				);
 
 			let unread: UInt = self
 				.services
 				.user
-				.notification_count(userid, &pdu.room_id)
+				.notification_count(&user_id, &pdu.room_id)
 				.await
 				.try_into()
 				.expect("notification count can't go that high");
@@ -695,19 +780,25 @@ impl Service {
 			let _response = self
 				.services
 				.pusher
-				.send_push_notice(userid, unread, &pusher, rules_for_user, &pdu)
+				.send_push_notice(&user_id, unread, &pusher, rules_for_user, &pdu)
 				.await
-				.map_err(|e| (dest.clone(), e));
+				.map_err(|e| (Destination::Push(user_id.clone(), pushkey.clone()), e));
 		}
 
-		Ok(dest.clone())
+		Ok(Destination::Push(user_id, pushkey))
 	}
 
-	#[tracing::instrument(skip(self, dest, events), name = "", level = "debug")]
-	async fn send_events_dest_normal(
+	#[tracing::instrument(
+		name = "fed",
+		level = "debug",
+		skip(self, events),
+		fields(
+			events = %events.len(),
+		),
+	)]
+	async fn send_events_dest_federation(
 		&self,
-		dest: &Destination,
-		server: &OwnedServerName,
+		server: OwnedServerName,
 		events: Vec<SendingEvent>,
 	) -> SendingResult {
 		let mut pdu_jsons = Vec::with_capacity(
@@ -759,7 +850,7 @@ impl Service {
 		};
 
 		let client = &self.services.client.sender;
-		self.send(client, server, request)
+		self.send(client, &server, request)
 			.await
 			.inspect(|response| {
 				response
@@ -770,8 +861,8 @@ impl Service {
 						|(pdu_id, res)| warn!(%txn_id, %server, "error sending PDU {pdu_id} to remote server: {res:?}"),
 					);
 			})
-			.map(|_| dest.clone())
-			.map_err(|e| (dest.clone(), e))
+			.map_err(|e| (Destination::Federation(server.clone()), e))
+			.map(|_| Destination::Federation(server))
 	}
 
 	/// This does not return a full `Pdu` it is only to satisfy ruma's types.