diff --git a/crates/irtt-client/src/probe.rs b/crates/irtt-client/src/probe.rs index eb05402..5cba0ba 100644 --- a/crates/irtt-client/src/probe.rs +++ b/crates/irtt-client/src/probe.rs @@ -103,7 +103,11 @@ impl TimedOutMap { } pub fn remove(&mut self, wire_seq: u32) -> Option { - self.map.remove(&wire_seq) + let removed = self.map.remove(&wire_seq); + if removed.is_some() { + self.insertion_order.retain(|seq| *seq != wire_seq); + } + removed } pub fn clear(&mut self) { @@ -116,6 +120,11 @@ impl TimedOutMap { self.map.len() } + #[cfg(test)] + fn insertion_order_len(&self) -> usize { + self.insertion_order.len() + } + fn evict_oldest(&mut self) { while let Some(oldest_key) = self.insertion_order.pop_front() { if self.map.remove(&oldest_key).is_some() { @@ -257,4 +266,17 @@ mod tests { assert!(map.remove(1).is_some()); assert!(map.remove(2).is_some()); } + + #[test] + fn timed_out_map_remove_prunes_insertion_order() { + let mut map = TimedOutMap::new(4); + let now = Instant::now(); + + for i in 0..20 { + map.insert(pending(i, u64::from(i), now)); + assert!(map.remove(i).is_some()); + assert_eq!(map.len(), 0); + assert_eq!(map.insertion_order_len(), 0); + } + } }