chore: Move health tracking to federation service

Otherwise federation would have to depend on sender

src/api/router/auth.rs

@@ -150,7 +150,7 @@ async fn verify<B: AsRef<[u8]> + Sync>(
 				}
 
 				// Ping the server as healthy
-				services.sending.mark_healthy(&output.origin);
+				services.federation.mark_healthy(&output.origin);
 
 				Ok(output.origin)
 			},

src/service/federation/mod.rs

@@ -1,13 +1,19 @@
 mod execute;
-use std::sync::Arc;
 
-use conduwuit::{Result, Server};
+use std::{collections::HashMap, sync::Arc, time::Duration};
+
+use conduwuit::{
+	Result, Server, SyncRwLock,
+	utils::{millis_since_unix_epoch, time::exponential_backoff::min_exp_backoff_duration},
+};
 pub(crate) use execute::FederationPathBuilderInput;
+use ruma::{OwnedServerName, ServerName};
 
 use crate::{Dep, client, moderation, server_keys};
 
 pub struct Service {
 	services: Services,
+	remote_health: SyncRwLock<HashMap<OwnedServerName, (u32, u64)>>,
 }
 
 struct Services {
@@ -26,8 +32,78 @@ fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
 				server_keys: args.depend::<server_keys::Service>("server_keys"),
 				moderation: args.depend::<moderation::Service>("moderation"),
 			},
+			remote_health: SyncRwLock::new(HashMap::new()),
 		}))
 	}
 
 	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
 }
+
+impl Service {
+	/// Checks if a remote is "healthy". "Healthy" is defined by either:
+	///
+	/// * The remote has not been marked as having a failed request, OR
+	/// * The next retry timestamp is in the past
+	pub fn is_healthy(&self, server_name: &ServerName) -> bool {
+		let map = self.remote_health.read();
+		let unix_now = millis_since_unix_epoch();
+		if let Some((_, next_retry)) = map.get(server_name) {
+			unix_now >= *next_retry
+		} else {
+			true
+		}
+	}
+
+	/// Returns how long the server should wait before attempting to contact the
+	/// remote again.
+	pub fn retry_after(&self, server_name: &ServerName) -> Option<Duration> {
+		let map = self.remote_health.read();
+		let unix_now = millis_since_unix_epoch();
+		map.get(server_name)
+			.map(|(_, next_retry)| Duration::from_millis(*next_retry - unix_now))
+	}
+
+	/// Marks or updates a remote's health status as unhealthy. If the remote is
+	/// not already marked as unhealthy, a new entry is created. Otherwise, the
+	/// retry count is incremented and
+	pub fn hit_unhealthy(&self, server_name: OwnedServerName) {
+		// TODO(nex): Can multiple concurrent failures cause this health monitor to
+		// rapidly max out?
+		//
+		// consider: a profile query and key claim query both go out at the same time,
+		// and both fail. They both then go to hit this (synchronous) function, which
+		// means one of them will increment the failure count by 1, and consequently
+		// increase the exp backoff. Then the second failure is allowed to call the
+		// function, at which point it increments the failure count AGAIN, increasing
+		// the backoff again.
+		// Technically, these are two distinct failures. However, from a UX perspective,
+		// they happened at the same time, so shouldn't incur a double penalty?
+		// Perhaps only incrementing the retry counter when the current next_retry is in
+		// the past might help.
+		//
+		// You know it's a banger thought process when the comment is longer than the
+		// code itself.
+		let unix_now = millis_since_unix_epoch();
+		let mut map = self.remote_health.write();
+		let (retries, next_retry) = map.entry(server_name).or_default();
+
+		let min = self.services.server.config.sender_timeout;
+		let max = self.services.server.config.sender_retry_backoff_limit;
+
+		*retries = retries.saturating_add(1);
+		*next_retry = unix_now.saturating_add(
+			u64::try_from(min_exp_backoff_duration(min, max, *retries).as_millis())
+				.expect("backoff milliseconds should not exceed u64::MAX"),
+		);
+	}
+
+	/// Marks a server as "healthy" by removing it from the health map.
+	///
+	/// TODO: flush senders too
+	pub fn mark_healthy(&self, server_name: &ServerName) {
+		// TODO: We need to make sure the sender flush DOESN'T trigger if this is called
+		// by the senders themselves.
+		let mut map = self.remote_health.write();
+		map.remove(server_name);
+	}
+}

src/service/sending/mod.rs

