refactor(probe): preflight pending insertion

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.

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<PendingProbe> {
@@ -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::<u32, PendingProbe>::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);

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 =