diff --git a/src/proto/discovery/limits.rs b/src/proto/discovery/limits.rs index c3e9dc6c..b92a4b1b 100644 --- a/src/proto/discovery/limits.rs +++ b/src/proto/discovery/limits.rs @@ -14,6 +14,7 @@ //! nodes generating fresh request_ids at high rate. use crate::NodeAddr; +use crate::proto::rate_limit::PerAddrRateLimiter; use alloc::collections::BTreeMap; // ============================================================================ @@ -34,9 +35,6 @@ const DEFAULT_BACKOFF_BASE_SECS: u64 = 0; /// Default maximum backoff cap. `0` = disabled. const DEFAULT_BACKOFF_MAX_SECS: u64 = 0; -/// Backoff multiplier per consecutive failure. -const BACKOFF_MULTIPLIER: u64 = 2; - /// Exponential backoff for failed discovery lookups. /// /// Tracks targets whose lookups have timed out and suppresses @@ -91,10 +89,11 @@ impl DiscoveryBackoff { pub fn record_failure(&mut self, target: &NodeAddr, now_ms: u64) { let failures = self.entries.get(target).map_or(0, |e| e.failures) + 1; - let backoff_ms = self - .base_ms - .saturating_mul(BACKOFF_MULTIPLIER.saturating_pow(failures.saturating_sub(1))) - .min(self.max_ms); + let backoff_ms = crate::proto::rate_limit::backoff_ms( + failures.saturating_sub(1), + self.base_ms, + self.max_ms, + ); self.entries.insert( *target, @@ -160,29 +159,20 @@ const FORWARD_MAX_AGE_MS: u64 = 60_000; /// Tracks the last time a LookupRequest was forwarded for each target /// and enforces a minimum interval to prevent floods from misbehaving /// nodes generating fresh request_ids. -pub struct DiscoveryForwardRateLimiter { - last_forwarded: BTreeMap, - min_interval_ms: u64, - max_age_ms: u64, -} +pub struct DiscoveryForwardRateLimiter(PerAddrRateLimiter); impl DiscoveryForwardRateLimiter { /// Create with default parameters (2s interval). pub fn new() -> Self { - Self { - last_forwarded: BTreeMap::new(), - min_interval_ms: DEFAULT_FORWARD_MIN_INTERVAL_MS, - max_age_ms: FORWARD_MAX_AGE_MS, - } + Self(PerAddrRateLimiter::new( + DEFAULT_FORWARD_MIN_INTERVAL_MS, + FORWARD_MAX_AGE_MS, + )) } /// Create with a custom minimum interval in milliseconds. pub fn with_interval_ms(min_interval_ms: u64) -> Self { - Self { - last_forwarded: BTreeMap::new(), - min_interval_ms, - max_age_ms: FORWARD_MAX_AGE_MS, - } + Self(PerAddrRateLimiter::new(min_interval_ms, FORWARD_MAX_AGE_MS)) } /// Check if we should forward a lookup for this target. @@ -190,32 +180,24 @@ impl DiscoveryForwardRateLimiter { /// Returns true if enough time has passed since the last forward /// for this target. Updates internal state when returning true. pub fn should_forward(&mut self, target: &NodeAddr, now_ms: u64) -> bool { - if let Some(&last) = self.last_forwarded.get(target) - && now_ms.saturating_sub(last) < self.min_interval_ms - { - return false; - } - - self.last_forwarded.insert(*target, now_ms); - self.cleanup(now_ms); - true + self.0.check_and_record(target, now_ms) } /// Replace the minimum interval in milliseconds (e.g., set to zero to disable). #[cfg(test)] pub fn set_interval_ms(&mut self, interval_ms: u64) { - self.min_interval_ms = interval_ms; + self.0.set_interval_ms(interval_ms); } /// Remove entries older than max_age. + #[cfg(test)] pub(crate) fn cleanup(&mut self, now_ms: u64) { - self.last_forwarded - .retain(|_, &mut last| now_ms.saturating_sub(last) < self.max_age_ms); + self.0.cleanup(now_ms); } #[cfg(test)] pub fn len(&self) -> usize { - self.last_forwarded.len() + self.0.len() } } diff --git a/src/proto/fmp/limits.rs b/src/proto/fmp/limits.rs index 24407ade..f304d2b2 100644 --- a/src/proto/fmp/limits.rs +++ b/src/proto/fmp/limits.rs @@ -9,8 +9,5 @@ /// Uses exponential backoff: `base_interval_ms * 2^retry_count`, capped at /// `max_backoff_ms`. pub(crate) fn backoff_ms(retry_count: u32, base_interval_ms: u64, max_backoff_ms: u64) -> u64 { - let multiplier = 1u64.checked_shl(retry_count).unwrap_or(u64::MAX); - base_interval_ms - .saturating_mul(multiplier) - .min(max_backoff_ms) + crate::proto::rate_limit::backoff_ms(retry_count, base_interval_ms, max_backoff_ms) } diff --git a/src/proto/mod.rs b/src/proto/mod.rs index 6cb6281c..e5fb961a 100644 --- a/src/proto/mod.rs +++ b/src/proto/mod.rs @@ -13,5 +13,6 @@ pub(crate) mod fmp; pub(crate) mod fsp; pub(crate) mod link; pub(crate) mod mmp; +pub(crate) mod rate_limit; pub(crate) mod routing; pub(crate) mod stp; diff --git a/src/proto/rate_limit.rs b/src/proto/rate_limit.rs new file mode 100644 index 00000000..2eeb83a7 --- /dev/null +++ b/src/proto/rate_limit.rs @@ -0,0 +1,60 @@ +//! Shared `proto` rate-limiting / backoff primitives: a per-address +//! minimum-interval limiter and an exponential backoff helper, hoisted out of +//! the subsystem `limits.rs` files. + +use crate::NodeAddr; +use alloc::collections::BTreeMap; + +/// Per-address minimum-interval rate limiter. Tracks the last event time per +/// address and enforces a minimum interval, evicting entries older than a max age. +pub(crate) struct PerAddrRateLimiter { + last: BTreeMap, + min_interval_ms: u64, + max_age_ms: u64, +} + +impl PerAddrRateLimiter { + pub(crate) fn new(min_interval_ms: u64, max_age_ms: u64) -> Self { + Self { + last: BTreeMap::new(), + min_interval_ms, + max_age_ms, + } + } + + /// Returns true (and records `now_ms`) if enough time has elapsed since the + /// last event for `addr`, or this is the first; false if within the interval. + pub(crate) fn check_and_record(&mut self, addr: &NodeAddr, now_ms: u64) -> bool { + if let Some(&last) = self.last.get(addr) + && now_ms.saturating_sub(last) < self.min_interval_ms + { + return false; + } + self.last.insert(*addr, now_ms); + self.cleanup(now_ms); + true + } + + pub(crate) fn cleanup(&mut self, now_ms: u64) { + self.last + .retain(|_, &mut last| now_ms.saturating_sub(last) < self.max_age_ms); + } + + #[cfg(test)] + pub(crate) fn set_interval_ms(&mut self, interval_ms: u64) { + self.min_interval_ms = interval_ms; + } + + #[cfg(test)] + pub(crate) fn len(&self) -> usize { + self.last.len() + } +} + +/// Exponential (base-2) backoff: `base_ms * 2^exponent`, saturating, capped at +/// `cap_ms`. Shared by the discovery originator backoff (exponent = failures-1) +/// and the FMP retry scheduler (exponent = retry_count). +pub(crate) fn backoff_ms(exponent: u32, base_ms: u64, cap_ms: u64) -> u64 { + let multiplier = 1u64.checked_shl(exponent).unwrap_or(u64::MAX); + base_ms.saturating_mul(multiplier).min(cap_ms) +} diff --git a/src/proto/routing/limits.rs b/src/proto/routing/limits.rs index d16cdfb3..01762d1f 100644 --- a/src/proto/routing/limits.rs +++ b/src/proto/routing/limits.rs @@ -9,7 +9,7 @@ //! portability and deterministic ordering. use crate::NodeAddr; -use alloc::collections::BTreeMap; +use crate::proto::rate_limit::PerAddrRateLimiter; /// Default minimum interval between error signals: 100 ms (max 10 errors/sec /// per destination). @@ -22,35 +22,19 @@ const MAX_AGE_MS: u64 = 10_000; /// /// Tracks the last time a routing error was sent for each destination /// address and enforces a minimum interval to prevent floods. -pub struct RoutingErrorRateLimiter { - /// Maps destination NodeAddr to the last time (Unix ms) we sent an error - /// about it. - last_sent: BTreeMap, - /// Minimum interval between error signals for the same destination (ms). - min_interval_ms: u64, - /// Maximum age of entries before cleanup (ms). - max_age_ms: u64, -} +pub struct RoutingErrorRateLimiter(PerAddrRateLimiter); impl RoutingErrorRateLimiter { /// Create a new rate limiter. /// /// Default: max 10 errors/sec per destination (100ms interval). pub fn new() -> Self { - Self { - last_sent: BTreeMap::new(), - min_interval_ms: DEFAULT_MIN_INTERVAL_MS, - max_age_ms: MAX_AGE_MS, - } + Self(PerAddrRateLimiter::new(DEFAULT_MIN_INTERVAL_MS, MAX_AGE_MS)) } /// Create a rate limiter with a custom minimum interval in milliseconds. pub fn with_interval_ms(min_interval_ms: u64) -> Self { - Self { - last_sent: BTreeMap::new(), - min_interval_ms, - max_age_ms: MAX_AGE_MS, - } + Self(PerAddrRateLimiter::new(min_interval_ms, MAX_AGE_MS)) } /// Check if we should send a routing error for this destination at @@ -60,26 +44,18 @@ impl RoutingErrorRateLimiter { /// this destination, or if this is the first error. Updates internal /// state when returning true. pub fn should_send(&mut self, dest_addr: &NodeAddr, now_ms: u64) -> bool { - if let Some(&last) = self.last_sent.get(dest_addr) - && now_ms.saturating_sub(last) < self.min_interval_ms - { - return false; - } - - self.last_sent.insert(*dest_addr, now_ms); - self.cleanup(now_ms); - true + self.0.check_and_record(dest_addr, now_ms) } /// Remove entries older than max_age. + #[cfg(test)] pub(crate) fn cleanup(&mut self, now_ms: u64) { - self.last_sent - .retain(|_, &mut last| now_ms.saturating_sub(last) < self.max_age_ms); + self.0.cleanup(now_ms); } #[cfg(test)] pub fn len(&self) -> usize { - self.last_sent.len() + self.0.len() } }