@@ -5,23 +5,17 @@
 mod sender;
 
 use std::{
-	collections::HashMap,
 	fmt::Debug,
 	hash::{DefaultHasher, Hash, Hasher},
 	iter::once,
 	sync::Arc,
-	time::SystemTime,
 };
 
 use async_trait::async_trait;
 use conduwuit::{
-	Result, Server, SyncRwLock, debug, debug_warn, err, error,
+	Result, Server, debug, debug_warn, err, error,
 	smallvec::SmallVec,
-	utils::{
-		ReadyExt, TryFutureExtExt, TryReadyExt, available_parallelism,
-		continue_exponential_backoff, math::usize_from_u64_truncated, millis_since_unix_epoch,
-		time::exponential_backoff::min_exp_backoff_duration,
-	},
+	utils::{ReadyExt, TryReadyExt, available_parallelism, math::usize_from_u64_truncated},
 	warn,
 };
 use futures::{FutureExt, Stream, StreamExt};
@@ -54,7 +48,6 @@ pub struct Service {
 	server: Arc<Server>,
 	services: Services,
 	channels: Vec<(loole::Sender<Msg>, loole::Receiver<Msg>)>,
-	remote_health: SyncRwLock<HashMap<OwnedServerName, (u32, u64)>>,
 }
 
 struct Services {
@@ -115,7 +108,6 @@ fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
 				federation: args.depend::<federation::Service>("federation"),
 			},
 			channels: (0..num_senders).map(|_| loole::unbounded()).collect(),
-			remote_health: SyncRwLock::new(HashMap::new()),
 		}))
 	}
 
@@ -435,64 +427,6 @@ pub(super) fn shard_id(&self, dest: &Destination) -> usize {
 		let chans = self.channels.len().max(1);
 		hash.overflowing_rem(chans).0
 	}
-
-	/// Checks if a remote is "healthy". "Healthy" is defined by either:
-	///
-	/// * The remote has not been marked as having a failed request, OR
-	/// * The next retry timestamp is in the past
-	pub fn is_healthy(&self, server_name: &ServerName) -> bool {
-		let map = self.remote_health.read();
-		let unix_now = millis_since_unix_epoch();
-		if let Some((_, next_retry)) = map.get(server_name) {
-			unix_now >= *next_retry
-		} else {
-			true
-		}
-	}
-
-	/// Marks or updates a remote's health status as unhealthy. If the remote is
-	/// not already marked as unhealthy, a new entry is created. Otherwise, the
-	/// retry count is incremented and
-	pub fn hit_unhealthy(&self, server_name: OwnedServerName) {
-		// TODO(nex): Can multiple concurrent failures cause this health monitor to
-		// rapidly max out?
-		//
-		// consider: a profile query and key claim query both go out at the same time,
-		// and both fail. They both then go to hit this (synchronous) function, which
-		// means one of them will increment the failure count by 1, and consequently
-		// increase the exp backoff. Then the second failure is allowed to call the
-		// function, at which point it increments the failure count AGAIN, increasing
-		// the backoff again.
-		// Technically, these are two distinct failures. However, from a UX perspective,
-		// they happened at the same time, so shouldn't incur a double penalty?
-		// Perhaps only incrementing the retry counter when the current next_retry is in
-		// the past might help.
-		//
-		// You know it's a banger thought process when the comment is longer than the
-		// code itself.
-		let unix_now = millis_since_unix_epoch();
-		let mut map = self.remote_health.write();
-		let (retries, next_retry) = map.entry(server_name).or_default();
-
-		let min = self.server.config.sender_timeout;
-		let max = self.server.config.sender_retry_backoff_limit;
-
-		*retries = retries.saturating_add(1);
-		*next_retry = unix_now.saturating_add(
-			u64::try_from(min_exp_backoff_duration(min, max, *retries).as_millis())
-				.expect("backoff milliseconds should not exceed u64::MAX"),
-		);
-	}
-
-	/// Marks a server as "healthy" by removing it from the health map.
-	///
-	/// TODO: flush senders too
-	pub fn mark_healthy(&self, server_name: &ServerName) {
-		// TODO: We need to make sure the sender flush DOESN'T trigger if this is called
-		// by the senders themselves.
-		let mut map = self.remote_health.write();
-		map.remove(server_name);
-	}
 }
 
 fn num_senders(args: &crate::Args<'_>) -> usize {

src/service/sending/sender.rs

@@ -146,7 +146,7 @@ async fn handle_response<'a>(
 				if e.status_code().is_server_error()
 					&& let Destination::Federation(dest) = &dest
 				{
-					self.hit_unhealthy(dest.clone());
+					self.services.federation.hit_unhealthy(dest.clone());
 				}
 				Self::handle_response_err(dest, statuses, &e);
 			},
@@ -179,7 +179,7 @@ async fn handle_response_ok<'a>(
 		let _cork = self.db.db.cork();
 		self.db.delete_all_active_requests_for(dest).await;
 		if let Destination::Federation(server_name) = dest {
-			self.mark_healthy(server_name);
+			self.services.federation.mark_healthy(server_name);
 		}
 
 		// Find events that have been added since starting the last request