feat: Add server health tracking

src/service/sending/mod.rs

@@ -5,17 +5,23 @@
 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, debug, debug_warn, err, error,
+	Result, Server, SyncRwLock, debug, debug_warn, err, error,
 	smallvec::SmallVec,
-	utils::{ReadyExt, TryReadyExt, available_parallelism, math::usize_from_u64_truncated},
+	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,
+	},
 	warn,
 };
 use futures::{FutureExt, Stream, StreamExt};
@@ -48,6 +54,7 @@ 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 {
@@ -108,6 +115,7 @@ 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()),
 		}))
 	}
 
@@ -427,6 +435,62 @@ 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) {
+		let mut map = self.remote_health.write();
+		map.remove(server_name);
+	}
 }
 
 fn num_senders(args: &crate::Args<'_>) -> usize {