diff --git a/Cargo.lock b/Cargo.lock index 9cb2617310c..8572cf1cd97 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5584,6 +5584,7 @@ dependencies = [ "prometheus-client", "proptest", "rand 0.9.2", + "rand_chacha 0.9.0", "regex", "serde", "sha2", diff --git a/beacon_node/lighthouse_network/Cargo.toml b/beacon_node/lighthouse_network/Cargo.toml index f69f13612ab..e7578bcff40 100644 --- a/beacon_node/lighthouse_network/Cargo.toml +++ b/beacon_node/lighthouse_network/Cargo.toml @@ -56,6 +56,7 @@ unsigned-varint = { version = "0.8", features = ["codec"] } async-channel = { workspace = true } logging = { workspace = true } proptest = { workspace = true } +rand_chacha = "0.9.0" tempfile = { workspace = true } [[test]] diff --git a/beacon_node/lighthouse_network/src/peer_manager/mod.rs b/beacon_node/lighthouse_network/src/peer_manager/mod.rs index 898b97a85f7..14f74e20743 100644 --- a/beacon_node/lighthouse_network/src/peer_manager/mod.rs +++ b/beacon_node/lighthouse_network/src/peer_manager/mod.rs @@ -8,6 +8,9 @@ use discv5::Enr; use libp2p::identify::Info as IdentifyInfo; use lru_cache::LRUTimeCache; use peerdb::{BanOperation, BanResult, ScoreUpdateResult}; +use rand::RngCore; +use rand::SeedableRng; +use rand::rngs::StdRng; use rand::seq::SliceRandom; use smallvec::SmallVec; use std::{ @@ -76,6 +79,14 @@ pub const PRIORITY_PEER_EXCESS: f32 = 0.2; /// The numbre of inbound libp2p peers we have seen before we consider our NAT to be open. pub const LIBP2P_NAT_OPEN_THRESHOLD: usize = 3; +/// The default RNG used by the [`PeerManager`] for peer-selection randomness in production. +/// +/// This is seeded from OS entropy (equivalent randomness to the previously-used thread-local RNG), +/// but is `Send` so it can live inside the `PeerManager`. +fn default_rng() -> Box { + Box::new(StdRng::from_os_rng()) +} + /// The main struct that handles peer's reputation and connection status. pub struct PeerManager { /// Storage of network globals to access the `PeerDB`. @@ -122,6 +133,13 @@ pub struct PeerManager { /// Keeps track of whether the QUIC protocol is enabled or not. quic_enabled: bool, trusted_peers: HashSet, + /// Random number generator used for peer-selection randomness (e.g. shuffling prune + /// candidates). + /// + /// In production this is backed by the thread-local RNG, preserving the previous behaviour. + /// Tests can replace it with a seeded RNG to make peer selection deterministic via + /// [`PeerManager::set_rng`]. + rng: Box, } /// The events that the `PeerManager` outputs (requests). @@ -203,9 +221,19 @@ impl PeerManager { metrics_enabled, quic_enabled, trusted_peers: Default::default(), + rng: default_rng(), }) } + /// Replace the RNG used for peer-selection randomness. + /// + /// This is intended for tests that want deterministic peer selection by supplying a seeded + /// RNG. Production code uses the default thread-local RNG installed in [`PeerManager::new`]. + #[cfg(test)] + pub(crate) fn set_rng(&mut self, rng: Box) { + self.rng = rng; + } + /* Public accessible functions */ /// The application layer wants to disconnect from a peer for a particular reason. @@ -1198,6 +1226,7 @@ impl PeerManager { /// Find the best candidate for removal from the densest custody subnet. /// /// Returns the PeerId of the candidate to remove, or None if no suitable candidate found. + #[allow(clippy::too_many_arguments)] fn find_prune_candidate( &self, column_subnet: DataColumnSubnetId, @@ -1206,12 +1235,13 @@ impl PeerManager { sampling_subnets: &HashSet, connected_outbound_peer_count: usize, outbound_peers_pruned: usize, + rng: &mut dyn RngCore, ) -> Option { let peers_on_subnet_clone = column_subnet_to_peers.get(&column_subnet)?.clone(); // Create a sorted list of peers prioritized for removal let mut sorted_peers = peers_on_subnet_clone; - sorted_peers.shuffle(&mut rand::rng()); + sorted_peers.shuffle(rng); sorted_peers.sort_by_key(|peer_id| { if let Some(peer_info) = peer_subnet_info.get(peer_id) { ( @@ -1354,6 +1384,10 @@ impl PeerManager { let mut peer_subnet_info = self.build_peer_subnet_info(&peers_to_prune); let mut custody_subnet_to_peers = Self::build_custody_subnet_lookup(&peer_subnet_info); + // Temporarily take the RNG out of `self` so we can pass it to `find_prune_candidate` + // (which borrows `&self`) without a borrow conflict. It is always put back below. + let mut rng = std::mem::replace(&mut self.rng, default_rng()); + // Attempt to prune peers to `target_peers`, or until we run out of peers to prune. while peers_to_prune.len() < connected_peer_count.saturating_sub(self.target_peers) { let custody_subnet_with_most_peers = custody_subnet_to_peers @@ -1374,6 +1408,7 @@ impl PeerManager { &sampling_subnets, connected_outbound_peer_count, outbound_peers_pruned, + rng.as_mut(), ) { // Update outbound peer count if needed if let Some(candidate_info) = peer_subnet_info.get(&candidate_peer) @@ -1399,6 +1434,9 @@ impl PeerManager { break; } } + + // Restore the RNG we temporarily took out above. + self.rng = rng; } // Disconnect the pruned peers. @@ -1422,6 +1460,15 @@ impl PeerManager { /// /// NOTE: Discovery will only add a new query if one isn't already queued. fn heartbeat(&mut self) { + self.heartbeat_at(Instant::now()); + } + + /// As [`PeerManager::heartbeat`], but uses the supplied `now` to drive the time-based parts of + /// the heartbeat (score decay / ban expiry). + /// + /// The production path (the libp2p poll loop) calls [`PeerManager::heartbeat`] with + /// `Instant::now()`. Tests can advance logical time deterministically by calling this directly. + fn heartbeat_at(&mut self, now: Instant) { // Optionally run a discovery query if we need more peers. self.maintain_peer_count(0); self.maintain_trusted_peers(); @@ -1434,7 +1481,7 @@ impl PeerManager { self.network_globals.peers.write().cleanup_dialing_peers(); // Updates peer's scores and unban any peers if required. - let actions = self.network_globals.peers.write().update_scores(); + let actions = self.network_globals.peers.write().update_scores_at(now); for (peer_id, action) in actions { self.handle_score_action(&peer_id, action, None); } @@ -3258,4 +3305,259 @@ mod tests { "Should generate discovery events when PeerDAS is enabled, but found no discovery events" ); } + + /// Deterministic ("Tier A") tests for the `PeerManager`. + /// + /// These tests mirror the approach used by the sync subsystem's `TestRig`: they avoid all + /// sources of non-determinism so that outcomes are reproducible. + /// + /// The two sources of non-determinism in the `PeerManager` are: + /// 1. Time: score decay and ban expiry. These are driven by an injectable `now` threaded + /// through `PeerManager::heartbeat_at`. Tests advance *logical* time (offsets from a + /// reference `Instant`) instead of sleeping, so no wall-clock time elapses. + /// 2. Randomness: peer selection during pruning shuffles candidates. The `PeerManager` holds + /// an injectable RNG (`PeerManager::set_rng`) which tests seed with a fixed + /// `ChaCha20Rng::from_seed([..])` to make selection reproducible. + mod deterministic_tests { + use super::*; + use crate::peer_manager::peerdb::score::ScoreState; + use rand_chacha::ChaCha20Rng; + + /// A small deterministic test harness for the `PeerManager`. + /// + /// It owns a `PeerManager` configured with a seeded RNG and a fixed reference `Instant` + /// (`t0`) from which all logical time offsets are computed. Helpers let a test connect + /// peers, report peer actions, advance logical time (by running a heartbeat at a given + /// offset), and inspect the resulting score / connection / ban state. + struct PeerManagerRig { + peer_manager: PeerManager, + /// Reference instant. All logical time is expressed as offsets from here. + t0: Instant, + } + + impl PeerManagerRig { + /// Build a rig with a seeded RNG so that any peer-selection randomness is reproducible. + async fn new(target_peer_count: usize, seed: [u8; 32]) -> Self { + let mut peer_manager = build_peer_manager(target_peer_count).await; + peer_manager.set_rng(Box::new(ChaCha20Rng::from_seed(seed))); + // Capture the reference instant *after* construction. New peers created below will + // anchor their score `last_updated` at roughly this instant. + let t0 = Instant::now(); + Self { peer_manager, t0 } + } + + /// Connect an inbound peer. + fn connect_ingoing(&mut self, peer_id: &PeerId) { + self.peer_manager.inject_connect_ingoing( + peer_id, + "/ip4/0.0.0.0".parse().unwrap(), + None, + ); + } + + /// Report a peer action (e.g. a downscore or a fatal action). + fn report(&mut self, peer_id: &PeerId, action: PeerAction) { + self.peer_manager.report_peer( + peer_id, + action, + ReportSource::PeerManager, + None, + "deterministic_test", + ); + } + + /// Run a heartbeat at logical time `t0 + offset`. This is how tests advance time + /// without sleeping: score decay / ban expiry are computed against this `now`. + fn heartbeat_after(&mut self, offset: Duration) { + self.peer_manager.heartbeat_at(self.t0 + offset); + } + + /// The current score-derived state of a peer (Healthy / ForcedDisconnect / Banned). + fn score_state(&self, peer_id: &PeerId) -> Option { + self.peer_manager + .network_globals + .peers + .read() + .peer_info(peer_id) + .map(|info| info.score_state()) + } + + /// The numeric score of a peer. + fn score(&self, peer_id: &PeerId) -> f64 { + self.peer_manager + .network_globals + .peers + .read() + .score(peer_id) + } + + /// Whether the peer is currently considered banned by its score. + fn is_banned(&self, peer_id: &PeerId) -> bool { + matches!( + self.peer_manager.ban_status(peer_id), + Some(BanResult::BadScore) + ) + } + } + + /// A peer reported with `PeerAction::Fatal` is immediately banned (no heartbeat / time + /// advance required). + #[tokio::test] + async fn fatal_action_bans_immediately() { + let peer = PeerId::random(); + let mut rig = PeerManagerRig::new(50, [0u8; 32]).await; + rig.connect_ingoing(&peer); + + assert_eq!(rig.score_state(&peer), Some(ScoreState::Healthy)); + assert!(!rig.is_banned(&peer)); + + rig.report(&peer, PeerAction::Fatal); + + assert_eq!(rig.score_state(&peer), Some(ScoreState::Banned)); + assert!(rig.is_banned(&peer)); + } + + /// A peer accumulates downscores across several reports and transitions + /// Healthy -> ForcedDisconnect -> Banned at the expected thresholds. + /// + /// `LowToleranceError` is -10 per report. Disconnect threshold is -20, ban threshold is + /// -50, so we expect: Healthy until the 2nd report (-20 -> ForcedDisconnect), and Banned + /// once the score reaches -50 (the 5th report). + #[tokio::test] + async fn downscores_transition_healthy_disconnect_banned() { + let peer = PeerId::random(); + let mut rig = PeerManagerRig::new(50, [1u8; 32]).await; + rig.connect_ingoing(&peer); + + // 1 report: score -10, still Healthy (disconnect threshold is -20). + rig.report(&peer, PeerAction::LowToleranceError); + assert!((rig.score(&peer) - -10.0).abs() < 1e-9); + assert_eq!(rig.score_state(&peer), Some(ScoreState::Healthy)); + + // 2 reports: score -20, now ForcedDisconnect. + rig.report(&peer, PeerAction::LowToleranceError); + assert!((rig.score(&peer) - -20.0).abs() < 1e-9); + assert_eq!(rig.score_state(&peer), Some(ScoreState::ForcedDisconnect)); + + // 4 reports: score -40, still ForcedDisconnect (ban threshold is -50). + rig.report(&peer, PeerAction::LowToleranceError); + rig.report(&peer, PeerAction::LowToleranceError); + assert!((rig.score(&peer) - -40.0).abs() < 1e-9); + assert_eq!(rig.score_state(&peer), Some(ScoreState::ForcedDisconnect)); + + // 5 reports: score -50, now Banned. + rig.report(&peer, PeerAction::LowToleranceError); + assert!((rig.score(&peer) - -50.0).abs() < 1e-9); + assert_eq!(rig.score_state(&peer), Some(ScoreState::Banned)); + assert!(rig.is_banned(&peer)); + } + + /// The same transition reachable via `MidToleranceError` (-5 per report). Ban requires 10 + /// reports (-50). + #[tokio::test] + async fn mid_tolerance_downscores_eventually_ban() { + let peer = PeerId::random(); + let mut rig = PeerManagerRig::new(50, [2u8; 32]).await; + rig.connect_ingoing(&peer); + + // 9 reports: -45, still ForcedDisconnect. + for _ in 0..9 { + rig.report(&peer, PeerAction::MidToleranceError); + } + assert!((rig.score(&peer) - -45.0).abs() < 1e-9); + assert_eq!(rig.score_state(&peer), Some(ScoreState::ForcedDisconnect)); + + // 10th report: -50, Banned. + rig.report(&peer, PeerAction::MidToleranceError); + assert!((rig.score(&peer) - -50.0).abs() < 1e-9); + assert_eq!(rig.score_state(&peer), Some(ScoreState::Banned)); + } + + /// Ban expiry: after advancing logical time past the ban duration, the peer's ban decays. + /// + /// When a peer is banned, its score `last_updated` is pushed `BANNED_BEFORE_DECAY` into the + /// future, so the score does not decay until that period elapses. We verify that: + /// - immediately after a ban, and right up to the ban-before-decay boundary, the peer + /// stays banned even when heartbeats run; + /// - once logical time passes the boundary, a heartbeat lets the score decay above the + /// ban threshold and the peer is no longer banned. + /// + /// All time advances are logical (offsets from `t0`); no real sleeping occurs. + #[tokio::test] + async fn ban_decays_after_logical_time_advance() { + use peerdb::score::testing::BANNED_BEFORE_DECAY; + + let peer = PeerId::random(); + let mut rig = PeerManagerRig::new(50, [3u8; 32]).await; + rig.connect_ingoing(&peer); + + // Ban the peer. + rig.report(&peer, PeerAction::Fatal); + assert!(rig.is_banned(&peer)); + + // Running a heartbeat just before the ban-before-decay boundary must NOT decay the + // ban: the peer is still banned. + rig.heartbeat_after(BANNED_BEFORE_DECAY.saturating_sub(Duration::from_secs(1))); + assert!( + rig.is_banned(&peer), + "peer should remain banned before the ban-before-decay period elapses" + ); + + // Advance logical time well past the ban-before-decay boundary and run a heartbeat. + // The score should now decay above the ban threshold and the peer unbans. + // We advance an extra ~5 half-lives (SCORE_HALFLIFE = 600s) to ensure the score (which + // starts at the floor of -100) decays above the -50 ban threshold. + rig.heartbeat_after(BANNED_BEFORE_DECAY + Duration::from_secs(3600)); + + assert!( + !rig.is_banned(&peer), + "peer should no longer be banned after ban decays; score = {}", + rig.score(&peer) + ); + assert!( + rig.score(&peer) > -50.0, + "score should have decayed above the ban threshold, got {}", + rig.score(&peer) + ); + } + + /// Determinism: running the same scenario twice with the same seed produces identical + /// outcomes (identical final scores for the same sequence of reports). + /// + /// We use the same fixed `PeerId`s across both runs so the comparison is well-defined. + #[tokio::test] + async fn same_seed_produces_identical_outcomes() { + // Fixed peer ids so both runs operate on identical inputs. + let peers: Vec = (0..6).map(|_| PeerId::random()).collect(); + let actions = [ + PeerAction::LowToleranceError, + PeerAction::MidToleranceError, + PeerAction::HighToleranceError, + ]; + + async fn run_scenario(peers: &[PeerId], actions: &[PeerAction]) -> Vec<(PeerId, f64)> { + let mut rig = PeerManagerRig::new(3, [7u8; 32]).await; + for peer in peers { + rig.connect_ingoing(peer); + } + // Apply a deterministic sequence of reports. + for (i, peer) in peers.iter().enumerate() { + rig.report(peer, actions[i % actions.len()]); + } + // Run a heartbeat at a fixed logical time (this exercises the seeded RNG via the + // pruning path, since we connected more peers than the target of 3). + rig.heartbeat_after(Duration::from_secs(0)); + + peers.iter().map(|p| (*p, rig.score(p))).collect() + } + + let run_a = run_scenario(&peers, &actions).await; + let run_b = run_scenario(&peers, &actions).await; + + assert_eq!( + run_a, run_b, + "the same seed and inputs must produce identical final scores" + ); + } + } } diff --git a/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs b/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs index 693fdebb69b..0f49207f756 100644 --- a/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs +++ b/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs @@ -472,7 +472,11 @@ impl PeerDB { /// NOTE: Peer scores cannot be penalized during the update, they can only increase. Therefore /// it not possible to ban peers when updating scores. #[must_use = "The unbanned peers must be sent to libp2p"] - pub(super) fn update_scores(&mut self) -> Vec<(PeerId, ScoreUpdateResult)> { + /// Applies score decay and ban-expiry to all peers using the supplied `now`. + /// + /// The production path (the peer manager heartbeat) passes `Instant::now()`; tests can advance + /// logical time deterministically by supplying their own `now`. + pub(super) fn update_scores_at(&mut self, now: Instant) -> Vec<(PeerId, ScoreUpdateResult)> { // Peer can be unbanned in this process. // We return the result, such that the peer manager can inform the swarm to lift the libp2p // ban on these peers. @@ -482,7 +486,7 @@ impl PeerDB { for (peer_id, info) in self.peers.iter_mut() { let previous_state = info.score_state(); // Update scores - info.score_update(); + info.score_update_at(now); match Self::handle_score_transition(previous_state, peer_id, info) { // A peer should not be able to be banned from a score update. diff --git a/beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs b/beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs index 658a6355e3f..80d788a0066 100644 --- a/beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs +++ b/beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs @@ -458,10 +458,13 @@ impl PeerInfo { self.partial_message_subnets.clear() } - /// Applies decay rates to a non-trusted peer's score. - pub(super) fn score_update(&mut self) { + /// Applies decay rates to a non-trusted peer's score using the supplied `now`. + /// + /// The production path passes `Instant::now()` (via the peer manager heartbeat); tests can + /// advance logical time deterministically instead of relying on the wall clock. + pub(super) fn score_update_at(&mut self, now: std::time::Instant) { if !self.is_trusted { - self.score.update() + self.score.update_at(now) } } diff --git a/beacon_node/lighthouse_network/src/peer_manager/peerdb/score.rs b/beacon_node/lighthouse_network/src/peer_manager/peerdb/score.rs index e57e7907db7..7659cb14e80 100644 --- a/beacon_node/lighthouse_network/src/peer_manager/peerdb/score.rs +++ b/beacon_node/lighthouse_network/src/peer_manager/peerdb/score.rs @@ -219,15 +219,12 @@ impl RealScore { self.update_state(); } - /// Applies time-based logic such as decay rates to the score. - /// This function should be called periodically. - pub fn update(&mut self) { - self.update_at(Instant::now()) - } - /// Applies time-based logic such as decay rates to the score with the given now value. - /// This private sub function is mainly used for testing. - fn update_at(&mut self, now: Instant) { + /// This function should be called periodically. + /// + /// The production path passes `Instant::now()` (via the peer manager heartbeat). Tests can + /// drive decay/ban-expiry deterministically by supplying their own `now`. + pub(crate) fn update_at(&mut self, now: Instant) { // Decay the current score // Using exponential decay based on a constant half life. @@ -288,7 +285,7 @@ macro_rules! apply { } apply!(apply_peer_action, peer_action: PeerAction); -apply!(update); +apply!(update_at, now: Instant); apply!(update_gossipsub_score, new_score: f64, ignore: bool); #[cfg(test)] apply!(test_add, score: f64); @@ -348,6 +345,14 @@ impl std::fmt::Display for Score { } } +/// Test-only re-exports of internal scoring constants so deterministic tests in other modules can +/// reason about ban-expiry timing without duplicating magic numbers. +#[cfg(test)] +pub mod testing { + /// The number of seconds we ban a peer for before their score begins to decay. + pub const BANNED_BEFORE_DECAY: super::Duration = super::BANNED_BEFORE_DECAY; +} + #[cfg(test)] mod tests { use super::*; diff --git a/beacon_node/lighthouse_network/src/rpc/rate_limiter.rs b/beacon_node/lighthouse_network/src/rpc/rate_limiter.rs index a5c98a4d309..605cee1f6fd 100644 --- a/beacon_node/lighthouse_network/src/rpc/rate_limiter.rs +++ b/beacon_node/lighthouse_network/src/rpc/rate_limiter.rs @@ -396,12 +396,33 @@ impl RPCRateLimiter { RPCRateLimiterBuilder::default() } + /// The instant the rate limiter was created. Logical time offsets used by the `*_at` methods + /// are measured relative to this. Exposed for deterministic tests. + #[cfg(test)] + pub(crate) fn init_time(&self) -> Instant { + self.init_time + } + pub fn allows( &mut self, peer_id: &PeerId, request: &Item, ) -> Result<(), RateLimitedErr> { - let time_since_start = self.init_time.elapsed(); + self.allows_at(Instant::now(), peer_id, request) + } + + /// As [`RPCRateLimiter::allows`], but uses the supplied `now` to compute the elapsed time + /// since the limiter was created. + /// + /// The production path (via [`RPCRateLimiter::allows`]) passes `Instant::now()`; tests can + /// supply their own `now` to advance logical time deterministically without sleeping. + pub(crate) fn allows_at( + &mut self, + now: Instant, + peer_id: &PeerId, + request: &Item, + ) -> Result<(), RateLimitedErr> { + let time_since_start = now.saturating_duration_since(self.init_time); let tokens = request .max_responses( self.fork_context.current_fork_epoch(), @@ -434,7 +455,14 @@ impl RPCRateLimiter { } pub fn prune(&mut self) { - let time_since_start = self.init_time.elapsed(); + self.prune_at(Instant::now()); + } + + /// As [`RPCRateLimiter::prune`], but uses the supplied `now` to compute the elapsed time since + /// the limiter was created. The production path passes `Instant::now()`; tests can advance + /// logical time deterministically. + pub(crate) fn prune_at(&mut self, now: Instant) { + let time_since_start = now.saturating_duration_since(self.init_time); let Self { prune_interval: _, @@ -679,4 +707,161 @@ mod tests { let result = limiter.allows(Duration::from_secs_f32(0.0), &10, tokens); assert!(matches!(result, Err(RateLimitedErr::TooLarge))); } + + /// A partial refill after a partial interval should make exactly the refilled tokens + /// available, and no more. + /// + /// Quota: 4 tokens per 2s => 1 token replenished every 0.5s. After consuming all 4 tokens at + /// t=0, at t=1.0s exactly 2 tokens have been replenished: two 1-token requests succeed and a + /// third fails. This is a pure-math (logical time) test of the token-bucket refill edge case. + #[test] + fn partial_refill_after_partial_interval() { + let mut limiter = Limiter::from_quota(Quota { + replenish_all_every: Duration::from_secs(2), + max_tokens: NonZeroU64::new(4).unwrap(), + }) + .unwrap(); + let key = 1; + + // Drain the full bucket at t=0. + assert!(limiter.allows(Duration::from_secs(0), &key, 4).is_ok()); + // Immediately after, the bucket is empty. + assert!(limiter.allows(Duration::from_secs(0), &key, 1).is_err()); + + // After 1.0s, exactly 2 tokens are replenished (1 token / 0.5s). + assert!( + limiter + .allows(Duration::from_secs_f32(1.0), &key, 1) + .is_ok() + ); + assert!( + limiter + .allows(Duration::from_secs_f32(1.0), &key, 1) + .is_ok() + ); + // The third token at t=1.0s is not yet available. + assert!( + limiter + .allows(Duration::from_secs_f32(1.0), &key, 1) + .is_err() + ); + } + + /// Deterministic ("Tier A") tests at the `RPCRateLimiter` level, exercising the injectable + /// time seam (`allows_at` / `init_time`). No real time elapses: all time is logical, expressed + /// as offsets from the limiter's `init_time`. + mod rpc_rate_limiter_deterministic { + use crate::rpc::config::RateLimiterConfig; + use crate::rpc::rate_limiter::{RPCRateLimiter, RateLimitedErr}; + use crate::rpc::{Ping, RequestType}; + use libp2p::PeerId; + use std::num::NonZeroU64; + use std::sync::Arc; + use std::time::Duration; + use types::{EthSpec, ForkContext, Hash256, MainnetEthSpec, Slot}; + + fn fork_context() -> Arc { + Arc::new(ForkContext::new::( + Slot::new(0), + Hash256::ZERO, + &MainnetEthSpec::default_spec(), + )) + } + + /// Build a limiter whose Ping protocol allows `max_tokens` tokens every `secs` seconds. + fn limiter_with_ping_quota(max_tokens: u64, secs: u64) -> RPCRateLimiter { + let config = RateLimiterConfig { + ping_quota: crate::rpc::rate_limiter::Quota::n_every( + NonZeroU64::new(max_tokens).unwrap(), + secs, + ), + ..Default::default() + }; + RPCRateLimiter::new_with_config(config, fork_context()).unwrap() + } + + fn ping(data: u64) -> RequestType { + RequestType::Ping(Ping { data }) + } + + /// The limiter allows up to capacity, then rejects once the bucket is empty. + #[tokio::test] + async fn allows_up_to_capacity_then_rejects() { + // 3 tokens (each Ping = 1 token) per 30s. + let mut limiter = limiter_with_ping_quota(3, 30); + let peer = PeerId::random(); + let t0 = limiter.init_time(); + + // All requests happen at the same logical instant t0. + assert!(limiter.allows_at(t0, &peer, &ping(0)).is_ok()); + assert!(limiter.allows_at(t0, &peer, &ping(1)).is_ok()); + assert!(limiter.allows_at(t0, &peer, &ping(2)).is_ok()); + + // The 4th request at the same instant is rejected (bucket empty). + assert!(matches!( + limiter.allows_at(t0, &peer, &ping(3)), + Err(RateLimitedErr::TooSoon(_)) + )); + } + + /// After advancing logical time enough to refill, requests are allowed again. + #[tokio::test] + async fn refills_after_logical_time_advance() { + // 1 token per 10s, so a full refill takes 10s. + let mut limiter = limiter_with_ping_quota(1, 10); + let peer = PeerId::random(); + let t0 = limiter.init_time(); + + // Consume the single token. + assert!(limiter.allows_at(t0, &peer, &ping(0)).is_ok()); + // Immediately rejected. + assert!(matches!( + limiter.allows_at(t0, &peer, &ping(1)), + Err(RateLimitedErr::TooSoon(_)) + )); + + // Still too soon after 9 logical seconds. + assert!(matches!( + limiter.allows_at(t0 + Duration::from_secs(9), &peer, &ping(2)), + Err(RateLimitedErr::TooSoon(_)) + )); + + // After 10 logical seconds the token is replenished and the request is allowed. + assert!( + limiter + .allows_at(t0 + Duration::from_secs(10), &peer, &ping(3)) + .is_ok() + ); + } + + /// `prune_at` removes keys whose bucket is full by the given logical time, but keeps keys + /// that are still rate limited. + #[tokio::test] + async fn prune_at_respects_logical_time() { + // 1 token per 10s. + let mut limiter = limiter_with_ping_quota(1, 10); + let peer = PeerId::random(); + let t0 = limiter.init_time(); + + // Consume the token, putting the peer's TAT into the future (bucket not yet full). + assert!(limiter.allows_at(t0, &peer, &ping(0)).is_ok()); + + // Pruning at t0 must NOT remove the peer: their bucket is not full yet, so a new + // request at t0 is still rate limited. + limiter.prune_at(t0); + assert!(matches!( + limiter.allows_at(t0, &peer, &ping(1)), + Err(RateLimitedErr::TooSoon(_)) + )); + + // Pruning well after the bucket has refilled removes the key; a subsequent request is + // treated as a fresh (full-bucket) peer and is allowed. + limiter.prune_at(t0 + Duration::from_secs(100)); + assert!( + limiter + .allows_at(t0 + Duration::from_secs(100), &peer, &ping(2)) + .is_ok() + ); + } + } } diff --git a/beacon_node/lighthouse_network/src/rpc/self_limiter.rs b/beacon_node/lighthouse_network/src/rpc/self_limiter.rs index 2a7ef955a19..1588c73290c 100644 --- a/beacon_node/lighthouse_network/src/rpc/self_limiter.rs +++ b/beacon_node/lighthouse_network/src/rpc/self_limiter.rs @@ -13,7 +13,7 @@ use std::{ collections::{HashMap, VecDeque, hash_map::Entry}, sync::Arc, task::{Context, Poll}, - time::Duration, + time::{Duration, Instant}, }; use tokio_util::time::DelayQueue; use tracing::debug; @@ -100,6 +100,7 @@ impl SelfRateLimiter { match Self::try_send_request( &mut self.active_requests, &mut self.rate_limiter, + Instant::now(), peer_id, request_id, req, @@ -125,6 +126,7 @@ impl SelfRateLimiter { fn try_send_request( active_requests: &mut HashMap>, rate_limiter: &mut Option, + now: Instant, peer_id: PeerId, request_id: Id, req: RequestType, @@ -149,7 +151,7 @@ impl SelfRateLimiter { } if let Some(limiter) = rate_limiter.as_mut() { - match limiter.allows(&peer_id, &req) { + match limiter.allows_at(now, &peer_id, &req) { Ok(()) => {} Err(e) => { let protocol = req.versioned_protocol(); @@ -190,6 +192,13 @@ impl SelfRateLimiter { /// When a peer and protocol are allowed to send a next request, this function checks the /// queued requests and attempts marking as ready as many as the limiter allows. fn next_peer_request_ready(&mut self, peer_id: PeerId, protocol: Protocol) { + self.next_peer_request_ready_at(Instant::now(), peer_id, protocol) + } + + /// As [`SelfRateLimiter::next_peer_request_ready`], but uses the supplied `now` to drive the + /// inner rate limiter. The production path passes `Instant::now()`; tests can advance logical + /// time deterministically. + fn next_peer_request_ready_at(&mut self, now: Instant, peer_id: PeerId, protocol: Protocol) { if let Entry::Occupied(mut entry) = self.delayed_requests.entry((peer_id, protocol)) { let queued_requests = entry.get_mut(); while let Some(QueuedRequest { @@ -201,6 +210,7 @@ impl SelfRateLimiter { match Self::try_send_request( &mut self.active_requests, &mut self.rate_limiter, + now, peer_id, request_id, req.clone(), @@ -275,6 +285,13 @@ impl SelfRateLimiter { failed_requests } + /// The `init_time` of the inner rate limiter, against which logical time offsets are measured. + /// Exposed for deterministic tests. + #[cfg(test)] + pub(crate) fn rate_limiter_init_time(&self) -> Option { + self.rate_limiter.as_ref().map(|rl| rl.init_time()) + } + /// Informs the limiter that a response has been received. pub fn request_completed(&mut self, peer_id: &PeerId, protocol: Protocol) { if let Some(active_requests) = self.active_requests.get_mut(peer_id) @@ -570,6 +587,94 @@ mod tests { ); } + /// Deterministic version of the "requests become ready after the tokens regenerate" scenario. + /// + /// Unlike `test_next_peer_request_ready`, this advances *logical* time via + /// `next_peer_request_ready_at(now)` instead of sleeping. It asserts that: + /// - before enough logical time has passed, no queued request becomes ready; + /// - after advancing logical time past the replenish interval, exactly one queued request + /// becomes ready (1 token / interval), matching the quota. + #[tokio::test] + async fn test_next_peer_request_ready_deterministic() { + use std::time::{Duration, Instant}; + + // 1 ping token per 2 seconds. + let config = OutboundRateLimiterConfig(RateLimiterConfig { + ping_quota: Quota::n_every(NonZeroU64::new(1).unwrap(), 2), + ..Default::default() + }); + let fork_context = std::sync::Arc::new(ForkContext::new::( + Slot::new(0), + Hash256::ZERO, + &MainnetEthSpec::default_spec(), + )); + let mut limiter: SelfRateLimiter = + SelfRateLimiter::new(Some(config), fork_context).unwrap(); + let peer_id = PeerId::random(); + let lookup_id = 0; + + // The first request is allowed (consumes the single token); the next four are queued. + for i in 1..=5u32 { + let _ = limiter.allows( + peer_id, + AppRequestId::Sync(SyncRequestId::SingleBlock { + id: SingleLookupReqId { + lookup_id, + req_id: i, + }, + }), + RequestType::Ping(Ping { data: i as u64 }), + ); + } + + let init_time = limiter + .rate_limiter_init_time() + .expect("rate limiter is configured"); + assert_eq!( + limiter + .delayed_requests + .get(&(peer_id, Protocol::Ping)) + .unwrap() + .len(), + 4 + ); + assert_eq!(limiter.ready_requests.len(), 0); + + // Advancing logical time only a little (well under the 2s replenish interval) does not make + // any queued request ready. We measure from `init_time` so the offset is comparable to the + // real instant the requests were enqueued. + let too_soon = init_time + Duration::from_millis(100); + // Guard: `too_soon` must not actually be in the past relative to wall clock (it won't be, + // since the requests above ran in well under 100ms), keeping the test deterministic. + assert!(too_soon <= Instant::now() + Duration::from_secs(2)); + limiter.next_peer_request_ready_at(too_soon, peer_id, Protocol::Ping); + assert_eq!( + limiter + .delayed_requests + .get(&(peer_id, Protocol::Ping)) + .unwrap() + .len(), + 4, + "no request should be ready before the replenish interval elapses" + ); + assert_eq!(limiter.ready_requests.len(), 0); + + // Advance logical time past the 2s replenish interval: exactly one token regenerates, so + // exactly one queued request becomes ready. + let after_refill = init_time + Duration::from_secs(3); + limiter.next_peer_request_ready_at(after_refill, peer_id, Protocol::Ping); + assert_eq!( + limiter + .delayed_requests + .get(&(peer_id, Protocol::Ping)) + .unwrap() + .len(), + 3, + "exactly one request should become ready after one token regenerates" + ); + assert_eq!(limiter.ready_requests.len(), 1); + } + /// Test that `peer_disconnected` returns the IDs of pending requests. #[tokio::test] async fn test_peer_disconnected_returns_failed_requests() {