diff --git a/src/service/sending/mod.rs b/src/service/sending/mod.rs index ac9920f76..6433fdd8e 100644 --- a/src/service/sending/mod.rs +++ b/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, services: Services, channels: Vec<(loole::Sender, loole::Receiver)>, + remote_health: SyncRwLock>, } struct Services { @@ -108,6 +115,7 @@ fn build(args: crate::Args<'_>) -> Result> { federation: args.depend::("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 {