Lochnair-Labs/irtt-rs
2
0 Fork 0
Code Issues 1 Pull Requests 4 Actions 8 Packages Projects Releases 6 Wiki Activity

fix(managed): bound timeout work per poll

Browse Source
This commit is contained in:
Nils Andreas Svee committed 2026-08-10 01:05:41 +07:00
1 parent 9769a1b49e
commit a67331a297
4 files changed
+217 -51
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/irtt-client/src/async_client.rs
+9 -1
View File
@@ -27,7 +27,7 @@ use crate::{
receive::try_recv_tokio_datagram,
session::machine::{
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_options::{apply_dscp_to_tokio_socket, clear_dscp_on_tokio_socket},
@@ -300,6 +300,14 @@ impl AsyncClient {
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.
///
/// DSCP is cleared only after write readiness and immediately before the
crates/irtt-client/src/managed/task.rs
+182 -37
View File
@@ -37,6 +37,7 @@ use super::{
const TARGET_WORK_BUDGET: usize = 128;
const COMMAND_WORK_BUDGET: usize = 32;
const TIMEOUT_WORK_BUDGET: usize = 128;
const POST_DEADLINE_RECEIVE_BUDGET: usize = TARGET_WORK_BUDGET;
// 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.
@@ -573,6 +574,7 @@ pub struct ManagedClientTask {
wake: Option<WakeFuture>,
timer: Option<Pin<Box<Sleep>>>,
cursor: usize,
timeout_cursor: usize,
scan_remaining: usize,
send_cursor: usize,
burst_remaining: usize,
@@ -1046,9 +1048,11 @@ impl ManagedClientTask {
let len = self.targets.len();
if len == 0 {
self.cursor = 0;
self.timeout_cursor = 0;
self.send_cursor = 0;
} else {
self.cursor %= len;
self.timeout_cursor %= len;
self.send_cursor %= len;
}
self.scan_remaining = self.scan_remaining.min(len);
@@ -1205,19 +1209,8 @@ impl ManagedClientTask {
.next_probe_timeout_deadline()
.is_some_and(|deadline| deadline <= now)
{
match client.poll_timeouts_at(now) {
Ok(events) => self.publish_client_events(index, events),
Err(error) => {
return self.begin_open_session_failure(
index,
client,
ManagedTargetFailurePhase::Timing,
error,
now,
OpenSessionFailureCleanup::Close,
);
}
}
self.targets[index].state = TargetState::Active { client };
return true;
}
let received = match client.poll_recv(cx) {
Poll::Pending => false,
@@ -1277,35 +1270,26 @@ impl ManagedClientTask {
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)]
let injected_receive_failure = self.take_drain_receive_failure();
#[cfg(not(test))]
let injected_receive_failure = None;
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) {
Poll::Pending => false,
Poll::Ready(Ok(events)) => {
@@ -1728,6 +1712,152 @@ impl ManagedClientTask {
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 due_targets = Vec::new();
for offset in 0..target_len {
let index = (self.timeout_cursor + offset) % target_len;
if self.target_has_due_timeout(index, now) {
due_targets.push(index);
}
}
if due_targets.is_empty() {
return false;
}
let mut remaining = TIMEOUT_WORK_BUDGET;
let mut position = 0;
let mut last_processed = None;
while remaining > 0 && !due_targets.is_empty() {
if position == due_targets.len() {
position = 0;
}
let index = due_targets[position];
let step = self.poll_target_timeout(index, now);
if step.expired == 0 {
due_targets.remove(position);
continue;
}
remaining -= step.expired;
last_processed = Some(index);
if step.more_due {
position += 1;
} else {
due_targets.remove(position);
}
}
if let Some(index) = last_processed {
self.timeout_cursor = (index + 1) % target_len;
}
!due_targets.is_empty()
}
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) => {
let expired = batch.events.len();
self.publish_client_events(index, batch.events);
self.targets[index].state = TargetState::Active { client };
TimeoutStep {
expired,
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) => {
let expired = batch.events.len();
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 {
expired,
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 {
self.targets
.iter()
@@ -1888,6 +2018,19 @@ enum SendResult {
Failed { accepted: bool },
}
#[derive(Clone, Copy)]
struct TimeoutStep {
expired: usize,
more_due: bool,
}
impl TimeoutStep {
const NONE: Self = Self {
expired: 0,
more_due: false,
};
}
impl SendResult {
fn accepted(self) -> bool {
matches!(
@@ -1922,6 +2065,7 @@ impl Future for ManagedClientTask {
this.observe_stop();
}
let now = Instant::now();
immediate |= this.poll_timeout_pass(now);
immediate |= this.poll_target_pass(cx, now);
this.prune_undesired_terminal();
if this.state == DriverState::Running {
@@ -2129,6 +2273,7 @@ fn build_task(
wake: Some(arm_wake(wake_receiver)),
timer: None,
cursor: 0,
timeout_cursor: 0,
scan_remaining: 0,
send_cursor: 0,
burst_remaining: 0,
crates/irtt-client/src/probe.rs
+5 -9
View File
@@ -82,12 +82,6 @@ impl PendingMap {
Some(probe)
}
pub fn drain_expired(&mut self, now: Instant) -> Vec<PendingProbe> {
let batch = self.drain_expired_bounded(now, usize::MAX);
debug_assert!(!batch.more_due);
batch.probes
}
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 {
@@ -447,7 +441,8 @@ mod tests {
}
}
let expected = exhaustive
.drain_expired(now + Duration::from_secs(4))
.drain_expired_bounded(now + Duration::from_secs(4), usize::MAX)
.probes
.into_iter()
.map(|probe| probe.wire_seq)
.collect::<Vec<_>>();
@@ -465,10 +460,11 @@ mod tests {
assert!(map.remove(7).is_some());
insert(&mut map, pending(7, now + Duration::from_secs(1)));
assert!(map.drain_expired(now).is_empty());
assert!(map.drain_expired_bounded(now, usize::MAX).probes.is_empty());
assert_eq!(map.deadline_index_len(), map.len());
assert_eq!(
map.drain_expired(now + Duration::from_secs(1))
map.drain_expired_bounded(now + Duration::from_secs(1), usize::MAX)
.probes
.into_iter()
.map(|probe| probe.wire_seq)
.collect::<Vec<_>>(),
crates/irtt-client/src/session/machine.rs
+21 -4
View File
@@ -157,6 +157,12 @@ pub(crate) struct ProbeCommit {
next_packets_sent: u64,
}
#[derive(Debug)]
pub(crate) struct TimeoutBatch {
pub events: Vec<ClientEvent>,
pub more_due: bool,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct ProbeSent {
pub(crate) seq: u32,
@@ -408,11 +414,19 @@ impl SessionMachine {
&mut self,
now: Instant,
) -> Result<Vec<ClientEvent>, ClientError> {
Ok(self.poll_timeouts_bounded_at(now, usize::MAX)?.events)
}
pub(crate) fn poll_timeouts_bounded_at(
&mut self,
now: Instant,
limit: usize,
) -> Result<TimeoutBatch, ClientError> {
let session = self.open_session_mut()?;
let expired = session.pending.drain_expired(now);
let mut events = Vec::with_capacity(expired.len());
for probe in expired {
let expired = session.pending.drain_expired_bounded(now, limit);
let mut events = Vec::with_capacity(expired.probes.len());
for probe in expired.probes {
events.push(ClientEvent::EchoLoss {
seq: probe.wire_seq,
sent_at: probe.sent_at,
@@ -421,7 +435,10 @@ impl SessionMachine {
session.timed_out.insert(probe);
}
Ok(events)
Ok(TimeoutBatch {
events,
more_due: expired.more_due,
})
}
pub(crate) fn prepare_close(&self) -> Result<PreparedClose<'_>, ClientError> {