Merge pull request #55 from Lochnair/fix/managed-timeout-budget

fix(managed): bound timeout work per poll
This commit was merged in pull request #55.
This commit is contained in:
Nils Andreas Svee authored and GitHub committed 2026-08-10 11:14:52 +07:00
commit 8c31f14419
7 files changed
+1692 -67

No files matched your search

+79 -2
View File
@@ -425,8 +425,26 @@ fn finite_multi_target_peer_close_is_accepted_as_completion() {
#[test] #[test]
fn continuous_all_peer_closed_targets_exit_nonzero() { fn continuous_all_peer_closed_targets_exit_nonzero() {
let params = test_params(None, Duration::from_millis(10)); let params = test_params(None, Duration::from_millis(10));
let a = start_peer_close_server(params.clone()); let (a_seen_tx, a_seen_rx) = mpsc::channel();
let b = start_peer_close_server(params); let (a_release_tx, a_release_rx) = mpsc::channel();
let (b_seen_tx, b_seen_rx) = mpsc::channel();
let (b_release_tx, b_release_rx) = mpsc::channel();
let a = start_synchronized_peer_close_server(params.clone(), a_seen_tx, a_release_rx);
let b = start_synchronized_peer_close_server(params, b_seen_tx, b_release_rx);
let coordinator = thread::spawn(move || {
a_seen_rx
.recv_timeout(Duration::from_secs(2))
.expect("timed out waiting for target A's first probe");
b_seen_rx
.recv_timeout(Duration::from_secs(2))
.expect("timed out waiting for target B's first probe");
a_release_tx
.send(())
.expect("target A stopped waiting for peer-close release");
b_release_tx
.send(())
.expect("target B stopped waiting for peer-close release");
});
let mut command = Command::new(env!("CARGO_BIN_EXE_irtt-cli")); let mut command = Command::new(env!("CARGO_BIN_EXE_irtt-cli"));
command.args([ command.args([
"--duration", "--duration",
@@ -446,6 +464,7 @@ fn continuous_all_peer_closed_targets_exit_nonzero() {
]); ]);
let output = run_with_timeout(command, Duration::from_secs(3)); let output = run_with_timeout(command, Duration::from_secs(3));
coordinator.join().unwrap();
a.join(); a.join();
b.join(); b.join();
@@ -860,6 +879,55 @@ fn start_peer_close_server(params: Params) -> FakeServer {
start_gated_peer_close_server(params, Vec::new()) start_gated_peer_close_server(params, Vec::new())
} }
fn start_synchronized_peer_close_server(
params: Params,
first_probe_seen_tx: mpsc::Sender<()>,
release_peer_close_rx: mpsc::Receiver<()>,
) -> FakeServer {
let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
socket
.set_read_timeout(Some(Duration::from_secs(2)))
.unwrap();
let addr = socket.local_addr().unwrap();
let done = thread::spawn(move || {
let (_, peer) = recv_request(&socket);
socket
.send_to(&open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params), peer)
.unwrap();
let (packet, peer) = recv_request(&socket);
let flags = packet[3];
assert_eq!(
flags, 0,
"target closed before synchronized first peer-close probe"
);
assert_eq!(
packet.len(),
echo_packet_len(false, &params).unwrap(),
"expected a complete echo probe before synchronized peer-close reply"
);
let seq = u32::from_le_bytes(packet[12..16].try_into().unwrap());
first_probe_seen_tx
.send(())
.expect("peer-close coordinator stopped waiting for first probe");
release_peer_close_rx
.recv_timeout(Duration::from_secs(2))
.expect("timed out waiting to release synchronized peer-close reply");
socket
.send_to(
&echo_reply_packet(
TOKEN,
seq,
&params,
&TimestampFields::default(),
FLAG_REPLY | FLAG_CLOSE,
),
peer,
)
.unwrap();
});
FakeServer { addr, done }
}
fn start_gated_peer_close_server( fn start_gated_peer_close_server(
params: Params, params: Params,
release_gates: Vec<mpsc::Receiver<()>>, release_gates: Vec<mpsc::Receiver<()>>,
@@ -875,10 +943,19 @@ fn start_gated_peer_close_server(
.send_to(&open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params), peer) .send_to(&open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params), peer)
.unwrap(); .unwrap();
let (packet, peer) = recv_request(&socket); let (packet, peer) = recv_request(&socket);
let flags = packet[3];
if flags & FLAG_CLOSE != 0 {
return;
}
for gate in release_gates { for gate in release_gates {
gate.recv_timeout(Duration::from_secs(2)) gate.recv_timeout(Duration::from_secs(2))
.expect("timed out waiting to release peer-close reply"); .expect("timed out waiting to release peer-close reply");
} }
assert_eq!(
packet.len(),
echo_packet_len(false, &params).unwrap(),
"expected a complete echo probe before peer-close reply"
);
let seq = u32::from_le_bytes(packet[12..16].try_into().unwrap()); let seq = u32::from_le_bytes(packet[12..16].try_into().unwrap());
socket socket
.send_to( .send_to(
+22 -1
View File
@@ -27,7 +27,7 @@ use crate::{
receive::try_recv_tokio_datagram, receive::try_recv_tokio_datagram,
session::machine::{ session::machine::{
recv_buffer_size, OpenDatagramDisposition, PreparedOpenAcceptance, PreparedOpenRequest, recv_buffer_size, OpenDatagramDisposition, PreparedOpenAcceptance, PreparedOpenRequest,
PreparedProbe, SessionMachine, MAX_OPEN_PACKET_SIZE, PreparedProbe, SessionMachine, TimeoutBatch, MAX_OPEN_PACKET_SIZE,
}, },
socket::{connect_tokio_udp_socket, resolve_remote_tokio, validate_open_timeouts}, socket::{connect_tokio_udp_socket, resolve_remote_tokio, validate_open_timeouts},
socket_options::{apply_dscp_to_tokio_socket, clear_dscp_on_tokio_socket}, socket_options::{apply_dscp_to_tokio_socket, clear_dscp_on_tokio_socket},
@@ -300,6 +300,14 @@ impl AsyncClient {
self.machine.poll_timeouts_at(now) self.machine.poll_timeouts_at(now)
} }
pub(crate) fn poll_timeouts_bounded_at(
&mut self,
now: Instant,
limit: usize,
) -> Result<TimeoutBatch, ClientError> {
self.machine.poll_timeouts_bounded_at(now, limit)
}
/// Send the retained close packet and commit local close exactly once. /// Send the retained close packet and commit local close exactly once.
/// ///
/// DSCP is cleared only after write readiness and immediately before the /// DSCP is cleared only after write readiness and immediately before the
@@ -765,6 +773,19 @@ 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)]
pub(crate) fn replace_pending_for_test(&mut self, probe: crate::probe::PendingProbe) {
self.machine.replace_pending_for_test(probe);
}
pub(crate) fn skip_missed_probe_slots_at(&mut self, now: Instant) -> Result<(), ClientError> { pub(crate) fn skip_missed_probe_slots_at(&mut self, now: Instant) -> Result<(), ClientError> {
let Some(schedule) = self.schedule.as_mut() else { let Some(schedule) = self.schedule.as_mut() else {
return Ok(()); return Ok(());
+33 -1
View File
@@ -16,7 +16,10 @@ use irtt_proto::{
use tokio::runtime::{Builder, Runtime}; use tokio::runtime::{Builder, Runtime};
use super::*; use super::*;
use crate::{socket_options::tokio_socket_traffic_class, Client, RunMode, SocketConfig}; use crate::{
probe::PendingProbe, socket_options::tokio_socket_traffic_class, Client, ClientTimestamp,
RunMode, SocketConfig,
};
const TOKEN: u64 = 0x1234_5678_90ab_cdef; const TOKEN: u64 = 0x1234_5678_90ab_cdef;
@@ -708,6 +711,35 @@ fn probe_false_readiness_resamples_all_timestamps() {
assert_eq!(server.finish().len(), 2); assert_eq!(server.finish().len(), 2);
} }
#[test]
fn public_timeout_polling_remains_exhaustive() {
let server = start_open_server(None, 0);
runtime().block_on(async {
let mut client = opened_client(server.addr, 0).await;
let now = Instant::now();
for seq in 0..3 {
client.replace_pending_for_test(PendingProbe {
wire_seq: seq,
sent_at: ClientTimestamp {
mono: now - Duration::from_secs(2),
wall: SystemTime::UNIX_EPOCH,
},
timeout_at: now - Duration::from_secs(1),
});
}
assert!(matches!(
client.poll_timeouts_at(now).unwrap().as_slice(),
[
ClientEvent::EchoLoss { seq: 0, .. },
ClientEvent::EchoLoss { seq: 1, .. },
ClientEvent::EchoLoss { seq: 2, .. },
]
));
});
assert_eq!(server.finish().len(), 1);
}
#[test] #[test]
fn successful_short_probe_commits_before_length_error() { fn successful_short_probe_commits_before_length_error() {
let server = start_open_server(None, 0); let server = start_open_server(None, 0);
+222 -38
View File
@@ -37,6 +37,7 @@ use super::{
const TARGET_WORK_BUDGET: usize = 128; const TARGET_WORK_BUDGET: usize = 128;
const COMMAND_WORK_BUDGET: usize = 32; const COMMAND_WORK_BUDGET: usize = 32;
const TIMEOUT_WORK_BUDGET: usize = 128;
const POST_DEADLINE_RECEIVE_BUDGET: usize = TARGET_WORK_BUDGET; const POST_DEADLINE_RECEIVE_BUDGET: usize = TARGET_WORK_BUDGET;
// Tokio 1.53 broadcast asserts this exact bound before rounding capacity to a // Tokio 1.53 broadcast asserts this exact bound before rounding capacity to a
// power of two; Tokio does not expose it as a public constant. // power of two; Tokio does not expose it as a public constant.
@@ -415,6 +416,12 @@ struct TargetRuntime {
state: TargetState, state: TargetState,
} }
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct TimeoutBacklogEntry {
index: usize,
generation: u64,
}
struct PlannedTarget { struct PlannedTarget {
target: ManagedTargetConfig, target: ManagedTargetConfig,
client_config: ClientConfig, client_config: ClientConfig,
@@ -573,6 +580,8 @@ pub struct ManagedClientTask {
wake: Option<WakeFuture>, wake: Option<WakeFuture>,
timer: Option<Pin<Box<Sleep>>>, timer: Option<Pin<Box<Sleep>>>,
cursor: usize, cursor: usize,
timeout_cursor: usize,
timeout_backlog: VecDeque<TimeoutBacklogEntry>,
scan_remaining: usize, scan_remaining: usize,
send_cursor: usize, send_cursor: usize,
burst_remaining: usize, burst_remaining: usize,
@@ -589,6 +598,10 @@ pub struct ManagedClientTask {
stagger_observations: Option<StaggerObservations>, stagger_observations: Option<StaggerObservations>,
#[cfg(test)] #[cfg(test)]
drain_test_hook: DrainTestHook, drain_test_hook: DrainTestHook,
#[cfg(test)]
timeout_inspections: usize,
#[cfg(test)]
deadline_inspections: usize,
} }
#[cfg(test)] #[cfg(test)]
@@ -1046,9 +1059,11 @@ impl ManagedClientTask {
let len = self.targets.len(); let len = self.targets.len();
if len == 0 { if len == 0 {
self.cursor = 0; self.cursor = 0;
self.timeout_cursor = 0;
self.send_cursor = 0; self.send_cursor = 0;
} else { } else {
self.cursor %= len; self.cursor %= len;
self.timeout_cursor %= len;
self.send_cursor %= len; self.send_cursor %= len;
} }
self.scan_remaining = self.scan_remaining.min(len); self.scan_remaining = self.scan_remaining.min(len);
@@ -1205,19 +1220,8 @@ impl ManagedClientTask {
.next_probe_timeout_deadline() .next_probe_timeout_deadline()
.is_some_and(|deadline| deadline <= now) .is_some_and(|deadline| deadline <= now)
{ {
match client.poll_timeouts_at(now) { self.targets[index].state = TargetState::Active { client };
Ok(events) => self.publish_client_events(index, events), return true;
Err(error) => {
return self.begin_open_session_failure(
index,
client,
ManagedTargetFailurePhase::Timing,
error,
now,
OpenSessionFailureCleanup::Close,
);
}
}
} }
let received = match client.poll_recv(cx) { let received = match client.poll_recv(cx) {
Poll::Pending => false, Poll::Pending => false,
@@ -1277,35 +1281,26 @@ impl ManagedClientTask {
return false; return false;
} }
if client
.next_probe_timeout_deadline()
.is_some_and(|timeout| timeout <= now)
{
self.targets[index].state = TargetState::Draining {
client,
drain_started_at,
deadline,
primary_end,
cleanup_failure,
post_deadline_receives_remaining,
};
return true;
}
#[cfg(test)] #[cfg(test)]
let injected_receive_failure = self.take_drain_receive_failure(); let injected_receive_failure = self.take_drain_receive_failure();
#[cfg(not(test))] #[cfg(not(test))]
let injected_receive_failure = None; let injected_receive_failure = None;
let mut retained_state_changed = false; let mut retained_state_changed = false;
if client
.next_probe_timeout_deadline()
.is_some_and(|timeout| timeout <= now)
{
match client.poll_timeouts_at(now) {
Ok(events) => {
self.publish_client_events(index, events);
retained_state_changed = true;
}
Err(error) => {
self.targets[index].counters.packets_sent = client.packets_sent();
cleanup_failure.get_or_insert_with(|| {
classify_client_error(ManagedTargetFailurePhase::Timing, &error)
});
return self.begin_close(
index,
client,
primary_end,
cleanup_failure,
now,
);
}
}
}
let received = match client.poll_recv(cx) { let received = match client.poll_recv(cx) {
Poll::Pending => false, Poll::Pending => false,
Poll::Ready(Ok(events)) => { Poll::Ready(Ok(events)) => {
@@ -1595,6 +1590,14 @@ impl ManagedClientTask {
self.targets[index].state = state; self.targets[index].state = state;
return SendResult::NotAttempted; return SendResult::NotAttempted;
}; };
if client
.next_probe_timeout_deadline()
.is_some_and(|deadline| deadline <= now)
{
self.targets[index].send_waiting = false;
self.targets[index].state = TargetState::Active { client };
return SendResult::NotAttempted;
}
if client if client
.next_send_deadline() .next_send_deadline()
.is_none_or(|deadline| deadline > now) .is_none_or(|deadline| deadline > now)
@@ -1728,13 +1731,176 @@ impl ManagedClientTask {
immediate || self.scan_remaining > 0 immediate || self.scan_remaining > 0
} }
fn poll_timeout_pass(&mut self, now: Instant) -> bool {
let target_len = self.targets.len();
if target_len == 0 {
return false;
}
let mut transitions_remaining = TIMEOUT_WORK_BUDGET;
for _ in 0..target_len.min(TIMEOUT_WORK_BUDGET) {
if transitions_remaining == 0 {
break;
}
let index = self.timeout_cursor;
self.timeout_cursor = (self.timeout_cursor + 1) % target_len;
#[cfg(test)]
{
self.timeout_inspections += 1;
}
if self.target_has_due_timeout(index, now) {
let generation = self.targets[index].instance.generation;
let step = self.poll_target_timeout(index, now);
transitions_remaining -= 1;
if step.more_due {
self.enqueue_timeout_backlog(index, generation);
}
}
}
while transitions_remaining > 0 {
let Some(entry) = self.timeout_backlog.pop_front() else {
break;
};
let valid = self
.targets
.get(entry.index)
.is_some_and(|target| target.instance.generation == entry.generation)
&& self.target_has_due_timeout(entry.index, now);
if !valid {
continue;
}
let step = self.poll_target_timeout(entry.index, now);
transitions_remaining -= 1;
if step.more_due {
self.enqueue_timeout_backlog(entry.index, entry.generation);
}
}
!self.timeout_backlog.is_empty()
}
fn enqueue_timeout_backlog(&mut self, index: usize, generation: u64) {
if self.timeout_backlog.len() == TIMEOUT_WORK_BUDGET
|| self
.timeout_backlog
.iter()
.any(|entry| entry.index == index && entry.generation == generation)
{
return;
}
self.timeout_backlog
.push_back(TimeoutBacklogEntry { index, generation });
}
fn target_has_due_timeout(&self, index: usize, now: Instant) -> bool {
match &self.targets[index].state {
TargetState::Active { client } => {
self.state != DriverState::Stopping
&& self.effective_retirement(index).is_none()
&& client
.next_probe_timeout_deadline()
.is_some_and(|deadline| deadline <= now)
}
TargetState::Draining {
client,
deadline: _deadline,
..
} => {
#[cfg(test)]
let defer_work = self.drain_test_hook.defer_work_until_deadline && now < *_deadline;
#[cfg(not(test))]
let defer_work = false;
!defer_work
&& client
.next_probe_timeout_deadline()
.is_some_and(|timeout| timeout <= now)
}
_ => false,
}
}
fn poll_target_timeout(&mut self, index: usize, now: Instant) -> TimeoutStep {
let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal);
match state {
TargetState::Active { mut client } => match client.poll_timeouts_bounded_at(now, 1) {
Ok(batch) => {
self.publish_client_events(index, batch.events);
self.targets[index].state = TargetState::Active { client };
TimeoutStep {
more_due: batch.more_due,
}
}
Err(error) => {
self.begin_open_session_failure(
index,
client,
ManagedTargetFailurePhase::Timing,
error,
now,
OpenSessionFailureCleanup::Close,
);
TimeoutStep::NONE
}
},
TargetState::Draining {
mut client,
drain_started_at,
deadline,
primary_end,
mut cleanup_failure,
post_deadline_receives_remaining,
} => match client.poll_timeouts_bounded_at(now, 1) {
Ok(batch) => {
self.publish_client_events(index, batch.events);
let deadline = match self.drain_deadline(&client, drain_started_at) {
Some(candidate) => deadline.min(candidate),
None => {
cleanup_failure.get_or_insert_with(duration_overflow_failure);
self.begin_close(index, client, primary_end, cleanup_failure, now);
return TimeoutStep::NONE;
}
};
self.targets[index].state = TargetState::Draining {
client,
drain_started_at,
deadline,
primary_end,
cleanup_failure,
post_deadline_receives_remaining,
};
TimeoutStep {
more_due: batch.more_due,
}
}
Err(error) => {
self.targets[index].counters.packets_sent = client.packets_sent();
cleanup_failure.get_or_insert_with(|| {
classify_client_error(ManagedTargetFailurePhase::Timing, &error)
});
self.begin_close(index, client, primary_end, cleanup_failure, now);
TimeoutStep::NONE
}
},
state => {
self.targets[index].state = state;
TimeoutStep::NONE
}
}
}
fn all_targets_terminal(&self) -> bool { fn all_targets_terminal(&self) -> bool {
self.targets self.targets
.iter() .iter()
.all(|target| matches!(target.state, TargetState::Terminal)) .all(|target| matches!(target.state, TargetState::Terminal))
} }
fn next_deadline(&self) -> Option<Instant> { fn next_deadline(&mut self) -> Option<Instant> {
#[cfg(test)]
{
self.deadline_inspections += self.targets.len();
}
let mut non_send_deadline = None; let mut non_send_deadline = None;
let mut send_deadline = None; let mut send_deadline = None;
for target in &self.targets { for target in &self.targets {
@@ -1888,6 +2054,15 @@ enum SendResult {
Failed { accepted: bool }, Failed { accepted: bool },
} }
#[derive(Clone, Copy)]
struct TimeoutStep {
more_due: bool,
}
impl TimeoutStep {
const NONE: Self = Self { more_due: false };
}
impl SendResult { impl SendResult {
fn accepted(self) -> bool { fn accepted(self) -> bool {
matches!( matches!(
@@ -1922,6 +2097,7 @@ impl Future for ManagedClientTask {
this.observe_stop(); this.observe_stop();
} }
let now = Instant::now(); let now = Instant::now();
immediate |= this.poll_timeout_pass(now);
immediate |= this.poll_target_pass(cx, now); immediate |= this.poll_target_pass(cx, now);
this.prune_undesired_terminal(); this.prune_undesired_terminal();
if this.state == DriverState::Running { if this.state == DriverState::Running {
@@ -1962,7 +2138,9 @@ impl Future for ManagedClientTask {
} }
immediate |= this.register_stop_wake(cx); immediate |= this.register_stop_wake(cx);
if !immediate {
immediate |= this.register_timer(cx); immediate |= this.register_timer(cx);
}
if immediate { if immediate {
cx.waker().wake_by_ref(); cx.waker().wake_by_ref();
} }
@@ -2129,6 +2307,8 @@ fn build_task(
wake: Some(arm_wake(wake_receiver)), wake: Some(arm_wake(wake_receiver)),
timer: None, timer: None,
cursor: 0, cursor: 0,
timeout_cursor: 0,
timeout_backlog: VecDeque::with_capacity(TIMEOUT_WORK_BUDGET),
scan_remaining: 0, scan_remaining: 0,
send_cursor: 0, send_cursor: 0,
burst_remaining: 0, burst_remaining: 0,
@@ -2145,6 +2325,10 @@ fn build_task(
stagger_observations: None, stagger_observations: None,
#[cfg(test)] #[cfg(test)]
drain_test_hook: DrainTestHook::default(), drain_test_hook: DrainTestHook::default(),
#[cfg(test)]
timeout_inspections: 0,
#[cfg(test)]
deadline_inspections: 0,
}; };
let handle = ManagedClientHandle { let handle = ManagedClientHandle {
stop, stop,
File diff suppressed because it is too large. Load diff
+307 -19
View File
@@ -12,9 +12,24 @@ pub(crate) struct PendingProbe {
pub timeout_at: Instant, pub timeout_at: Instant,
} }
#[derive(Debug)]
struct PendingEntry {
probe: PendingProbe,
previous: Option<u32>,
next: Option<u32>,
}
#[derive(Debug)]
pub(crate) struct ExpiredBatch {
pub probes: Vec<PendingProbe>,
pub more_due: bool,
}
#[derive(Debug)] #[derive(Debug)]
pub(crate) struct PendingMap { pub(crate) struct PendingMap {
map: HashMap<u32, PendingProbe>, map: HashMap<u32, PendingEntry>,
first: Option<u32>,
last: Option<u32>,
max_capacity: usize, max_capacity: usize,
} }
@@ -22,6 +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(),
first: None,
last: None,
max_capacity, max_capacity,
} }
} }
@@ -42,43 +59,140 @@ impl PendingMap {
} }
pub fn commit_insert(&mut self, probe: PendingProbe) { pub fn commit_insert(&mut self, probe: PendingProbe) {
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"); debug_assert!(replaced.is_none(), "preflight rejected pending collision");
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> {
self.map.remove(&wire_seq) 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;
} }
pub fn drain_expired(&mut self, now: Instant) -> Vec<PendingProbe> { self.assert_links_consistent();
let expired_keys: Vec<u32> = self 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(wire_seq) = self.first else {
break;
};
let timeout_at = self
.map .map
.iter() .get(&wire_seq)
.filter(|(_, probe)| probe.timeout_at <= now) .expect("pending list head remains present")
.map(|(key, _)| *key) .probe
.collect(); .timeout_at;
let mut expired = Vec::with_capacity(expired_keys.len()); if timeout_at > now {
for key in expired_keys { break;
if let Some(probe) = self.map.remove(&key) {
expired.push(probe);
} }
let probe = self
.remove(wire_seq)
.expect("pending list head remains present");
probes.push(probe);
} }
expired.sort_by_key(|p| p.sent_at.mono); let more_due = self
expired .first
.and_then(|wire_seq| self.map.get(&wire_seq))
.is_some_and(|entry| entry.probe.timeout_at <= now);
ExpiredBatch { probes, more_due }
} }
pub fn len(&self) -> usize { pub fn len(&self) -> usize {
self.map.len() self.map.len()
} }
#[cfg(feature = "tokio")] #[cfg(any(feature = "tokio", test))]
pub fn next_timeout_deadline(&self) -> Option<Instant> { pub fn next_timeout_deadline(&self) -> Option<Instant> {
self.map.values().map(|probe| probe.timeout_at).min() self.first
.and_then(|wire_seq| self.map.get(&wire_seq))
.map(|entry| entry.probe.timeout_at)
} }
#[cfg(feature = "tokio")] #[cfg(any(feature = "tokio", test))]
pub fn latest_timeout_deadline(&self) -> Option<Instant> { pub fn latest_timeout_deadline(&self) -> Option<Instant> {
self.map.values().map(|probe| probe.timeout_at).max() self.last
.and_then(|wire_seq| self.map.get(&wire_seq))
.map(|entry| entry.probe.timeout_at)
}
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)] #[cfg(test)]
@@ -90,6 +204,18 @@ impl PendingMap {
pub fn capacity(&self) -> usize { pub fn capacity(&self) -> usize {
self.map.capacity() self.map.capacity()
} }
#[cfg(test)]
pub fn linkage(&self) -> (Option<u32>, Option<u32>) {
(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))
}
} }
fn pending_allocation_failed(source: TryReserveError) -> ClientError { fn pending_allocation_failed(source: TryReserveError) -> ClientError {
@@ -244,6 +370,11 @@ mod tests {
} }
} }
fn insert(map: &mut PendingMap, probe: PendingProbe) {
map.preflight_insert(probe.wire_seq).unwrap();
map.commit_insert(probe);
}
#[test] #[test]
fn pending_map_rejects_over_capacity() { fn pending_map_rejects_over_capacity() {
let mut map = PendingMap::new(2); let mut map = PendingMap::new(2);
@@ -303,6 +434,163 @@ mod tests {
)); ));
} }
#[test]
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(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(),
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(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))
);
assert_eq!(
map.latest_timeout_deadline(),
Some(now + Duration::from_secs(3))
);
assert_eq!(map.linkage(), (Some(3), Some(3)));
assert_eq!(map.entry_links(3), Some((None, None)));
assert!(map.remove(3).is_some());
assert_eq!(map.next_timeout_deadline(), None);
assert_eq!(map.latest_timeout_deadline(), None);
assert_eq!(map.linkage(), (None, None));
}
#[test]
fn pending_map_bounded_expiration_matches_exhaustive_order() {
let now = Instant::now();
let probes = [
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());
for probe in probes.iter().cloned() {
insert(&mut bounded, probe.clone());
insert(&mut exhaustive, probe);
}
let empty = bounded.drain_expired_bounded(now + Duration::from_secs(4), 0);
assert!(empty.probes.is_empty());
assert!(empty.more_due);
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));
if !batch.more_due {
break;
}
}
let expected = exhaustive
.drain_expired_bounded(now + Duration::from_secs(4), usize::MAX)
.probes
.into_iter()
.map(|probe| probe.wire_seq)
.collect::<Vec<_>>();
assert_eq!(actual, expected);
assert_eq!(actual, vec![1, 2, 3, 4]);
assert_eq!(bounded.linkage(), (None, None));
}
#[test]
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));
assert!(map.remove(7).is_some());
insert(&mut map, pending(7, now + Duration::from_secs(1)));
assert!(map.drain_expired_bounded(now, usize::MAX).probes.is_empty());
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
.into_iter()
.map(|probe| probe.wire_seq)
.collect::<Vec<_>>(),
vec![7]
);
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]
fn bounded_probe_tracking_evicts_oldest_entries() { fn bounded_probe_tracking_evicts_oldest_entries() {
let mut set = CompletedSet::new(3); let mut set = CompletedSet::new(3);
+70 -4
View File
@@ -157,6 +157,12 @@ pub(crate) struct ProbeCommit {
next_packets_sent: u64, next_packets_sent: u64,
} }
#[derive(Debug)]
pub(crate) struct TimeoutBatch {
pub events: Vec<ClientEvent>,
pub more_due: bool,
}
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub(crate) struct ProbeSent { pub(crate) struct ProbeSent {
pub(crate) seq: u32, pub(crate) seq: u32,
@@ -408,11 +414,21 @@ impl SessionMachine {
&mut self, &mut self,
now: Instant, now: Instant,
) -> Result<Vec<ClientEvent>, ClientError> { ) -> Result<Vec<ClientEvent>, ClientError> {
let batch = self.poll_timeouts_bounded_at(now, usize::MAX)?;
debug_assert!(!batch.more_due);
Ok(batch.events)
}
pub(crate) fn poll_timeouts_bounded_at(
&mut self,
now: Instant,
limit: usize,
) -> Result<TimeoutBatch, ClientError> {
let session = self.open_session_mut()?; let session = self.open_session_mut()?;
let expired = session.pending.drain_expired(now); let expired = session.pending.drain_expired_bounded(now, limit);
let mut events = Vec::with_capacity(expired.len()); let mut events = Vec::with_capacity(expired.probes.len());
for probe in expired { for probe in expired.probes {
events.push(ClientEvent::EchoLoss { events.push(ClientEvent::EchoLoss {
seq: probe.wire_seq, seq: probe.wire_seq,
sent_at: probe.sent_at, sent_at: probe.sent_at,
@@ -421,7 +437,10 @@ impl SessionMachine {
session.timed_out.insert(probe); session.timed_out.insert(probe);
} }
Ok(events) Ok(TimeoutBatch {
events,
more_due: expired.more_due,
})
} }
pub(crate) fn prepare_close(&self) -> Result<PreparedClose<'_>, ClientError> { pub(crate) fn prepare_close(&self) -> Result<PreparedClose<'_>, ClientError> {
@@ -1019,6 +1038,25 @@ impl SessionMachine {
} }
} }
#[cfg(all(test, feature = "tokio"))]
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) {
let session = self
.open_session_mut()
.expect("test pending probes require an open session");
session.pending.remove(probe.wire_seq);
session.pending.preflight_insert(probe.wire_seq).unwrap();
session.pending.commit_insert(probe);
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@@ -1803,6 +1841,34 @@ mod tests {
assert_eq!(active(&machine).pending.len(), 0); assert_eq!(active(&machine).pending.len(), 0);
} }
#[test]
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, 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: sent_at.mono + Duration::from_secs(1),
});
}
let events = machine.poll_timeouts_at(now).unwrap();
assert!(matches!(
events.as_slice(),
[
ClientEvent::EchoLoss { seq: 2, .. },
ClientEvent::EchoLoss { seq: 1, .. },
ClientEvent::EchoLoss { seq: 0, .. },
]
));
assert_eq!(active(&machine).pending.len(), 0);
assert_eq!(active(&machine).timed_out.len(), 3);
}
#[test] #[test]
fn wrapping_sequence_from_max_to_zero_remains_valid() { fn wrapping_sequence_from_max_to_zero_remains_valid() {
let mut machine = open_machine(2, Duration::from_secs(1)); let mut machine = open_machine(2, Duration::from_secs(1));