fix(managed): bound timeout work per poll

crates/irtt-client/src/async_client.rs

@@ -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

@@ -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) {
+                    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) {
                     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

@@ -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

@@ -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 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 {
                 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> {