From 99b6ddfead671157a97a765b2f6519ccffca523c Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Wed, 5 Aug 2026 00:09:46 +0700 Subject: [PATCH] fix(managed): preserve drain and pacing invariants --- crates/irtt-client/src/managed/task.rs | 210 ++++++++++++++++++------ crates/irtt-client/src/managed/tests.rs | 190 +++++++++++++++++---- 2 files changed, 317 insertions(+), 83 deletions(-) diff --git a/crates/irtt-client/src/managed/task.rs b/crates/irtt-client/src/managed/task.rs index ec3e8be..a0bf27d 100644 --- a/crates/irtt-client/src/managed/task.rs +++ b/crates/irtt-client/src/managed/task.rs @@ -171,6 +171,7 @@ enum TargetState { }, Draining { client: AsyncClient, + drain_started_at: Instant, deadline: Instant, primary_end: ManagedTargetEndReason, cleanup_failure: Option, @@ -213,7 +214,6 @@ struct TargetRuntime { server_addr: Arc, remote: Option, counters: TargetCounters, - latest_committed_timeout: Option, send_waiting: bool, active: bool, state: TargetState, @@ -302,6 +302,12 @@ enum DriverState { Finished, } +#[derive(Clone, Copy)] +enum OpenSessionFailureCleanup { + Drain, + Close, +} + /// The sole authoritative zero-or-many target driver. #[must_use = "ManagedClientTask must be awaited or deliberately dropped"] pub struct ManagedClientTask { @@ -416,11 +422,14 @@ impl ManagedClientTask { fn install_target_state(&mut self, index: usize, state: TargetState) { let lifecycle = state.lifecycle(); let active = matches!(state, TargetState::Active { .. }); - let membership_changed = self.targets[index].active != active; + let was_active = self.targets[index].active; self.targets[index].active = active; self.targets[index].state = state; - if membership_changed { - self.recompute_stagger_gate(Instant::now()); + let now = Instant::now(); + match (was_active, active) { + (false, true) => self.stagger_target_added(now), + (true, false) => self.stagger_target_removed(now), + _ => {} } self.replace_status(); self.publish_event(ManagedEvent::TargetStateChanged { @@ -493,6 +502,26 @@ impl ManagedClientTask { self.finish_target(index, ManagedTargetEndReason::Failed(failure), None); } + fn begin_open_session_failure( + &mut self, + index: usize, + mut client: AsyncClient, + phase: ManagedTargetFailurePhase, + error: ClientError, + now: Instant, + cleanup: OpenSessionFailureCleanup, + ) -> bool { + let primary_end = ManagedTargetEndReason::Failed(classify_client_error(phase, &error)); + client.discard_prepared_probe(); + self.targets[index].counters.packets_sent = client.packets_sent(); + match cleanup { + OpenSessionFailureCleanup::Drain => self.begin_drain(index, client, primary_end, now), + OpenSessionFailureCleanup::Close => { + self.begin_close(index, client, primary_end, None, now) + } + } + } + fn begin_running(&mut self) -> Result<(), ManagedDriverFailure> { tokio::runtime::Handle::try_current().map_err(|_| ManagedDriverFailure::NoTokioRuntime)?; self.state = DriverState::Running; @@ -631,9 +660,14 @@ impl ManagedClientTask { match client.poll_timeouts_at(now) { Ok(events) => self.publish_client_events(index, events), Err(error) => { - self.targets[index].counters.packets_sent = client.packets_sent(); - self.fail_target(index, ManagedTargetFailurePhase::Timing, error); - return false; + return self.begin_open_session_failure( + index, + client, + ManagedTargetFailurePhase::Timing, + error, + now, + OpenSessionFailureCleanup::Close, + ); } } } @@ -644,9 +678,14 @@ impl ManagedClientTask { true } Poll::Ready(Err(error)) => { - self.targets[index].counters.packets_sent = client.packets_sent(); - self.fail_target(index, ManagedTargetFailurePhase::Receiving, error); - return false; + return self.begin_open_session_failure( + index, + client, + ManagedTargetFailurePhase::Receiving, + error, + now, + OpenSessionFailureCleanup::Close, + ); } }; if client.is_peer_closed() { @@ -668,6 +707,7 @@ impl ManagedClientTask { } TargetState::Draining { mut client, + drain_started_at, deadline, primary_end, mut cleanup_failure, @@ -680,6 +720,7 @@ impl ManagedClientTask { if defer_work { self.targets[index].state = TargetState::Draining { client, + drain_started_at, deadline, primary_end, cleanup_failure, @@ -692,12 +733,16 @@ impl ManagedClientTask { 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), + 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(|| { @@ -717,6 +762,7 @@ impl ManagedClientTask { Poll::Pending => false, Poll::Ready(Ok(events)) => { self.publish_client_events(index, events); + retained_state_changed = true; true } Poll::Ready(Err(error)) => { @@ -739,11 +785,29 @@ impl ManagedClientTask { self.finish_target(index, ManagedTargetEndReason::PeerClosed, cleanup_failure); return false; } + let deadline = if retained_state_changed { + match self.drain_deadline(&client, drain_started_at) { + Some(candidate) => deadline.min(candidate), + None => { + cleanup_failure.get_or_insert_with(duration_overflow_failure); + return self.begin_close( + index, + client, + primary_end, + cleanup_failure, + now, + ); + } + } + } else { + deadline + }; if now >= deadline { if received && post_deadline_receives_remaining > 1 { post_deadline_receives_remaining -= 1; self.targets[index].state = TargetState::Draining { client, + drain_started_at, deadline, primary_end, cleanup_failure, @@ -756,6 +820,7 @@ impl ManagedClientTask { } else { self.targets[index].state = TargetState::Draining { client, + drain_started_at, deadline, primary_end, cleanup_failure, @@ -824,23 +889,24 @@ impl ManagedClientTask { index: usize, client: AsyncClient, primary_end: ManagedTargetEndReason, - now: Instant, + _now: Instant, ) -> bool { - let latest = self.targets[index] - .latest_committed_timeout - .into_iter() - .chain(client.latest_probe_timeout_deadline()) - .max() - .unwrap_or(now) - .max(now); - let Some(deadline) = latest.checked_add(self.config.final_drain) else { + let drain_started_at = Instant::now(); + let Some(deadline) = self.drain_deadline(&client, drain_started_at) else { let cleanup_failure = Some(duration_overflow_failure()); - return self.begin_close(index, client, primary_end, cleanup_failure, now); + return self.begin_close( + index, + client, + primary_end, + cleanup_failure, + drain_started_at, + ); }; self.install_target_state( index, TargetState::Draining { client, + drain_started_at, deadline, primary_end, cleanup_failure: None, @@ -850,6 +916,14 @@ impl ManagedClientTask { true } + fn drain_deadline(&self, client: &AsyncClient, drain_started_at: Instant) -> Option { + client + .latest_probe_timeout_deadline() + .unwrap_or(drain_started_at) + .max(drain_started_at) + .checked_add(self.config.final_drain) + } + fn begin_close( &mut self, index: usize, @@ -914,15 +988,50 @@ impl ManagedClientTask { minimum.map(|interval| (active, stagger_spacing(interval, active))) } - fn recompute_stagger_gate(&mut self, now: Instant) { - self.send_gate = self + fn stagger_target_added(&mut self, now: Instant) { + let Some(existing) = self.send_gate.filter(|gate| *gate > now) else { + self.send_gate = None; + return; + }; + let candidate = self .last_stagger_send .zip(self.active_stagger_spacing()) .and_then(|(last, (_, spacing))| last.checked_add(spacing)) .filter(|gate| *gate > now); + self.send_gate = candidate.map(|candidate| existing.min(candidate)); } - fn poll_one_send(&mut self, index: usize, cx: &mut Context<'_>, now: Instant) -> SendResult { + fn stagger_target_removed(&mut self, now: Instant) { + self.send_gate = self.send_gate.filter(|gate| *gate > now); + } + + fn record_stagger_acceptance( + &mut self, + result: SendResult, + stagger_spacing: Option<(usize, Duration)>, + accepted_at: Instant, + ) { + if !result.accepted() { + return; + } + let Some((_active, spacing)) = stagger_spacing else { + return; + }; + self.last_stagger_send = Some(accepted_at); + self.send_gate = accepted_at.checked_add(spacing); + #[cfg(test)] + if let Some(observations) = &self.stagger_observations { + observations.lock().unwrap().push((_active, spacing)); + } + } + + fn poll_one_send( + &mut self, + index: usize, + cx: &mut Context<'_>, + now: Instant, + stagger_spacing: Option<(usize, Duration)>, + ) -> SendResult { let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal); let TargetState::Active { mut client } = state else { self.targets[index].state = state; @@ -940,18 +1049,17 @@ impl ManagedClientTask { let result = client.poll_send_probe(cx); let after = client.packets_sent(); let accepted = after > before; + let send_result = match &result { + Poll::Pending => SendResult::Pending, + Poll::Ready(Ok(_)) => SendResult::Ready { accepted }, + Poll::Ready(Err(_)) => SendResult::Failed { accepted }, + }; + self.record_stagger_acceptance(send_result, stagger_spacing, Instant::now()); let schedule_error = accepted .then(|| client.skip_missed_probe_slots_at(Instant::now())) .transpose() .err(); self.targets[index].counters.packets_sent = after; - if accepted { - self.targets[index].latest_committed_timeout = client - .latest_probe_timeout_deadline() - .into_iter() - .chain(self.targets[index].latest_committed_timeout) - .max(); - } match result { Poll::Pending => { self.targets[index].send_waiting = true; @@ -963,7 +1071,19 @@ impl ManagedClientTask { self.targets[index].state = TargetState::Active { client }; self.publish_client_events(index, events); if let Some(error) = schedule_error { - self.fail_target(index, ManagedTargetFailurePhase::Timing, error); + let TargetState::Active { client } = + mem::replace(&mut self.targets[index].state, TargetState::Terminal) + else { + unreachable!("probe sender remained active while publishing send events"); + }; + self.begin_open_session_failure( + index, + client, + ManagedTargetFailurePhase::Timing, + error, + now, + OpenSessionFailureCleanup::Drain, + ); SendResult::Failed { accepted } } else { SendResult::Ready { accepted } @@ -971,7 +1091,14 @@ impl ManagedClientTask { } Poll::Ready(Err(error)) => { self.targets[index].send_waiting = false; - self.fail_target(index, ManagedTargetFailurePhase::Sending, error); + self.begin_open_session_failure( + index, + client, + ManagedTargetFailurePhase::Sending, + error, + now, + OpenSessionFailureCleanup::Drain, + ); SendResult::Failed { accepted } } } @@ -992,24 +1119,12 @@ impl ManagedClientTask { if !matches!(self.targets[index].state, TargetState::Active { .. }) { continue; } - #[cfg(test)] let stagger_spacing = self.active_stagger_spacing(); - let result = self.poll_one_send(index, cx, now); + let result = self.poll_one_send(index, cx, now, stagger_spacing); if result == SendResult::NotAttempted { continue; } self.stagger_remaining = 0; - if result.accepted() { - #[cfg(test)] - if let Some((_active, spacing)) = stagger_spacing { - if let Some(observations) = &self.stagger_observations { - observations.lock().unwrap().push((_active, spacing)); - } - } - let accepted_at = Instant::now(); - self.last_stagger_send = Some(accepted_at); - self.recompute_stagger_gate(accepted_at); - } return matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. }); } self.stagger_remaining > 0 @@ -1028,7 +1143,7 @@ impl ManagedClientTask { let index = self.send_cursor % self.targets.len(); self.send_cursor = (self.send_cursor + 1) % self.targets.len(); self.burst_remaining -= 1; - let result = self.poll_one_send(index, cx, now); + let result = self.poll_one_send(index, cx, now, None); immediate |= matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. }); } immediate || self.burst_remaining > 0 @@ -1319,7 +1434,6 @@ fn build_task( server_addr: Arc::from(target.server_addr), remote: None, counters: TargetCounters::default(), - latest_committed_timeout: None, send_waiting: false, active: false, state: TargetState::Pending { client_config }, diff --git a/crates/irtt-client/src/managed/tests.rs b/crates/irtt-client/src/managed/tests.rs index 23e11f1..bbc8e35 100644 --- a/crates/irtt-client/src/managed/tests.rs +++ b/crates/irtt-client/src/managed/tests.rs @@ -1,5 +1,6 @@ use std::{ future::{poll_fn, Future}, + mem, net::{SocketAddr, UdpSocket}, pin::Pin, sync::{Arc, Condvar, Mutex}, @@ -524,14 +525,19 @@ fn pre_poll_stop() { } #[test] -fn finite_one_target_lifecycle() { +fn finite_immediate_replies_use_retained_drain_deadline() { let server = start_server(ServerBehavior::Echo, None); - let (task, _) = ManagedClient::task( - config(ManagedPacing::Staggered), - vec![target("one", server.addr)], - ) - .unwrap(); - let outcome = runtime().block_on(task); + let mut managed = config(ManagedPacing::Staggered); + managed.client.probe_timeout = Duration::from_secs(3); + managed.final_drain = Duration::from_millis(30); + let (task, _) = ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap(); + let started_at = Instant::now(); + let outcome = runtime().block_on(async { + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("replied probes must not hold the drain until their obsolete timeout") + }); + assert!(started_at.elapsed() < Duration::from_secs(1)); let records = server.finish(); assert_eq!(outcome.end_reason, ManagedEndReason::TargetsComplete); assert!(matches!( @@ -793,7 +799,8 @@ fn staggered_pacing_uses_all_active_negotiated_intervals() { #[test] fn stagger_gate_tracks_active_membership() { - let first = start_server_negotiating(ServerBehavior::Echo, None, Some(Duration::from_secs(1))); + let first = + start_server_negotiating(ServerBehavior::Echo, None, Some(Duration::from_millis(100))); let second_gate = Arc::new(PacketGate::default()); let second = start_server_with_gates( ServerBehavior::Echo, @@ -803,7 +810,7 @@ fn stagger_gate_tracks_active_membership() { None, ); let mut managed = config(ManagedPacing::Staggered); - managed.client.duration = Some(Duration::from_secs(2)); + managed.client.duration = Some(Duration::from_millis(1_250)); managed.client.interval = Duration::from_secs(1); managed.client.negotiation_policy = NegotiationPolicy::Loose; managed.client.open_timeouts = vec![Duration::from_secs(2)]; @@ -818,19 +825,44 @@ fn stagger_gate_tracks_active_membership() { task.targets[0].counters.packets_sent == 1 }); let one_target_last = task.last_stagger_send.unwrap(); + let one_target_gate = task.send_gate.unwrap(); assert_eq!( task.send_gate, - one_target_last.checked_add(Duration::from_secs(1)) + one_target_last.checked_add(Duration::from_millis(100)) ); second_gate.release(); drive_task_until(&runtime, &mut task, |task| task.active_count() == 2); let last = task.last_stagger_send.unwrap(); let gate = task.send_gate.unwrap(); - assert_eq!(gate, last.checked_add(Duration::from_millis(500)).unwrap()); - assert!(gate > last && gate < last.checked_add(Duration::from_secs(1)).unwrap()); + assert_eq!(gate, last.checked_add(Duration::from_millis(50)).unwrap()); + assert!(gate <= one_target_gate); + assert!(gate > last && gate < last.checked_add(Duration::from_millis(100)).unwrap()); drive_task_until(&runtime, &mut task, |task| { task.targets[1].counters.packets_sent == 1 }); + + let accepted_at = Instant::now(); + task.record_stagger_acceptance( + SendResult::Failed { accepted: true }, + Some((2, Duration::from_millis(50))), + accepted_at, + ); + assert_eq!( + task.send_gate, + accepted_at.checked_add(Duration::from_millis(50)) + ); + let preserved_gate = task.send_gate; + let TargetState::Active { client } = + mem::replace(&mut task.targets[0].state, TargetState::Terminal) + else { + panic!("fast target was not active before removal"); + }; + assert!(task.begin_drain(0, client, ManagedTargetEndReason::Stopped, Instant::now())); + assert_eq!(task.send_gate, preserved_gate); + drive_task_until(&runtime, &mut task, |task| { + task.targets[1].counters.packets_sent == 2 + }); + drop(handle.stop()); let outcome = runtime.block_on(task); assert_eq!(outcome.failed_target_outcomes, 0); @@ -907,37 +939,125 @@ fn overdue_queued_reply_is_lost_then_late() { } #[test] -fn committed_probe_graceful_drain() { - let server = start_server(ServerBehavior::DelayedEcho(Duration::from_millis(25)), None); - let mut config = config(ManagedPacing::Staggered); - config.completion = ManagedCompletionPolicy::ExplicitStop; - config.client.duration = None; - config.client.interval = Duration::from_millis(100); - config.client.probe_timeout = Duration::from_millis(80); - config.final_drain = Duration::from_millis(5); - let (task, handle) = ManagedClient::task(config, vec![target("one", server.addr)]).unwrap(); - let stop_handle = handle.clone(); - let seen = Arc::clone(&server.probe_seen); - let stopper = thread::spawn(move || { - let (flag, ready) = &*seen; - let guard = flag.lock().unwrap(); - let (guard, timeout) = ready - .wait_timeout_while(guard, Duration::from_secs(2), |seen| !*seen) - .unwrap(); - let probe_seen = *guard; - drop(guard); - drop(stop_handle.stop()); - assert!(probe_seen && !timeout.timed_out()); +fn retained_probe_deadline_shortens_after_late_reply() { + let reply_gate = Arc::new(PacketGate::default()); + let server = start_server_with_gates( + ServerBehavior::Echo, + None, + None, + None, + Some(Arc::clone(&reply_gate)), + ); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + managed.client.interval = Duration::from_secs(1); + managed.client.probe_timeout = Duration::from_millis(100); + managed.final_drain = Duration::from_millis(200); + let final_drain = managed.final_drain; + let (task, handle) = ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap(); + let runtime = runtime(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| { + task.targets[0].counters.packets_sent == 1 }); - let outcome = runtime().block_on(task); - stopper.join().unwrap(); + drop(handle.stop()); + drive_task_until(&runtime, &mut task, |task| { + matches!(task.targets[0].state, TargetState::Draining { .. }) + }); + let (drain_started_at, initial_deadline, timeout_at) = match &task.targets[0].state { + TargetState::Draining { + client, + drain_started_at, + deadline, + .. + } => ( + *drain_started_at, + *deadline, + client.next_probe_timeout_deadline().unwrap(), + ), + _ => unreachable!(), + }; + assert_eq!( + initial_deadline, + timeout_at.checked_add(final_drain).unwrap() + ); + thread::sleep(timeout_at.saturating_duration_since(Instant::now())); + while Instant::now() <= timeout_at { + thread::yield_now(); + } + drive_task_until(&runtime, &mut task, |task| match &task.targets[0].state { + TargetState::Draining { client, .. } => { + client.next_probe_timeout_deadline().is_none() + && client.latest_probe_timeout_deadline() == Some(timeout_at) + } + _ => false, + }); + let retained_deadline = match &task.targets[0].state { + TargetState::Draining { deadline, .. } => *deadline, + _ => unreachable!(), + }; + assert_eq!(retained_deadline, initial_deadline); + reply_gate.release(); + wait_flag(&server.reply_sent); + let shortened_deadline = drain_started_at.checked_add(final_drain).unwrap(); + drive_task_until(&runtime, &mut task, |task| { + matches!( + task.targets[0].state, + TargetState::Draining { deadline, .. } if deadline == shortened_deadline + ) + }); + assert!(shortened_deadline < initial_deadline); + let outcome = runtime.block_on(task); let records = server.finish(); let target = &outcome.recent_target_outcomes[0]; assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); assert!(matches!(target.end_reason, ManagedTargetEndReason::Stopped)); assert_eq!(target.packets_sent, 1); + assert_eq!(target.replies_received, 0); + assert_eq!(target.late, 1); + assert_eq!(probes(&records).len(), 1); +} + +#[test] +fn pending_limit_failure_drains_reply_and_closes_session() { + let key = b"managed-send-failure-cleanup".to_vec(); + let server = start_server( + ServerBehavior::DelayedEcho(Duration::from_millis(80)), + Some(key.clone()), + ); + let mut managed = config(ManagedPacing::Staggered); + managed.client.duration = Some(Duration::from_millis(250)); + managed.client.interval = Duration::from_millis(20); + managed.client.probe_timeout = Duration::from_secs(2); + managed.client.max_pending_probes = 1; + managed.final_drain = Duration::from_millis(50); + let mut configured = target("one", server.addr); + configured.auth = Some(ClientAuthConfig { + hmac_key: Some(key), + }); + let (task, _) = ManagedClient::task(managed, vec![configured]).unwrap(); + let outcome = runtime().block_on(async { + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("send failure cleanup must shorten after draining the committed reply") + }); + let records = server.finish(); + let target = &outcome.recent_target_outcomes[0]; + assert_eq!(outcome.failed_target_outcomes, 1); + assert!(matches!( + target.end_reason, + ManagedTargetEndReason::Failed(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Sending, + kind: ManagedTargetFailureKind::ResourceExhausted, + .. + }) + )); + assert_eq!(target.cleanup_failure, None); + assert_eq!(target.packets_sent, 1); assert_eq!(target.replies_received, 1); assert_eq!(probes(&records).len(), 1); + assert!(has_close(&records)); } #[test]