From 1af8afcfae5d95fa2bc82fba746cdb45d0463887 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 22:33:00 +0700 Subject: [PATCH] fix(managed): correct pacing and drain cleanup --- crates/irtt-client/src/async_client.rs | 7 + crates/irtt-client/src/client/schedule.rs | 5 + crates/irtt-client/src/managed/task.rs | 195 +++++++++++++++-- crates/irtt-client/src/managed/tests.rs | 247 +++++++++++++++++++++- 4 files changed, 429 insertions(+), 25 deletions(-) diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs index 4c561c3..5051eca 100644 --- a/crates/irtt-client/src/async_client.rs +++ b/crates/irtt-client/src/async_client.rs @@ -663,6 +663,13 @@ impl AsyncClient { self.schedule.as_ref()?.next_send_deadline() } + pub(crate) fn probe_interval(&self) -> Option { + if !self.machine.is_open() { + return None; + } + self.schedule.as_ref().map(ProbeSchedule::interval) + } + /// Return the configured local probe timeout. pub fn probe_timeout(&self) -> Duration { self.machine.probe_timeout() diff --git a/crates/irtt-client/src/client/schedule.rs b/crates/irtt-client/src/client/schedule.rs index 19fb7b3..8779b98 100644 --- a/crates/irtt-client/src/client/schedule.rs +++ b/crates/irtt-client/src/client/schedule.rs @@ -50,6 +50,11 @@ impl ProbeSchedule { self.next_send_at } + #[cfg(feature = "tokio")] + pub(crate) fn interval(&self) -> Duration { + self.interval + } + pub(crate) fn permit_probe_at(&mut self, now: Instant) -> bool { if self.next_send_at.is_none() { return false; diff --git a/crates/irtt-client/src/managed/task.rs b/crates/irtt-client/src/managed/task.rs index 411ef70..30052cc 100644 --- a/crates/irtt-client/src/managed/task.rs +++ b/crates/irtt-client/src/managed/task.rs @@ -35,6 +35,7 @@ use super::{ }; const TARGET_WORK_BUDGET: usize = 128; +const POST_DEADLINE_RECEIVE_BUDGET: usize = TARGET_WORK_BUDGET; /// Entry point for constructing a unified Tokio managed task. #[derive(Debug, Default)] @@ -144,6 +145,8 @@ type ConnectFuture = type WakeFuture = Pin>> + Send + 'static>>; #[cfg(test)] type EventObservations = Arc)>>>; +#[cfg(test)] +type StaggerObservations = Arc>>; fn arm_wake(mut receiver: watch::Receiver<()>) -> WakeFuture { Box::pin(async move { @@ -170,11 +173,14 @@ enum TargetState { client: AsyncClient, deadline: Instant, primary_end: ManagedTargetEndReason, + cleanup_failure: Option, + post_deadline_receives_remaining: usize, }, Closing { client: AsyncClient, deadline: Instant, primary_end: ManagedTargetEndReason, + cleanup_failure: Option, }, Terminal, } @@ -317,6 +323,18 @@ pub struct ManagedClientTask { _next_generation: u64, #[cfg(test)] event_observations: Option, + #[cfg(test)] + stagger_observations: Option, + #[cfg(test)] + drain_test_hook: DrainTestHook, +} + +#[cfg(test)] +#[derive(Default)] +struct DrainTestHook { + defer_work_until_deadline: bool, + fail_receive: bool, + fail_close: bool, } impl ManagedClientTask { @@ -616,7 +634,28 @@ impl ManagedClientTask { mut client, deadline, primary_end, + mut cleanup_failure, + mut post_deadline_receives_remaining, } => { + #[cfg(test)] + let defer_work = self.drain_test_hook.defer_work_until_deadline && now < deadline; + #[cfg(not(test))] + let defer_work = false; + if defer_work { + self.targets[index].state = TargetState::Draining { + client, + deadline, + primary_end, + cleanup_failure, + post_deadline_receives_remaining, + }; + return false; + } + + #[cfg(test)] + let injected_receive_failure = self.take_drain_receive_failure(); + #[cfg(not(test))] + let injected_receive_failure = None; let received = match client.poll_recv(cx) { Poll::Pending => false, Poll::Ready(Ok(events)) => { @@ -625,13 +664,22 @@ impl ManagedClientTask { } Poll::Ready(Err(error)) => { self.targets[index].counters.packets_sent = client.packets_sent(); - self.fail_target(index, ManagedTargetFailurePhase::Receiving, error); - return false; + cleanup_failure.get_or_insert_with(|| { + classify_client_error(ManagedTargetFailurePhase::Receiving, &error) + }); + return self.begin_close(index, client, primary_end, cleanup_failure, now); } }; + if let Some(error) = injected_receive_failure { + self.targets[index].counters.packets_sent = client.packets_sent(); + cleanup_failure.get_or_insert_with(|| { + classify_client_error(ManagedTargetFailurePhase::Receiving, &error) + }); + return self.begin_close(index, client, primary_end, cleanup_failure, now); + } if client.is_peer_closed() { self.targets[index].counters.packets_sent = client.packets_sent(); - self.finish_target(index, ManagedTargetEndReason::PeerClosed, None); + self.finish_target(index, primary_end, cleanup_failure); return false; } if client @@ -642,27 +690,40 @@ impl ManagedClientTask { 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; + cleanup_failure.get_or_insert_with(|| { + classify_client_error(ManagedTargetFailurePhase::Timing, &error) + }); + return self.begin_close( + index, + client, + primary_end, + cleanup_failure, + now, + ); } } } if now >= deadline { - let close_deadline = now.checked_add(self.config.final_drain).unwrap_or(now); - self.install_target_state( - index, - TargetState::Closing { + if received && post_deadline_receives_remaining > 1 { + post_deadline_receives_remaining -= 1; + self.targets[index].state = TargetState::Draining { client, - deadline: close_deadline, + deadline, primary_end, - }, - ); - true + cleanup_failure, + post_deadline_receives_remaining, + }; + true + } else { + self.begin_close(index, client, primary_end, cleanup_failure, now) + } } else { self.targets[index].state = TargetState::Draining { client, deadline, primary_end, + cleanup_failure, + post_deadline_receives_remaining, }; received } @@ -671,17 +732,23 @@ impl ManagedClientTask { mut client, deadline, primary_end, + cleanup_failure, } => match client.poll_close(cx) { Poll::Pending => { if now >= deadline { self.targets[index].counters.packets_sent = client.packets_sent(); - self.finish_target(index, primary_end, Some(close_timeout_failure())); + self.finish_target( + index, + primary_end, + cleanup_failure.or_else(|| Some(close_timeout_failure())), + ); false } else { self.targets[index].state = TargetState::Closing { client, deadline, primary_end, + cleanup_failure, }; false } @@ -689,13 +756,23 @@ impl ManagedClientTask { Poll::Ready(Ok(events)) => { self.targets[index].counters.packets_sent = client.packets_sent(); self.publish_client_events(index, events); - self.finish_target(index, primary_end, None); + #[cfg(test)] + let injected_close_failure = self.take_drain_close_failure().map(|error| { + classify_client_error(ManagedTargetFailurePhase::Closing, &error) + }); + #[cfg(not(test))] + let injected_close_failure = None; + self.finish_target( + index, + primary_end, + cleanup_failure.or(injected_close_failure), + ); false } Poll::Ready(Err(error)) => { self.targets[index].counters.packets_sent = client.packets_sent(); let cleanup = classify_client_error(ManagedTargetFailurePhase::Closing, &error); - self.finish_target(index, primary_end, Some(cleanup)); + self.finish_target(index, primary_end, cleanup_failure.or(Some(cleanup))); false } }, @@ -734,11 +811,54 @@ impl ManagedClientTask { client, deadline, primary_end, + cleanup_failure: None, + post_deadline_receives_remaining: POST_DEADLINE_RECEIVE_BUDGET, }, ); true } + fn begin_close( + &mut self, + index: usize, + client: AsyncClient, + primary_end: ManagedTargetEndReason, + cleanup_failure: Option, + now: Instant, + ) -> bool { + let deadline = now.checked_add(self.config.final_drain).unwrap_or(now); + self.install_target_state( + index, + TargetState::Closing { + client, + deadline, + primary_end, + cleanup_failure, + }, + ); + true + } + + #[cfg(test)] + fn take_drain_receive_failure(&mut self) -> Option { + if !mem::take(&mut self.drain_test_hook.fail_receive) { + return None; + } + Some(ClientError::Socket(std::io::Error::other( + "injected drain receive failure", + ))) + } + + #[cfg(test)] + fn take_drain_close_failure(&mut self) -> Option { + if !mem::take(&mut self.drain_test_hook.fail_close) { + return None; + } + Some(ClientError::Socket(std::io::Error::other( + "injected drain close failure", + ))) + } + fn active_count(&self) -> usize { self.targets .iter() @@ -746,6 +866,22 @@ impl ManagedClientTask { .count() } + fn active_stagger_spacing(&self) -> Option<(usize, Duration)> { + let mut active = 0; + let mut minimum: Option = None; + for target in &self.targets { + let TargetState::Active { client } = &target.state else { + continue; + }; + active += 1; + let interval = client + .probe_interval() + .expect("active managed targets have a committed probe schedule"); + minimum = minimum.into_iter().chain(Some(interval)).min(); + } + minimum.map(|interval| (active, stagger_spacing(interval, active))) + } + fn poll_one_send(&mut self, index: usize, cx: &mut Context<'_>, now: Instant) -> SendResult { let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal); let TargetState::Active { mut client } = state else { @@ -802,8 +938,7 @@ impl ManagedClientTask { } fn poll_staggered_send(&mut self, cx: &mut Context<'_>, now: Instant) -> bool { - let active = self.active_count(); - if active == 0 || self.send_gate.is_some_and(|gate| gate > now) { + if self.active_count() == 0 || self.send_gate.is_some_and(|gate| gate > now) { return false; } if self.stagger_remaining == 0 { @@ -817,14 +952,20 @@ impl ManagedClientTask { if !matches!(self.targets[index].state, TargetState::Active { .. }) { continue; } + let stagger_spacing = self.active_stagger_spacing(); let result = self.poll_one_send(index, cx, now); if result == SendResult::NotAttempted { continue; } self.stagger_remaining = 0; if result.accepted() { - self.send_gate = Instant::now() - .checked_add(stagger_spacing(self.config.client.interval, active)); + if let Some((_active, spacing)) = stagger_spacing { + #[cfg(test)] + if let Some(observations) = &self.stagger_observations { + observations.lock().unwrap().push((_active, spacing)); + } + self.send_gate = Instant::now().checked_add(spacing); + } } return matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. }); } @@ -894,9 +1035,17 @@ impl ManagedClientTask { TargetState::Draining { client, deadline, .. } => { + #[cfg(test)] + let defer_work = self.drain_test_hook.defer_work_until_deadline; + #[cfg(not(test))] + let defer_work = false; non_send_deadline = non_send_deadline .into_iter() - .chain(client.next_probe_timeout_deadline()) + .chain( + (!defer_work) + .then(|| client.next_probe_timeout_deadline()) + .flatten(), + ) .chain(Some(*deadline)) .min(); } @@ -1193,6 +1342,10 @@ fn build_task( _next_generation: next_generation, #[cfg(test)] event_observations: None, + #[cfg(test)] + stagger_observations: None, + #[cfg(test)] + drain_test_hook: DrainTestHook::default(), }; let handle = ManagedClientHandle { stop, diff --git a/crates/irtt-client/src/managed/tests.rs b/crates/irtt-client/src/managed/tests.rs index 2ce1d95..e349972 100644 --- a/crates/irtt-client/src/managed/tests.rs +++ b/crates/irtt-client/src/managed/tests.rs @@ -15,7 +15,7 @@ use irtt_proto::{ use tokio::runtime::{Builder, Runtime}; use super::*; -use crate::{socket::resolution_call_counts, ClientAuthConfig, RunMode}; +use crate::{socket::resolution_call_counts, ClientAuthConfig, NegotiationPolicy, RunMode}; const TOKEN: u64 = 0x1234_5678_90ab_cdef; @@ -25,6 +25,7 @@ enum ServerBehavior { NoTest, PeerClose, DelayedEcho(Duration), + DeferredBurst(usize), } #[derive(Clone, Debug)] @@ -55,6 +56,14 @@ impl TestServer { } fn start_server(behavior: ServerBehavior, key: Option>) -> TestServer { + start_server_negotiating(behavior, key, None) +} + +fn start_server_negotiating( + behavior: ServerBehavior, + key: Option>, + interval: Option, +) -> TestServer { let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); socket .set_read_timeout(Some(Duration::from_secs(3))) @@ -72,6 +81,10 @@ fn start_server(behavior: ServerBehavior, key: Option>) -> TestServer { at: Instant::now(), }); let request = decode_open_request(&buffer[..open_len], key.as_deref()).unwrap(); + let mut negotiated = request.params.clone(); + if let Some(interval) = interval { + negotiated.interval_ns = i64::try_from(interval.as_nanos()).unwrap(); + } let no_test = matches!(behavior, ServerBehavior::NoTest); let open_reply = encode_open_reply( &OpenReply { @@ -79,7 +92,7 @@ fn start_server(behavior: ServerBehavior, key: Option>) -> TestServer { | flags::FLAG_REPLY | if no_test { flags::FLAG_CLOSE } else { 0 }, token: if no_test { 0 } else { TOKEN }, - params: request.params.clone(), + params: negotiated.clone(), }, key.as_deref(), ) @@ -89,12 +102,13 @@ fn start_server(behavior: ServerBehavior, key: Option>) -> TestServer { return; } + let mut deferred_burst_sent = false; loop { let Ok((len, packet_peer)) = socket.recv_from(&mut buffer) else { return; }; let packet = &buffer[..len]; - if let Ok(probe) = decode_echo_request(packet, &request.params, key.as_deref()) { + if let Ok(probe) = decode_echo_request(packet, &negotiated, key.as_deref()) { thread_records.lock().unwrap().push(PacketRecord { kind: PacketKind::Probe, at: Instant::now(), @@ -105,6 +119,35 @@ fn start_server(behavior: ServerBehavior, key: Option>) -> TestServer { if let ServerBehavior::DelayedEcho(delay) = behavior { thread::sleep(delay); } + if let ServerBehavior::DeferredBurst(count) = behavior { + if deferred_burst_sent { + continue; + } + deferred_burst_sent = true; + thread::sleep(Duration::from_millis(30)); + for reply_index in 0..count { + let reply = encode_echo_reply( + &EchoReply { + flags: flags::FLAG_REPLY, + token: TOKEN, + sequence: if reply_index < 2 { + probe.sequence + } else { + u32::MAX + }, + recv_count: None, + recv_window: None, + timestamps: TimestampFields::default(), + payload: Vec::new(), + }, + &negotiated, + key.as_deref(), + ) + .unwrap(); + socket.send_to(&reply, packet_peer).unwrap(); + } + continue; + } let peer_close = matches!(behavior, ServerBehavior::PeerClose); let reply = encode_echo_reply( &EchoReply { @@ -116,7 +159,7 @@ fn start_server(behavior: ServerBehavior, key: Option>) -> TestServer { timestamps: TimestampFields::default(), payload: Vec::new(), }, - &request.params, + &negotiated, key.as_deref(), ) .unwrap(); @@ -223,6 +266,75 @@ fn assert_rotated_without_catchup(sends: &[(TargetInstance, u32, Instant, Instan } } +fn run_negotiated_stagger_case( + requested: Duration, + negotiated: &[Duration], +) -> (Vec<(usize, Duration)>, Vec) { + let servers = negotiated + .iter() + .map(|interval| start_server_negotiating(ServerBehavior::Echo, None, Some(*interval))) + .collect::>(); + let mut managed = config(ManagedPacing::Staggered); + managed.client.duration = Some(Duration::from_millis(190)); + managed.client.interval = requested; + managed.client.negotiation_policy = NegotiationPolicy::Loose; + let targets = servers + .iter() + .enumerate() + .map(|(index, server)| target(&format!("target-{index}"), server.addr)) + .collect(); + let (mut task, _) = ManagedClient::task(managed, targets).unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + task.stagger_observations = Some(Arc::clone(&observations)); + let outcome = runtime().block_on(task); + let packets_sent = (0..servers.len()) + .map(|index| { + outcome + .recent_target_outcomes + .iter() + .find(|target| target.target.id.as_ref() == format!("target-{index}")) + .unwrap() + .packets_sent + }) + .collect(); + for server in servers { + server.finish(); + } + let spacings = observations.lock().unwrap().clone(); + (spacings, packets_sent) +} + +fn run_deferred_drain_burst(packet_count: usize) -> ManagedOutcome { + let server = start_server(ServerBehavior::DeferredBurst(packet_count), None); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + managed.client.interval = Duration::from_millis(100); + managed.client.probe_timeout = Duration::from_millis(60); + managed.final_drain = Duration::from_millis(20); + let (mut task, handle) = + ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap(); + task.drain_test_hook.defer_work_until_deadline = true; + 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(); + assert!(*guard && !timeout.timed_out()); + drop(handle.stop()); + }); + let outcome = runtime().block_on(async { + tokio::time::timeout(Duration::from_secs(2), task) + .await + .unwrap() + }); + stopper.join().unwrap(); + server.finish(); + outcome +} + #[test] fn construction_is_runtime_and_io_free() { let before = resolution_call_counts(); @@ -446,6 +558,45 @@ fn staggered_pacing_fairness() { second.finish(); } +#[test] +fn staggered_pacing_uses_all_active_negotiated_intervals() { + let requested = Duration::from_millis(120); + + let (single, single_packets) = + run_negotiated_stagger_case(requested, &[Duration::from_millis(40)]); + assert!(single + .iter() + .all(|entry| *entry == (1, Duration::from_millis(40)))); + assert!(single_packets[0] >= 4); + + let (different, different_packets) = run_negotiated_stagger_case( + requested, + &[Duration::from_millis(30), Duration::from_millis(90)], + ); + let while_both_active = different + .iter() + .filter(|(active, _)| *active == 2) + .collect::>(); + assert!(!while_both_active.is_empty()); + assert!(while_both_active + .iter() + .all(|(_, spacing)| *spacing == Duration::from_millis(15))); + assert!(different_packets[0] > different_packets[1]); + + let (equal, _) = run_negotiated_stagger_case( + requested, + &[Duration::from_millis(60), Duration::from_millis(60)], + ); + let while_both_active = equal + .iter() + .filter(|(active, _)| *active == 2) + .collect::>(); + assert!(!while_both_active.is_empty()); + assert!(while_both_active + .iter() + .all(|(_, spacing)| *spacing == Duration::from_millis(30))); +} + #[test] fn committed_probe_graceful_drain() { let server = start_server(ServerBehavior::DelayedEcho(Duration::from_millis(25)), None); @@ -480,6 +631,94 @@ fn committed_probe_graceful_drain() { assert_eq!(probes(&records).len(), 1); } +#[test] +fn drain_failures_preserve_primary_outcome_and_first_cleanup_failure() { + let completed_server = start_server(ServerBehavior::Echo, None); + let (mut completed_task, _) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("completed", completed_server.addr)], + ) + .unwrap(); + completed_task.drain_test_hook.fail_receive = true; + completed_task.drain_test_hook.fail_close = true; + let completed = runtime().block_on(completed_task); + let completed_target = &completed.recent_target_outcomes[0]; + assert!(matches!( + completed_target.end_reason, + ManagedTargetEndReason::TestComplete + )); + assert_eq!(completed.failed_target_outcomes, 0); + assert!(matches!( + completed_target.cleanup_failure, + Some(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Receiving, + kind: ManagedTargetFailureKind::Socket, + .. + }) + )); + completed_server.finish(); + + let stopped_server = start_server(ServerBehavior::Echo, None); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + let (mut stopped_task, handle) = + ManagedClient::task(managed, vec![target("stopped", stopped_server.addr)]).unwrap(); + stopped_task.drain_test_hook.fail_receive = true; + let seen = Arc::clone(&stopped_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(); + assert!(*guard && !timeout.timed_out()); + drop(handle.stop()); + }); + let stopped = runtime().block_on(stopped_task); + stopper.join().unwrap(); + let stopped_target = &stopped.recent_target_outcomes[0]; + assert!(matches!( + stopped_target.end_reason, + ManagedTargetEndReason::Stopped + )); + assert_eq!(stopped.failed_target_outcomes, 0); + assert!(matches!( + stopped_target.cleanup_failure, + Some(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Receiving, + kind: ManagedTargetFailureKind::Socket, + .. + }) + )); + stopped_server.finish(); +} + +#[test] +fn post_deadline_receive_sweep_drains_queued_packets_and_is_bounded() { + let drained = run_deferred_drain_burst(3); + let drained_target = &drained.recent_target_outcomes[0]; + assert!(matches!( + drained_target.end_reason, + ManagedTargetEndReason::Stopped + )); + assert_eq!(drained_target.replies_received, 1); + assert_eq!(drained_target.duplicates, 1); + assert_eq!(drained_target.late, 1); + + let bounded = run_deferred_drain_burst(POST_DEADLINE_RECEIVE_BUDGET + 64); + let bounded_target = &bounded.recent_target_outcomes[0]; + assert!(matches!( + bounded_target.end_reason, + ManagedTargetEndReason::Stopped + )); + assert_eq!( + bounded_target.replies_received + bounded_target.duplicates + bounded_target.late, + u64::try_from(POST_DEADLINE_RECEIVE_BUDGET).unwrap() + ); + assert_eq!(bounded.failed_target_outcomes, 0); +} + #[test] fn event_loss_independence() { let server = start_server(ServerBehavior::Echo, None);