diff --git a/src/api/router/auth.rs b/src/api/router/auth.rs index 8906277b8..f8f72519c 100644 --- a/src/api/router/auth.rs +++ b/src/api/router/auth.rs @@ -150,7 +150,7 @@ async fn verify + Sync>( } // Ping the server as healthy - services.sending.mark_healthy(&output.origin); + services.federation.mark_healthy(&output.origin); Ok(output.origin) }, diff --git a/src/service/federation/mod.rs b/src/service/federation/mod.rs index 703f939b2..e15f7f960 100644 --- a/src/service/federation/mod.rs +++ b/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>, } struct Services { @@ -26,8 +32,78 @@ fn build(args: crate::Args<'_>) -> Result> { server_keys: args.depend::("server_keys"), moderation: args.depend::("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 { + 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); + } +} diff --git a/src/service/sending/mod.rs b/src/service/sending/mod.rs index 52e56d50e..ac9920f76 100644 --- a/src/service/sending/mod.rs +++ b/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, services: Services, channels: Vec<(loole::Sender, loole::Receiver)>, - remote_health: SyncRwLock>, } struct Services { @@ -115,7 +108,6 @@ fn build(args: crate::Args<'_>) -> Result> { federation: args.depend::("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 { diff --git a/src/service/sending/sender.rs b/src/service/sending/sender.rs index 683131e47..d1940ffb3 100644 --- a/src/service/sending/sender.rs +++ b/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