From 9769a1b49eb22d24a80d39c5245a469ded4f1750 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Mon, 10 Aug 2026 01:00:50 +0700 Subject: [PATCH] refactor(client): index pending probe deadlines --- crates/irtt-client/src/probe.rs | 210 ++++++++++++++++++++-- crates/irtt-client/src/session/machine.rs | 28 +++ 2 files changed, 220 insertions(+), 18 deletions(-) diff --git a/crates/irtt-client/src/probe.rs b/crates/irtt-client/src/probe.rs index 682bcfb..f39fa37 100644 --- a/crates/irtt-client/src/probe.rs +++ b/crates/irtt-client/src/probe.rs @@ -1,5 +1,5 @@ use std::{ - collections::{HashMap, HashSet, TryReserveError, VecDeque}, + collections::{BTreeSet, HashMap, HashSet, TryReserveError, VecDeque}, time::Instant, }; @@ -12,9 +12,33 @@ pub(crate) struct PendingProbe { pub timeout_at: Instant, } +#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)] +struct DeadlineKey { + timeout_at: Instant, + sent_at: Instant, + wire_seq: u32, +} + +impl From<&PendingProbe> for DeadlineKey { + fn from(probe: &PendingProbe) -> Self { + Self { + timeout_at: probe.timeout_at, + sent_at: probe.sent_at.mono, + wire_seq: probe.wire_seq, + } + } +} + +#[derive(Debug)] +pub(crate) struct ExpiredBatch { + pub probes: Vec, + pub more_due: bool, +} + #[derive(Debug)] pub(crate) struct PendingMap { map: HashMap, + deadlines: BTreeSet, max_capacity: usize, } @@ -22,6 +46,7 @@ impl PendingMap { pub fn new(max_capacity: usize) -> Self { Self { map: HashMap::new(), + deadlines: BTreeSet::new(), max_capacity, } } @@ -42,43 +67,75 @@ impl PendingMap { } pub fn commit_insert(&mut self, probe: PendingProbe) { + let inserted = self.deadlines.insert(DeadlineKey::from(&probe)); + debug_assert!(inserted, "pending deadline index already contained probe"); let replaced = self.map.insert(probe.wire_seq, probe); debug_assert!(replaced.is_none(), "preflight rejected pending collision"); + self.assert_index_consistent(); } pub fn remove(&mut self, wire_seq: u32) -> Option { - self.map.remove(&wire_seq) + let probe = self.map.remove(&wire_seq)?; + let removed = self.deadlines.remove(&DeadlineKey::from(&probe)); + debug_assert!(removed, "pending deadline index lost removed probe"); + self.assert_index_consistent(); + Some(probe) } pub fn drain_expired(&mut self, now: Instant) -> Vec { - let expired_keys: Vec = self - .map - .iter() - .filter(|(_, probe)| probe.timeout_at <= now) - .map(|(key, _)| *key) - .collect(); - let mut expired = Vec::with_capacity(expired_keys.len()); - for key in expired_keys { - if let Some(probe) = self.map.remove(&key) { - expired.push(probe); + let batch = self.drain_expired_bounded(now, usize::MAX); + debug_assert!(!batch.more_due); + batch.probes + } + + pub(crate) fn drain_expired_bounded(&mut self, now: Instant, limit: usize) -> ExpiredBatch { + let mut probes = Vec::with_capacity(limit.min(self.map.len())); + while probes.len() < limit { + let Some(key) = self.deadlines.first().copied() else { + break; + }; + if key.timeout_at > now { + break; } + + let removed_key = self + .deadlines + .take(&key) + .expect("first pending deadline remains present"); + debug_assert_eq!(removed_key, key); + let probe = self + .map + .remove(&key.wire_seq) + .expect("pending deadline always has a pending probe"); + debug_assert_eq!(DeadlineKey::from(&probe), key); + probes.push(probe); } - expired.sort_by_key(|p| p.sent_at.mono); - expired + self.assert_index_consistent(); + let more_due = self + .deadlines + .first() + .is_some_and(|key| key.timeout_at <= now); + ExpiredBatch { probes, more_due } } pub fn len(&self) -> usize { self.map.len() } - #[cfg(feature = "tokio")] pub fn next_timeout_deadline(&self) -> Option { - self.map.values().map(|probe| probe.timeout_at).min() + self.deadlines.first().map(|key| key.timeout_at) } - #[cfg(feature = "tokio")] pub fn latest_timeout_deadline(&self) -> Option { - self.map.values().map(|probe| probe.timeout_at).max() + self.deadlines.last().map(|key| key.timeout_at) + } + + fn assert_index_consistent(&self) { + debug_assert_eq!( + self.deadlines.len(), + self.map.len(), + "pending deadline index cardinality diverged" + ); } #[cfg(test)] @@ -90,6 +147,11 @@ impl PendingMap { pub fn capacity(&self) -> usize { self.map.capacity() } + + #[cfg(test)] + pub fn deadline_index_len(&self) -> usize { + self.deadlines.len() + } } fn pending_allocation_failed(source: TryReserveError) -> ClientError { @@ -244,6 +306,11 @@ mod tests { } } + fn insert(map: &mut PendingMap, probe: PendingProbe) { + map.preflight_insert(probe.wire_seq).unwrap(); + map.commit_insert(probe); + } + #[test] fn pending_map_rejects_over_capacity() { let mut map = PendingMap::new(2); @@ -303,6 +370,113 @@ mod tests { )); } + #[test] + fn pending_map_indexes_deadline_extrema_and_removals() { + let mut map = PendingMap::new(4); + let now = Instant::now(); + insert(&mut map, pending(1, now + Duration::from_secs(3))); + insert(&mut map, pending(2, now + Duration::from_secs(1))); + insert(&mut map, pending(3, now + Duration::from_secs(2))); + + assert_eq!( + map.next_timeout_deadline(), + Some(now + Duration::from_secs(1)) + ); + assert_eq!( + map.latest_timeout_deadline(), + Some(now + Duration::from_secs(3)) + ); + assert_eq!(map.deadline_index_len(), map.len()); + + assert!(map.remove(3).is_some()); + assert_eq!( + map.next_timeout_deadline(), + Some(now + Duration::from_secs(1)) + ); + assert_eq!( + map.latest_timeout_deadline(), + Some(now + Duration::from_secs(3)) + ); + assert_eq!(map.deadline_index_len(), map.len()); + + assert!(map.remove(2).is_some()); + assert_eq!( + map.next_timeout_deadline(), + Some(now + Duration::from_secs(3)) + ); + assert_eq!( + map.latest_timeout_deadline(), + Some(now + Duration::from_secs(3)) + ); + assert_eq!(map.deadline_index_len(), map.len()); + + assert!(map.remove(1).is_some()); + assert_eq!(map.next_timeout_deadline(), None); + assert_eq!(map.latest_timeout_deadline(), None); + assert_eq!(map.deadline_index_len(), map.len()); + } + + #[test] + fn pending_map_bounded_expiration_matches_exhaustive_order() { + let now = Instant::now(); + let probes = [ + pending(4, now + Duration::from_secs(4)), + pending(2, now + Duration::from_secs(2)), + pending(1, now + Duration::from_secs(1)), + pending(3, now + Duration::from_secs(3)), + ]; + let mut bounded = PendingMap::new(probes.len()); + let mut exhaustive = PendingMap::new(probes.len()); + for probe in probes.iter().cloned() { + insert(&mut bounded, probe.clone()); + insert(&mut exhaustive, probe); + } + + let empty = bounded.drain_expired_bounded(now + Duration::from_secs(4), 0); + assert!(empty.probes.is_empty()); + assert!(empty.more_due); + assert_eq!(bounded.deadline_index_len(), bounded.len()); + + let mut actual = Vec::new(); + loop { + let batch = bounded.drain_expired_bounded(now + Duration::from_secs(4), 2); + actual.extend(batch.probes.into_iter().map(|probe| probe.wire_seq)); + assert_eq!(bounded.deadline_index_len(), bounded.len()); + if !batch.more_due { + break; + } + } + let expected = exhaustive + .drain_expired(now + Duration::from_secs(4)) + .into_iter() + .map(|probe| probe.wire_seq) + .collect::>(); + + assert_eq!(actual, expected); + assert_eq!(actual, vec![1, 2, 3, 4]); + assert_eq!(bounded.deadline_index_len(), 0); + } + + #[test] + fn pending_map_sequence_reuse_leaves_no_stale_deadline() { + let mut map = PendingMap::new(1); + let now = Instant::now(); + insert(&mut map, pending(7, now)); + assert!(map.remove(7).is_some()); + insert(&mut map, pending(7, now + Duration::from_secs(1))); + + assert!(map.drain_expired(now).is_empty()); + assert_eq!(map.deadline_index_len(), map.len()); + assert_eq!( + map.drain_expired(now + Duration::from_secs(1)) + .into_iter() + .map(|probe| probe.wire_seq) + .collect::>(), + vec![7] + ); + assert_eq!(map.deadline_index_len(), map.len()); + } + #[test] fn bounded_probe_tracking_evicts_oldest_entries() { let mut set = CompletedSet::new(3); diff --git a/crates/irtt-client/src/session/machine.rs b/crates/irtt-client/src/session/machine.rs index 7d3d8f9..a180909 100644 --- a/crates/irtt-client/src/session/machine.rs +++ b/crates/irtt-client/src/session/machine.rs @@ -1803,6 +1803,34 @@ mod tests { assert_eq!(active(&machine).pending.len(), 0); } + #[test] + fn exhaustive_timeout_polling_returns_every_due_loss() { + let mut machine = open_machine(4, Duration::from_secs(1)); + let now = Instant::now(); + for seq in [2, 0, 1] { + let sent_at = timestamp(now - Duration::from_secs(u64::from(seq) + 1)); + let session = active_mut(&mut machine); + session.pending.preflight_insert(seq).unwrap(); + session.pending.commit_insert(PendingProbe { + wire_seq: seq, + sent_at, + timeout_at: now, + }); + } + + let events = machine.poll_timeouts_at(now).unwrap(); + assert!(matches!( + events.as_slice(), + [ + ClientEvent::EchoLoss { seq: 2, .. }, + ClientEvent::EchoLoss { seq: 1, .. }, + ClientEvent::EchoLoss { seq: 0, .. }, + ] + )); + assert_eq!(active(&machine).pending.len(), 0); + assert_eq!(active(&machine).timed_out.len(), 3); + } + #[test] fn wrapping_sequence_from_max_to_zero_remains_valid() { let mut machine = open_machine(2, Duration::from_secs(1));