From b5d0b1e52efc69e5ba9b42dcfa17b033e11a51eb Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 00:53:26 +0700 Subject: [PATCH] refactor(probe): preflight pending insertion --- crates/irtt-client/src/error.rs | 16 +++- crates/irtt-client/src/probe.rs | 89 ++++++++++++++++++++--- crates/irtt-client/src/session/machine.rs | 4 +- 3 files changed, 95 insertions(+), 14 deletions(-) diff --git a/crates/irtt-client/src/error.rs b/crates/irtt-client/src/error.rs index ace8630..9aa8fe3 100644 --- a/crates/irtt-client/src/error.rs +++ b/crates/irtt-client/src/error.rs @@ -1,4 +1,4 @@ -use std::{io, time::Duration}; +use std::{collections::TryReserveError, io, time::Duration}; use thiserror::Error; @@ -146,6 +146,20 @@ pub enum ClientError { /// sending again. #[error("pending probe limit exceeded ({limit})")] PendingLimitExceeded { limit: usize }, + /// A wrapping wire sequence is still represented by an in-flight probe. + /// + /// The existing probe must be replied to or timed out before that sequence + /// can be reused. + #[error("wire sequence {seq} is still pending")] + PendingSequenceCollision { seq: u32 }, + /// Reserving memory required by an operation failed before any datagram was + /// sent or logical protocol state was changed. + #[error("failed to reserve memory for {operation}: {source}")] + AllocationFailed { + operation: &'static str, + #[source] + source: TryReserveError, + }, /// Client configuration failed validation before the requested operation. /// /// The reported reason names a value that must be changed before retrying. diff --git a/crates/irtt-client/src/probe.rs b/crates/irtt-client/src/probe.rs index b26e20b..26ae1b2 100644 --- a/crates/irtt-client/src/probe.rs +++ b/crates/irtt-client/src/probe.rs @@ -1,5 +1,5 @@ use std::{ - collections::{HashMap, HashSet, VecDeque}, + collections::{HashMap, HashSet, TryReserveError, VecDeque}, time::Instant, }; @@ -26,19 +26,24 @@ impl PendingMap { } } - pub fn check_capacity(&self) -> Result<(), ClientError> { + pub fn preflight_insert(&mut self, wire_seq: u32) -> Result<(), ClientError> { + if self.map.contains_key(&wire_seq) { + return Err(ClientError::PendingSequenceCollision { seq: wire_seq }); + } if self.map.len() >= self.max_capacity { return Err(ClientError::PendingLimitExceeded { limit: self.max_capacity, }); } + if self.map.len() == self.map.capacity() { + self.map.try_reserve(1).map_err(pending_allocation_failed)?; + } Ok(()) } - pub fn insert(&mut self, probe: PendingProbe) -> Result<(), ClientError> { - self.check_capacity()?; - self.map.insert(probe.wire_seq, probe); - Ok(()) + pub fn commit_insert(&mut self, probe: PendingProbe) { + let replaced = self.map.insert(probe.wire_seq, probe); + debug_assert!(replaced.is_none(), "preflight rejected pending collision"); } pub fn remove(&mut self, wire_seq: u32) -> Option { @@ -65,6 +70,23 @@ impl PendingMap { pub fn len(&self) -> usize { self.map.len() } + + #[cfg(test)] + fn capacity(&self) -> usize { + self.map.capacity() + } + + #[cfg(test)] + fn contains(&self, wire_seq: u32) -> bool { + self.map.contains_key(&wire_seq) + } +} + +fn pending_allocation_failed(source: TryReserveError) -> ClientError { + ClientError::AllocationFailed { + operation: "pending probe storage", + source, + } } #[derive(Debug)] @@ -197,16 +219,61 @@ mod tests { fn pending_map_rejects_over_capacity() { let mut map = PendingMap::new(2); let now = Instant::now(); - map.insert(pending(0, now + Duration::from_secs(4))) - .unwrap(); - map.insert(pending(1, now + Duration::from_secs(4))) - .unwrap(); + map.preflight_insert(0).unwrap(); + map.commit_insert(pending(0, now + Duration::from_secs(4))); + map.preflight_insert(1).unwrap(); + map.commit_insert(pending(1, now + Duration::from_secs(4))); assert!(matches!( - map.insert(pending(2, now + Duration::from_secs(4))), + map.preflight_insert(2), Err(ClientError::PendingLimitExceeded { limit: 2 }) )); } + #[test] + fn pending_map_preflight_reserves_without_changing_contents() { + let mut map = PendingMap::new(2); + let initial_capacity = map.capacity(); + + map.preflight_insert(7).unwrap(); + let reserved_capacity = map.capacity(); + assert!(reserved_capacity >= initial_capacity); + assert_eq!(map.len(), 0); + assert!(!map.contains(7)); + + map.preflight_insert(7).unwrap(); + assert_eq!(map.capacity(), reserved_capacity); + assert_eq!(map.len(), 0); + assert!(!map.contains(7)); + } + + #[test] + fn pending_map_preflight_rejects_sequence_collision() { + let mut map = PendingMap::new(2); + let now = Instant::now(); + map.preflight_insert(9).unwrap(); + map.commit_insert(pending(9, now)); + + assert!(matches!( + map.preflight_insert(9), + Err(ClientError::PendingSequenceCollision { seq: 9 }) + )); + assert_eq!(map.len(), 1); + } + + #[test] + fn pending_map_allocation_failure_maps_to_dedicated_error() { + let mut map = HashMap::::new(); + let source = map.try_reserve(usize::MAX).unwrap_err(); + + assert!(matches!( + pending_allocation_failed(source), + ClientError::AllocationFailed { + operation: "pending probe storage", + .. + } + )); + } + #[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 86d862b..ba56f41 100644 --- a/crates/irtt-client/src/session/machine.rs +++ b/crates/irtt-client/src/session/machine.rs @@ -211,7 +211,7 @@ impl SessionMachine { MachineState::NoTestCompleted => return Err(ClientError::AlreadyCompleted), }; - session.pending.check_capacity()?; + session.pending.preflight_insert(session.next_wire_seq)?; let wire_seq = session.next_wire_seq; @@ -232,7 +232,7 @@ impl SessionMachine { .checked_add(probe_timeout) .ok_or(ClientError::DurationOverflow)?, }; - session.pending.insert(pending)?; + session.pending.commit_insert(pending); session.next_wire_seq = session.next_wire_seq.wrapping_add(1); session.packets_sent =