fix(client): bound managed drain deadline bookkeeping

Reuse deadline-ordered linked storage in the separate timed-out retention map so each drain lookup reads its tail and late replies unlink directly. Preserve FIFO capacity eviction, original timeouts, classification, and the managed work budget; cover bounded inspection work and retention lifecycle transitions.

crates/irtt-client/src/probe.rs

@@ -242,63 +242,61 @@ impl PendingMap {
 
 #[derive(Debug)]
 pub(crate) struct TimedOutMap {
-    map: HashMap<u32, PendingProbe>,
+    // Timeout transitions consume PendingMap's head in deadline order. Reuse
-    insertion_order: VecDeque<u32>,
+    // its linked storage in a separate map: the head owns FIFO eviction and
-    max_capacity: usize,
+    // the tail owns the latest retained timeout, even after arbitrary replies.
+    probes: PendingMap,
+    #[cfg(all(test, feature = "tokio"))]
+    deadline_inspections: std::cell::Cell<usize>,
 }
 
 impl TimedOutMap {
     pub fn new(max_capacity: usize) -> Self {
         Self {
-            map: HashMap::new(),
+            probes: PendingMap::new(max_capacity),
-            insertion_order: VecDeque::new(),
+            #[cfg(all(test, feature = "tokio"))]
-            max_capacity,
+            deadline_inspections: std::cell::Cell::new(0),
         }
     }
 
     pub fn insert(&mut self, probe: PendingProbe) {
-        if self.max_capacity == 0 {
+        if self.probes.max_capacity == 0 {
             return;
         }
-        if let std::collections::hash_map::Entry::Occupied(mut entry) =
+        // A sequence cannot be both pending and retained: committing its next
-            self.map.entry(probe.wire_seq)
+        // generation removes the old retention before adding the pending probe.
-        {
+        debug_assert!(!self.probes.map.contains_key(&probe.wire_seq));
-            entry.insert(probe);
+        if self.probes.map.len() == self.probes.max_capacity {
-            return;
+            self.probes
+                .remove(self.probes.first.expect("full retention has a head"));
         }
-        while self.map.len() >= self.max_capacity {
+        // Unlike pending admission, retention always makes room by eviction.
-            self.evict_oldest();
+        // Allocation here remains infallible as it was for the previous map.
-        }
+        self.probes.commit_insert(probe);
-        self.insertion_order.push_back(probe.wire_seq);
-        self.map.insert(probe.wire_seq, probe);
     }
 
     pub fn remove(&mut self, wire_seq: u32) -> Option<PendingProbe> {
-        let removed = self.map.remove(&wire_seq);
+        self.probes.remove(wire_seq)
-        if removed.is_some() {
-            self.insertion_order.retain(|seq| *seq != wire_seq);
-        }
-        removed
     }
 
     /// Mutable access to a still-retained timed-out probe by wire sequence,
     /// without disturbing eviction order. Used to attach an observed kernel
     /// TX timestamp that arrives after the probe has already timed out.
     pub fn get_mut(&mut self, wire_seq: u32) -> Option<&mut PendingProbe> {
-        self.map.get_mut(&wire_seq)
+        self.probes.get_mut(wire_seq)
     }
 
     #[cfg(feature = "tokio")]
     pub fn latest_timeout_deadline(&self) -> Option<Instant> {
-        self.map.values().map(|probe| probe.timeout_at).max()
+        self.probes
-    }
+            .last
-
+            .and_then(|wire_seq| self.probes.map.get(&wire_seq))
-    fn evict_oldest(&mut self) {
+            .map(|entry| {
-        while let Some(oldest_key) = self.insertion_order.pop_front() {
+                #[cfg(test)]
-            if self.map.remove(&oldest_key).is_some() {
+                self.deadline_inspections
-                break;
+                    .set(self.deadline_inspections.get() + 1);
-            }
+                entry.probe.timeout_at
-        }
+            })
     }
 }
 
@@ -427,17 +425,89 @@ mod tests {
         }
     }
 
+    // Measure deadline inspection work directly: wall-clock thresholds cannot
+    // reliably distinguish a bounded lookup from a retention scan.
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn timeout_retention_deadlines_have_bounded_lookup_work() {
+        const COUNT: u32 = 4096;
+        let start = Instant::now();
+        let deadline = |seq| start + Duration::from_micros(u64::from(seq / 2));
+        let mut pending_map = PendingMap::new(COUNT as usize);
+        let mut timed_out = TimedOutMap::new(COUNT as usize);
+        for seq in 0..COUNT {
+            pending_map.preflight_insert(seq).unwrap();
+            pending_map.commit_insert(pending(seq, deadline(seq)));
+        }
+        let last_deadline = deadline(COUNT - 1);
+        for seq in 0..COUNT {
+            let batch = pending_map.drain_expired_bounded(last_deadline, 1);
+            assert_eq!(batch.probes.len(), 1);
+            assert_eq!(batch.probes[0].wire_seq, seq);
+            assert_eq!(batch.more_due, seq + 1 < COUNT);
+            timed_out.insert(batch.probes.into_iter().next().unwrap());
+            assert_eq!(timed_out.latest_timeout_deadline(), Some(deadline(seq)));
+            assert_eq!(
+                pending_map
+                    .latest_timeout_deadline()
+                    .into_iter()
+                    .chain(timed_out.latest_timeout_deadline())
+                    .max(),
+                Some(last_deadline)
+            );
+        }
+        assert!(
+            timed_out.deadline_inspections.replace(0) <= 2 * COUNT as usize,
+            "draining must not scan growing timeout retention after each transition"
+        );
+        // A middle/head removal leaves the latest deadline intact. Removing
+        // the newest survivors must walk back through equal deadlines too.
+        assert!(timed_out.remove(COUNT / 2).is_some());
+        assert!(timed_out.remove(0).is_some());
+        for seq in (1..COUNT).rev().filter(|seq| *seq != COUNT / 2) {
+            assert_eq!(timed_out.latest_timeout_deadline(), Some(deadline(seq)));
+            assert!(timed_out.remove(seq).is_some());
+        }
+        assert_eq!(timed_out.latest_timeout_deadline(), None);
+        assert!(timed_out.deadline_inspections.get() <= COUNT as usize);
+    }
+
     // Public loss/reply events cannot reveal leaked retention bookkeeping.
     #[test]
-    fn timed_out_retention_stays_bounded_after_repeated_removals() {
+    fn timed_out_retention_preserves_fifo_eviction_and_stays_bounded() {
         let mut map = TimedOutMap::new(4);
         let now = Instant::now();
 
         for i in 0..20 {
             map.insert(pending(i, now));
             assert!(map.remove(i).is_some());
-            assert!(map.map.len() <= 4);
+            assert!(map.probes.map.len() <= 4);
-            assert!(map.insertion_order.len() <= 4);
+        }
+
+        // Equal deadlines across wire-sequence wrap must still evict by
+        // insertion order. Removal at any position frees capacity immediately.
+        let sequences = [u32::MAX - 1, u32::MAX, 0, 1];
+        for removed in sequences {
+            let mut map = TimedOutMap::new(4);
+            for seq in sequences {
+                map.insert(pending(seq, now));
+            }
+            assert!(map.remove(removed).is_some());
+            map.insert(pending(2, now));
+            let oldest = sequences.into_iter().find(|seq| *seq != removed).unwrap();
+            map.insert(pending(3, now));
+            assert!(map.remove(oldest).is_none(), "oldest survivor is evicted");
+            for seq in sequences
+                .into_iter()
+                .filter(|seq| *seq != removed && *seq != oldest)
+            {
+                assert!(map.remove(seq).is_some());
+            }
+            assert!(map.remove(2).is_some());
+            assert!(map.remove(3).is_some());
+            assert!(map.probes.map.is_empty());
+            #[cfg(feature = "tokio")]
+            assert_eq!(map.latest_timeout_deadline(), None);
         }
     }
 }

crates/irtt-client/src/session/machine.rs

@@ -1245,6 +1245,86 @@ mod tests {
         events.into_iter().next().unwrap()
     }
 
+    // The latest retained deadline is internal to managed draining. Exercise
+    // ordinary session transitions to check its classification/retention policy.
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn retained_deadline_tracks_timeouts_late_replies_eviction_and_close() {
+        let mut machine = opened();
+        assert_eq!(send(&mut machine).seq, 0);
+        let first = machine.latest_probe_timeout_deadline().unwrap();
+        assert_eq!(send(&mut machine).seq, 1);
+        let second = machine.latest_probe_timeout_deadline().unwrap();
+        let batch = machine.poll_timeouts_bounded_at(second, 1).unwrap();
+        assert!(
+            matches!(batch.events[0], ClientEvent::EchoLoss { seq: 0, timeout_at, .. } if timeout_at == first)
+        );
+        assert!(batch.more_due);
+        assert_eq!(machine.latest_probe_timeout_deadline(), Some(second));
+        let batch = machine.poll_timeouts_bounded_at(second, 1).unwrap();
+        assert!(
+            matches!(batch.events[0], ClientEvent::EchoLoss { seq: 1, timeout_at, .. } if timeout_at == second)
+        );
+        assert!(!batch.more_due);
+        assert!(machine.pending_is_empty());
+        assert_eq!(machine.next_probe_timeout_deadline(), None);
+        assert_eq!(machine.latest_probe_timeout_deadline(), Some(second));
+        assert!(matches!(
+            reply(&mut machine, 1),
+            ClientEvent::LateReply {
+                sent_at: Some(_),
+                rtt: Some(_),
+                ..
+            }
+        ));
+        assert_eq!(machine.latest_probe_timeout_deadline(), Some(first));
+        assert!(matches!(
+            reply(&mut machine, 1),
+            ClientEvent::DuplicateReply { .. }
+        ));
+
+        for seq in 2..4 {
+            assert_eq!(send(&mut machine).seq, seq);
+            let deadline = machine.latest_probe_timeout_deadline().unwrap();
+            machine.poll_timeouts_at(deadline).unwrap();
+            assert_eq!(machine.latest_probe_timeout_deadline(), Some(deadline));
+        }
+        // Capacity eviction removes seq 0's measurement, not late classification.
+        assert!(matches!(
+            reply(&mut machine, 0),
+            ClientEvent::LateReply {
+                sent_at: None,
+                rtt: None,
+                ..
+            }
+        ));
+        assert!(matches!(
+            reply(&mut machine, 3),
+            ClientEvent::LateReply {
+                sent_at: Some(_),
+                rtt: Some(_),
+                ..
+            }
+        ));
+        assert!(machine.latest_probe_timeout_deadline().is_some());
+        assert!(matches!(
+            reply(&mut machine, 2),
+            ClientEvent::LateReply {
+                sent_at: Some(_),
+                rtt: Some(_),
+                ..
+            }
+        ));
+        assert_eq!(machine.latest_probe_timeout_deadline(), None);
+
+        assert_eq!(send(&mut machine).seq, 4);
+        let deadline = machine.next_probe_timeout_deadline().unwrap();
+        machine.poll_timeouts_at(deadline).unwrap();
+        let commit = machine.prepare_close().unwrap().commit;
+        machine.commit_local_close(commit, ClientTimestamp::now());
+        assert_eq!(machine.latest_probe_timeout_deadline(), None);
+    }
+
     // Advancing a real client through 2^32 submissions is unreasonable. Only
     // move the send/receive sequence epoch; construct completion/timeout history through
     // ordinary machine transitions and assert emitted behavior on reuse.
@@ -1278,6 +1358,8 @@ mod tests {
                 reply(&mut machine, 0),
                 ClientEvent::EchoReply { seq: 0, .. }
             ));
+            #[cfg(feature = "tokio")]
+            assert_eq!(machine.latest_probe_timeout_deadline(), None);
             assert!(matches!(
                 reply(&mut machine, 0),
                 ClientEvent::DuplicateReply { seq: 0, .. }