fix(managed): bound timeout work per poll #55

Merged
Lochnair merged 10 commits from fix/managed-timeout-budget into main 2026-08-10 04:14:52 +00:00
4 changed files with 236 additions and 89 deletions
Showing only changes of commit f65998ff70 - Show all commits

No files matched your search

+8
View File
@@ -773,6 +773,14 @@ impl AsyncClient {
self.socket.poll_recv_ready(cx) self.socket.poll_recv_ready(cx)
} }
#[cfg(test)]
pub(crate) fn remove_pending_for_test(
&mut self,
wire_seq: u32,
) -> Option<crate::probe::PendingProbe> {
self.machine.remove_pending_for_test(wire_seq)
}
#[cfg(test)] #[cfg(test)]
pub(crate) fn replace_pending_for_test(&mut self, probe: crate::probe::PendingProbe) { pub(crate) fn replace_pending_for_test(&mut self, probe: crate::probe::PendingProbe) {
self.machine.replace_pending_for_test(probe); self.machine.replace_pending_for_test(probe);
+26 -10
View File
@@ -359,6 +359,20 @@ fn replace_pending_for_timeout_test(
client.replace_pending_for_test(probe); 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<Box<ManagedClientTask>>) { fn poll_task_once(runtime: &Runtime, task: &mut Pin<Box<ManagedClientTask>>) {
runtime.block_on(poll_fn(|cx| { runtime.block_on(poll_fn(|cx| {
assert!(task.as_mut().poll(cx).is_pending()); 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); wait_flag(&server.probe_seen);
let now = Instant::now(); let now = Instant::now();
replace_pending_for_timeout_test( let _ = remove_pending_for_timeout_test(&mut task, 0, 0);
&mut task,
0,
timeout_probe(0, now - Duration::from_secs(1), now),
);
for seq in 1..=u32::try_from(TIMEOUT_WORK_BUDGET).unwrap() { for seq in 1..=u32::try_from(TIMEOUT_WORK_BUDGET).unwrap() {
replace_pending_for_timeout_test( replace_pending_for_timeout_test(
&mut task, &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(); reply_gate.release();
wait_flag(&server.reply_sent); wait_flag(&server.reply_sent);
@@ -1973,11 +1988,7 @@ fn partial_timeout_batch_blocks_draining_receive_until_loss_commits() {
wait_flag(&server.probe_seen); wait_flag(&server.probe_seen);
let now = Instant::now(); let now = Instant::now();
replace_pending_for_timeout_test( let _ = remove_pending_for_timeout_test(&mut task, 0, 0);
&mut task,
0,
timeout_probe(0, now - Duration::from_secs(1), now),
);
for seq in 1..=u32::try_from(TIMEOUT_WORK_BUDGET).unwrap() { for seq in 1..=u32::try_from(TIMEOUT_WORK_BUDGET).unwrap() {
replace_pending_for_timeout_test( replace_pending_for_timeout_test(
&mut task, &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 } = let TargetState::Active { client } =
mem::replace(&mut task.targets[0].state, TargetState::Terminal) mem::replace(&mut task.targets[0].state, TargetState::Terminal)
else { else {
+193 -77
View File
@@ -1,5 +1,5 @@
use std::{ use std::{
collections::{BTreeSet, HashMap, HashSet, TryReserveError, VecDeque}, collections::{HashMap, HashSet, TryReserveError, VecDeque},
time::Instant, time::Instant,
}; };
@@ -12,21 +12,11 @@ pub(crate) struct PendingProbe {
pub timeout_at: Instant, pub timeout_at: Instant,
} }
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)] #[derive(Debug)]
struct DeadlineKey { struct PendingEntry {
timeout_at: Instant, probe: PendingProbe,
sent_at: Instant, previous: Option<u32>,
wire_seq: u32, next: Option<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)] #[derive(Debug)]
@@ -37,8 +27,9 @@ pub(crate) struct ExpiredBatch {
#[derive(Debug)] #[derive(Debug)]
pub(crate) struct PendingMap { pub(crate) struct PendingMap {
map: HashMap<u32, PendingProbe>, map: HashMap<u32, PendingEntry>,
deadlines: BTreeSet<DeadlineKey>, first: Option<u32>,
last: Option<u32>,
max_capacity: usize, max_capacity: usize,
} }
@@ -46,7 +37,8 @@ impl PendingMap {
pub fn new(max_capacity: usize) -> Self { pub fn new(max_capacity: usize) -> Self {
Self { Self {
map: HashMap::new(), map: HashMap::new(),
deadlines: BTreeSet::new(), first: None,
last: None,
max_capacity, max_capacity,
} }
} }
@@ -67,48 +59,98 @@ impl PendingMap {
} }
pub fn commit_insert(&mut self, probe: PendingProbe) { pub fn commit_insert(&mut self, probe: PendingProbe) {
let inserted = self.deadlines.insert(DeadlineKey::from(&probe)); let wire_seq = probe.wire_seq;
debug_assert!(inserted, "pending deadline index already contained probe"); let previous = self.last;
let replaced = self.map.insert(probe.wire_seq, probe); if let Some(previous) = previous {
let previous = self
.map
.get_mut(&previous)
chatgpt-codex-connector[bot] commented 2026-08-09 18:43:40 +00:00 (Migrated from github.com)
Review

P2 Badge Keep probe commits allocation-free

When the deadline index needs a new B-tree node, deadlines.insert allocates during commit_insert, but preflight_insert only fallibly reserves the hash map. This commit runs after the UDP send has already been accepted, so allocation failure now aborts or panics instead of returning ClientError::AllocationFailed before transmission, breaking the transactional send path that reserves storage during preflight. Use an index with fallible capacity reservation and reserve it alongside the map before allowing the send.

Useful? React with 👍 / 👎.

**<sub><sub>![P2 Badge](https://img.shields.io/badge/P2-yellow?style=flat)</sub></sub> Keep probe commits allocation-free** When the deadline index needs a new B-tree node, `deadlines.insert` allocates during `commit_insert`, but `preflight_insert` only fallibly reserves the hash map. This commit runs after the UDP send has already been accepted, so allocation failure now aborts or panics instead of returning `ClientError::AllocationFailed` before transmission, breaking the transactional send path that reserves storage during preflight. Use an index with fallible capacity reservation and reserve it alongside the map before allowing the send. Useful? React with 👍 / 👎.
.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"); 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<PendingProbe> { pub fn remove(&mut self, wire_seq: u32) -> Option<PendingProbe> {
let probe = self.map.remove(&wire_seq)?; let PendingEntry {
let removed = self.deadlines.remove(&DeadlineKey::from(&probe)); probe,
debug_assert!(removed, "pending deadline index lost removed probe"); previous,
self.assert_index_consistent(); 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) Some(probe)
} }
pub(crate) fn drain_expired_bounded(&mut self, now: Instant, limit: usize) -> ExpiredBatch { pub(crate) fn drain_expired_bounded(&mut self, now: Instant, limit: usize) -> ExpiredBatch {
let mut probes = Vec::with_capacity(limit.min(self.map.len())); let mut probes = Vec::with_capacity(limit.min(self.map.len()));
while probes.len() < limit { while probes.len() < limit {
let Some(key) = self.deadlines.first().copied() else { let Some(wire_seq) = self.first else {
break; 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; break;
} }
let removed_key = self
.deadlines
.take(&key)
.expect("first pending deadline remains present");
debug_assert_eq!(removed_key, key);
let probe = self let probe = self
.map .remove(wire_seq)
.remove(&key.wire_seq) .expect("pending list head remains present");
.expect("pending deadline always has a pending probe");
debug_assert_eq!(DeadlineKey::from(&probe), key);
probes.push(probe); probes.push(probe);
} }
self.assert_index_consistent();
let more_due = self let more_due = self
.deadlines .first
.first() .and_then(|wire_seq| self.map.get(&wire_seq))
.is_some_and(|key| key.timeout_at <= now); .is_some_and(|entry| entry.probe.timeout_at <= now);
ExpiredBatch { probes, more_due } ExpiredBatch { probes, more_due }
} }
@@ -118,21 +160,40 @@ impl PendingMap {
#[cfg(any(feature = "tokio", test))] #[cfg(any(feature = "tokio", test))]
pub fn next_timeout_deadline(&self) -> Option<Instant> { pub fn next_timeout_deadline(&self) -> Option<Instant> {
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))] #[cfg(any(feature = "tokio", test))]
pub fn latest_timeout_deadline(&self) -> Option<Instant> { pub fn latest_timeout_deadline(&self) -> Option<Instant> {
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) { 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!( debug_assert_eq!(
self.deadlines.len(), self.map
self.map.len(), .get(&first)
"pending deadline index cardinality diverged" .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)] #[cfg(test)]
pub fn contains(&self, wire_seq: u32) -> bool { pub fn contains(&self, wire_seq: u32) -> bool {
@@ -145,8 +206,15 @@ impl PendingMap {
} }
#[cfg(test)] #[cfg(test)]
pub fn deadline_index_len(&self) -> usize { pub fn linkage(&self) -> (Option<u32>, Option<u32>) {
self.deadlines.len() (self.first, self.last)
}
#[cfg(test)]
pub fn entry_links(&self, wire_seq: u32) -> Option<(Option<u32>, Option<u32>)> {
self.map
.get(&wire_seq)
.map(|entry| (entry.previous, entry.next))
} }
} }
@@ -367,12 +435,12 @@ mod tests {
} }
#[test] #[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 mut map = PendingMap::new(4);
let now = Instant::now(); let now = Instant::now();
insert(&mut map, pending(1, now + Duration::from_secs(3))); insert(&mut map, pending(1, now + Duration::from_secs(1)));
insert(&mut map, pending(2, now + Duration::from_secs(1))); insert(&mut map, pending(2, now + Duration::from_secs(2)));
insert(&mut map, pending(3, now + Duration::from_secs(2))); insert(&mut map, pending(3, now + Duration::from_secs(3)));
assert_eq!( assert_eq!(
map.next_timeout_deadline(), map.next_timeout_deadline(),
@@ -382,20 +450,25 @@ mod tests {
map.latest_timeout_deadline(), map.latest_timeout_deadline(),
Some(now + Duration::from_secs(3)) 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!(map.remove(3).is_some()); assert_eq!(map.entry_links(2), Some((Some(1), Some(3))));
assert_eq!( assert_eq!(map.entry_links(3), Some((Some(2), None)));
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!(map.remove(2).is_some()); 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!( assert_eq!(
map.next_timeout_deadline(), map.next_timeout_deadline(),
Some(now + Duration::from_secs(3)) Some(now + Duration::from_secs(3))
@@ -404,22 +477,23 @@ mod tests {
map.latest_timeout_deadline(), map.latest_timeout_deadline(),
Some(now + Duration::from_secs(3)) 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.next_timeout_deadline(), None);
assert_eq!(map.latest_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] #[test]
fn pending_map_bounded_expiration_matches_exhaustive_order() { fn pending_map_bounded_expiration_matches_exhaustive_order() {
let now = Instant::now(); let now = Instant::now();
let probes = [ let probes = [
pending(4, now + Duration::from_secs(4)),
pending(2, now + Duration::from_secs(2)),
pending(1, now + Duration::from_secs(1)), pending(1, now + Duration::from_secs(1)),
pending(2, now + Duration::from_secs(2)),
pending(3, now + Duration::from_secs(3)), pending(3, now + Duration::from_secs(3)),
pending(4, now + Duration::from_secs(4)),
]; ];
let mut bounded = PendingMap::new(probes.len()); let mut bounded = PendingMap::new(probes.len());
let mut exhaustive = 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); let empty = bounded.drain_expired_bounded(now + Duration::from_secs(4), 0);
assert!(empty.probes.is_empty()); assert!(empty.probes.is_empty());
assert!(empty.more_due); 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(); let mut actual = Vec::new();
loop { loop {
let batch = bounded.drain_expired_bounded(now + Duration::from_secs(4), 2); let batch = bounded.drain_expired_bounded(now + Duration::from_secs(4), 2);
actual.extend(batch.probes.into_iter().map(|probe| probe.wire_seq)); actual.extend(batch.probes.into_iter().map(|probe| probe.wire_seq));
assert_eq!(bounded.deadline_index_len(), bounded.len());
if !batch.more_due { if !batch.more_due {
break; break;
} }
@@ -451,11 +524,11 @@ mod tests {
assert_eq!(actual, expected); assert_eq!(actual, expected);
assert_eq!(actual, vec![1, 2, 3, 4]); assert_eq!(actual, vec![1, 2, 3, 4]);
assert_eq!(bounded.deadline_index_len(), 0); assert_eq!(bounded.linkage(), (None, None));
} }
#[test] #[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 mut map = PendingMap::new(1);
let now = Instant::now(); let now = Instant::now();
insert(&mut map, pending(7, now)); insert(&mut map, pending(7, now));
@@ -463,7 +536,8 @@ mod tests {
insert(&mut map, pending(7, now + Duration::from_secs(1))); insert(&mut map, pending(7, now + Duration::from_secs(1)));
assert!(map.drain_expired_bounded(now, usize::MAX).probes.is_empty()); 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!( assert_eq!(
map.drain_expired_bounded(now + Duration::from_secs(1), usize::MAX) map.drain_expired_bounded(now + Duration::from_secs(1), usize::MAX)
.probes .probes
@@ -472,7 +546,49 @@ mod tests {
.collect::<Vec<_>>(), .collect::<Vec<_>>(),
vec![7] 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<_>>(),
vec![1, 3]
);
assert_eq!(map.linkage(), (Some(4), Some(4)));
assert!(map.remove(4).is_some());
assert_eq!(map.linkage(), (None, None));
} }
#[test] #[test]
+9 -2
View File
@@ -1040,6 +1040,13 @@ impl SessionMachine {
#[cfg(all(test, feature = "tokio"))] #[cfg(all(test, feature = "tokio"))]
impl SessionMachine { impl SessionMachine {
pub(crate) fn remove_pending_for_test(&mut self, wire_seq: u32) -> Option<PendingProbe> {
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) { pub(crate) fn replace_pending_for_test(&mut self, probe: PendingProbe) {
let session = self let session = self
.open_session_mut() .open_session_mut()
@@ -1838,14 +1845,14 @@ mod tests {
fn exhaustive_timeout_polling_returns_every_due_loss() { fn exhaustive_timeout_polling_returns_every_due_loss() {
let mut machine = open_machine(4, Duration::from_secs(1)); let mut machine = open_machine(4, Duration::from_secs(1));
let now = Instant::now(); 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 sent_at = timestamp(now - Duration::from_secs(u64::from(seq) + 1));
let session = active_mut(&mut machine); let session = active_mut(&mut machine);
session.pending.preflight_insert(seq).unwrap(); session.pending.preflight_insert(seq).unwrap();
session.pending.commit_insert(PendingProbe { session.pending.commit_insert(PendingProbe {
wire_seq: seq, wire_seq: seq,
sent_at, sent_at,
timeout_at: now, timeout_at: sent_at.mono + Duration::from_secs(1),
}); });
} }