diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs index 2ffba62..f693b85 100644 --- a/crates/irtt-client/src/async_client.rs +++ b/crates/irtt-client/src/async_client.rs @@ -773,6 +773,14 @@ impl AsyncClient { self.socket.poll_recv_ready(cx) } + #[cfg(test)] + pub(crate) fn remove_pending_for_test( + &mut self, + wire_seq: u32, + ) -> Option { + self.machine.remove_pending_for_test(wire_seq) + } + #[cfg(test)] pub(crate) fn replace_pending_for_test(&mut self, probe: crate::probe::PendingProbe) { self.machine.replace_pending_for_test(probe); diff --git a/crates/irtt-client/src/managed/tests.rs b/crates/irtt-client/src/managed/tests.rs index 05e22ce..076ae83 100644 --- a/crates/irtt-client/src/managed/tests.rs +++ b/crates/irtt-client/src/managed/tests.rs @@ -359,6 +359,20 @@ fn replace_pending_for_timeout_test( client.replace_pending_for_test(probe); } +fn remove_pending_for_timeout_test( + task: &mut ManagedClientTask, + index: usize, + wire_seq: u32, +) -> PendingProbe { + let client = match &mut task.targets[index].state { + TargetState::Active { client } | TargetState::Draining { client, .. } => client, + _ => panic!("timeout test target must be active or draining"), + }; + client + .remove_pending_for_test(wire_seq) + .expect("timeout test pending probe must exist") +} + fn poll_task_once(runtime: &Runtime, task: &mut Pin>) { runtime.block_on(poll_fn(|cx| { assert!(task.as_mut().poll(cx).is_pending()); @@ -1880,11 +1894,7 @@ fn partial_timeout_batch_blocks_active_receive_until_loss_commits() { wait_flag(&server.probe_seen); let now = Instant::now(); - replace_pending_for_timeout_test( - &mut task, - 0, - timeout_probe(0, now - Duration::from_secs(1), now), - ); + let _ = remove_pending_for_timeout_test(&mut task, 0, 0); for seq in 1..=u32::try_from(TIMEOUT_WORK_BUDGET).unwrap() { replace_pending_for_timeout_test( &mut task, @@ -1896,6 +1906,11 @@ fn partial_timeout_batch_blocks_active_receive_until_loss_commits() { ), ); } + replace_pending_for_timeout_test( + &mut task, + 0, + timeout_probe(0, now - Duration::from_secs(1), now), + ); reply_gate.release(); wait_flag(&server.reply_sent); @@ -1973,11 +1988,7 @@ fn partial_timeout_batch_blocks_draining_receive_until_loss_commits() { wait_flag(&server.probe_seen); let now = Instant::now(); - replace_pending_for_timeout_test( - &mut task, - 0, - timeout_probe(0, now - Duration::from_secs(1), now), - ); + let _ = remove_pending_for_timeout_test(&mut task, 0, 0); for seq in 1..=u32::try_from(TIMEOUT_WORK_BUDGET).unwrap() { replace_pending_for_timeout_test( &mut task, @@ -1989,6 +2000,11 @@ fn partial_timeout_batch_blocks_draining_receive_until_loss_commits() { ), ); } + replace_pending_for_timeout_test( + &mut task, + 0, + timeout_probe(0, now - Duration::from_secs(1), now), + ); let TargetState::Active { client } = mem::replace(&mut task.targets[0].state, TargetState::Terminal) else { diff --git a/crates/irtt-client/src/probe.rs b/crates/irtt-client/src/probe.rs index 60b73a6..951fb1e 100644 --- a/crates/irtt-client/src/probe.rs +++ b/crates/irtt-client/src/probe.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeSet, HashMap, HashSet, TryReserveError, VecDeque}, + collections::{HashMap, HashSet, TryReserveError, VecDeque}, time::Instant, }; @@ -12,21 +12,11 @@ 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)] +struct PendingEntry { + probe: PendingProbe, + previous: Option, + next: Option, } #[derive(Debug)] @@ -37,8 +27,9 @@ pub(crate) struct ExpiredBatch { #[derive(Debug)] pub(crate) struct PendingMap { - map: HashMap, - deadlines: BTreeSet, + map: HashMap, + first: Option, + last: Option, max_capacity: usize, } @@ -46,7 +37,8 @@ impl PendingMap { pub fn new(max_capacity: usize) -> Self { Self { map: HashMap::new(), - deadlines: BTreeSet::new(), + first: None, + last: None, max_capacity, } } @@ -67,48 +59,98 @@ 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); + let wire_seq = probe.wire_seq; + let previous = self.last; + if let Some(previous) = previous { + let previous = self + .map + .get_mut(&previous) + .expect("pending list tail remains present"); + // probe_timeout is fixed for an open session and committed sends + // have nondecreasing monotonic timestamps, so deadlines append. + debug_assert!(previous.probe.timeout_at <= probe.timeout_at); + debug_assert!( + previous.next.is_none(), + "pending list tail has no successor" + ); + previous.next = Some(wire_seq); + } else { + debug_assert!(self.first.is_none(), "empty pending list has no head"); + self.first = Some(wire_seq); + } + + let replaced = self.map.insert( + wire_seq, + PendingEntry { + probe, + previous, + next: None, + }, + ); debug_assert!(replaced.is_none(), "preflight rejected pending collision"); - self.assert_index_consistent(); + self.last = Some(wire_seq); + self.assert_links_consistent(); } pub fn remove(&mut self, wire_seq: u32) -> Option { - 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(); + let PendingEntry { + probe, + previous, + next, + } = self.map.remove(&wire_seq)?; + + if let Some(previous) = previous { + let previous = self + .map + .get_mut(&previous) + .expect("pending list predecessor remains present"); + debug_assert_eq!(previous.next, Some(wire_seq)); + previous.next = next; + } else { + debug_assert_eq!(self.first, Some(wire_seq)); + self.first = next; + } + if let Some(next) = next { + let next = self + .map + .get_mut(&next) + .expect("pending list successor remains present"); + debug_assert_eq!(next.previous, Some(wire_seq)); + next.previous = previous; + } else { + debug_assert_eq!(self.last, Some(wire_seq)); + self.last = previous; + } + + self.assert_links_consistent(); Some(probe) } 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 { + let Some(wire_seq) = self.first else { break; }; - if key.timeout_at > now { + let timeout_at = self + .map + .get(&wire_seq) + .expect("pending list head remains present") + .probe + .timeout_at; + if 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); + .remove(wire_seq) + .expect("pending list head remains present"); probes.push(probe); } - self.assert_index_consistent(); let more_due = self - .deadlines - .first() - .is_some_and(|key| key.timeout_at <= now); + .first + .and_then(|wire_seq| self.map.get(&wire_seq)) + .is_some_and(|entry| entry.probe.timeout_at <= now); ExpiredBatch { probes, more_due } } @@ -118,20 +160,39 @@ impl PendingMap { #[cfg(any(feature = "tokio", test))] pub fn next_timeout_deadline(&self) -> Option { - self.deadlines.first().map(|key| key.timeout_at) + self.first + .and_then(|wire_seq| self.map.get(&wire_seq)) + .map(|entry| entry.probe.timeout_at) } #[cfg(any(feature = "tokio", test))] pub fn latest_timeout_deadline(&self) -> Option { - self.deadlines.last().map(|key| key.timeout_at) + self.last + .and_then(|wire_seq| self.map.get(&wire_seq)) + .map(|entry| entry.probe.timeout_at) } - fn assert_index_consistent(&self) { - debug_assert_eq!( - self.deadlines.len(), - self.map.len(), - "pending deadline index cardinality diverged" - ); + fn assert_links_consistent(&self) { + debug_assert_eq!(self.first.is_none(), self.map.is_empty()); + debug_assert_eq!(self.last.is_none(), self.map.is_empty()); + if let Some(first) = self.first { + debug_assert_eq!( + self.map + .get(&first) + .expect("pending list head remains present") + .previous, + None + ); + } + if let Some(last) = self.last { + debug_assert_eq!( + self.map + .get(&last) + .expect("pending list tail remains present") + .next, + None + ); + } } #[cfg(test)] @@ -145,8 +206,15 @@ impl PendingMap { } #[cfg(test)] - pub fn deadline_index_len(&self) -> usize { - self.deadlines.len() + pub fn linkage(&self) -> (Option, Option) { + (self.first, self.last) + } + + #[cfg(test)] + pub fn entry_links(&self, wire_seq: u32) -> Option<(Option, Option)> { + self.map + .get(&wire_seq) + .map(|entry| (entry.previous, entry.next)) } } @@ -367,12 +435,12 @@ mod tests { } #[test] - fn pending_map_indexes_deadline_extrema_and_removals() { + fn pending_map_unlinks_head_middle_and_tail_and_updates_extrema() { 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))); + insert(&mut map, pending(1, now + Duration::from_secs(1))); + insert(&mut map, pending(2, now + Duration::from_secs(2))); + insert(&mut map, pending(3, now + Duration::from_secs(3))); assert_eq!( map.next_timeout_deadline(), @@ -382,20 +450,25 @@ mod tests { 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_eq!(map.linkage(), (Some(1), Some(3))); + assert_eq!(map.entry_links(1), Some((None, Some(2)))); + assert_eq!(map.entry_links(2), Some((Some(1), Some(3)))); + assert_eq!(map.entry_links(3), Some((Some(2), None))); assert!(map.remove(2).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.linkage(), (Some(1), Some(3))); + assert_eq!(map.entry_links(1), Some((None, Some(3)))); + assert_eq!(map.entry_links(3), Some((Some(1), None))); + + assert!(map.remove(1).is_some()); assert_eq!( map.next_timeout_deadline(), Some(now + Duration::from_secs(3)) @@ -404,22 +477,23 @@ mod tests { map.latest_timeout_deadline(), Some(now + Duration::from_secs(3)) ); - assert_eq!(map.deadline_index_len(), map.len()); + assert_eq!(map.linkage(), (Some(3), Some(3))); + assert_eq!(map.entry_links(3), Some((None, None))); - assert!(map.remove(1).is_some()); + assert!(map.remove(3).is_some()); assert_eq!(map.next_timeout_deadline(), None); assert_eq!(map.latest_timeout_deadline(), None); - assert_eq!(map.deadline_index_len(), map.len()); + assert_eq!(map.linkage(), (None, None)); } #[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(2, now + Duration::from_secs(2)), pending(3, now + Duration::from_secs(3)), + pending(4, now + Duration::from_secs(4)), ]; let mut bounded = PendingMap::new(probes.len()); let mut exhaustive = PendingMap::new(probes.len()); @@ -431,13 +505,12 @@ mod tests { 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()); + assert_eq!(bounded.linkage(), (Some(1), Some(4))); 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; } @@ -451,11 +524,11 @@ mod tests { assert_eq!(actual, expected); assert_eq!(actual, vec![1, 2, 3, 4]); - assert_eq!(bounded.deadline_index_len(), 0); + assert_eq!(bounded.linkage(), (None, None)); } #[test] - fn pending_map_sequence_reuse_leaves_no_stale_deadline() { + fn pending_map_sequence_reuse_leaves_no_stale_link() { let mut map = PendingMap::new(1); let now = Instant::now(); insert(&mut map, pending(7, now)); @@ -463,7 +536,8 @@ mod tests { insert(&mut map, pending(7, now + Duration::from_secs(1))); assert!(map.drain_expired_bounded(now, usize::MAX).probes.is_empty()); - assert_eq!(map.deadline_index_len(), map.len()); + assert_eq!(map.linkage(), (Some(7), Some(7))); + assert_eq!(map.entry_links(7), Some((None, None))); assert_eq!( map.drain_expired_bounded(now + Duration::from_secs(1), usize::MAX) .probes @@ -472,7 +546,49 @@ mod tests { .collect::>(), vec![7] ); - assert_eq!(map.deadline_index_len(), map.len()); + assert_eq!(map.linkage(), (None, None)); + } + + #[test] + fn pending_map_commit_uses_only_preflighted_map_capacity() { + let mut map = PendingMap::new(2); + let now = Instant::now(); + map.preflight_insert(7).unwrap(); + let capacity = map.capacity(); + + map.commit_insert(pending(7, now)); + + assert_eq!(map.capacity(), capacity); + assert_eq!(map.len(), 1); + assert_eq!(map.linkage(), (Some(7), Some(7))); + } + + #[test] + fn pending_map_mixed_removals_preserve_links() { + let mut map = PendingMap::new(4); + let now = Instant::now(); + for seq in 1..=4 { + insert( + &mut map, + pending(seq, now + Duration::from_secs(u64::from(seq))), + ); + } + + assert!(map.remove(2).is_some()); + assert_eq!(map.entry_links(1), Some((None, Some(3)))); + assert_eq!(map.entry_links(3), Some((Some(1), Some(4)))); + let expired = map.drain_expired_bounded(now + Duration::from_secs(3), usize::MAX); + assert_eq!( + expired + .probes + .into_iter() + .map(|probe| probe.wire_seq) + .collect::>(), + vec![1, 3] + ); + assert_eq!(map.linkage(), (Some(4), Some(4))); + assert!(map.remove(4).is_some()); + assert_eq!(map.linkage(), (None, None)); } #[test] diff --git a/crates/irtt-client/src/session/machine.rs b/crates/irtt-client/src/session/machine.rs index 57f9528..7eade8b 100644 --- a/crates/irtt-client/src/session/machine.rs +++ b/crates/irtt-client/src/session/machine.rs @@ -1040,6 +1040,13 @@ impl SessionMachine { #[cfg(all(test, feature = "tokio"))] impl SessionMachine { + pub(crate) fn remove_pending_for_test(&mut self, wire_seq: u32) -> Option { + self.open_session_mut() + .expect("test pending probes require an open session") + .pending + .remove(wire_seq) + } + pub(crate) fn replace_pending_for_test(&mut self, probe: PendingProbe) { let session = self .open_session_mut() @@ -1838,14 +1845,14 @@ mod tests { 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] { + for seq in [2, 1, 0] { 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, + timeout_at: sent_at.mono + Duration::from_secs(1), }); }