diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs index 5051eca..8523746 100644 --- a/crates/irtt-client/src/async_client.rs +++ b/crates/irtt-client/src/async_client.rs @@ -86,6 +86,8 @@ pub(crate) struct AsyncOpenState { deadline: Option, deadline_timer: Option>>, request_submitted: bool, + stop_after_current_attempt: bool, + stopped_after_current_attempt: bool, buffer: [u8; MAX_OPEN_PACKET_SIZE], cleanup: Option, } @@ -97,6 +99,8 @@ impl AsyncOpenState { deadline: None, deadline_timer: None, request_submitted: false, + stop_after_current_attempt: false, + stopped_after_current_attempt: false, buffer: [0_u8; MAX_OPEN_PACKET_SIZE], cleanup: None, } @@ -117,6 +121,24 @@ impl AsyncOpenState { self.deadline = None; self.deadline_timer = None; self.request_submitted = false; + self.stopped_after_current_attempt = self.stop_after_current_attempt; + } + + pub(crate) fn has_in_flight_work(&self) -> bool { + self.request_submitted || self.cleanup.is_some() + } + + pub(crate) fn request_stop_after_current_attempt(&mut self) { + self.stop_after_current_attempt = true; + } + + pub(crate) fn stopped_after_current_attempt(&self) -> bool { + self.stopped_after_current_attempt + } + + #[cfg(test)] + pub(crate) fn has_retained_cleanup(&self) -> bool { + self.cleanup.is_some() } fn deadline(&self) -> Instant { @@ -149,6 +171,7 @@ struct AsyncClientTestHooks { fail_cleanup_send: Cell, fail_peer_close_dscp: Cell, fail_dscp_restore: Cell, + pause_open_cleanup_before_send: Cell, pause_open_before_writable: Cell, pause_open_before_readable: Cell, pause_probe_before_writable: Cell, @@ -301,6 +324,10 @@ impl AsyncClient { if state.attempt >= attempt_count { return Poll::Ready(Err(ClientError::OpenTimeout)); } + if state.stop_after_current_attempt && state.deadline.is_none() { + state.stopped_after_current_attempt = true; + return Poll::Ready(Err(ClientError::OpenTimeout)); + } if state.deadline.is_none() { let timeout = self.machine.config().open_timeouts[state.attempt]; if let Err(error) = state.start_attempt(timeout) { @@ -333,10 +360,10 @@ impl AsyncClient { let send_result = self.socket.try_send(&request.bytes); match send_result { Ok(bytes) => { + state.request_submitted = true; if let Err(error) = validate_datagram_length(request.bytes.len(), bytes) { return Poll::Ready(Err(error)); } - state.request_submitted = true; } Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, Err(error) => return Poll::Ready(Err(ClientError::Socket(error))), @@ -352,6 +379,9 @@ impl AsyncClient { Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), Poll::Ready(Ok(false)) => { state.finish_attempt(); + if state.stopped_after_current_attempt() { + return Poll::Ready(Err(ClientError::OpenTimeout)); + } continue; } Poll::Ready(Ok(true)) => {} @@ -364,6 +394,9 @@ impl AsyncClient { }; if datagram.received_at.mono > state.deadline() { state.finish_attempt(); + if state.stopped_after_current_attempt() { + return Poll::Ready(Err(ClientError::OpenTimeout)); + } continue; } @@ -709,6 +742,17 @@ impl AsyncClient { self.prepared_probe = None; } + #[cfg(test)] + pub(crate) fn retain_open_cleanup_once(&self) { + self.test_hooks.fail_open_dscp.set(true); + self.test_hooks.pause_open_cleanup_before_send.set(true); + } + + #[cfg(test)] + pub(crate) fn poll_recv_ready_for_test(&self, cx: &mut Context<'_>) -> Poll> { + self.socket.poll_recv_ready(cx) + } + pub(crate) fn skip_missed_probe_slots_at(&mut self, now: Instant) -> Result<(), ClientError> { let Some(schedule) = self.schedule.as_mut() else { return Ok(()); @@ -821,6 +865,15 @@ impl AsyncClient { return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); }; + #[cfg(test)] + if self + .test_hooks + .pause_open_cleanup_before_send + .replace(false) + { + return Poll::Pending; + } + #[cfg(test)] if self.test_hooks.fail_cleanup_send.replace(false) { return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); diff --git a/crates/irtt-client/src/managed/model.rs b/crates/irtt-client/src/managed/model.rs index 80bf437..230559e 100644 --- a/crates/irtt-client/src/managed/model.rs +++ b/crates/irtt-client/src/managed/model.rs @@ -260,6 +260,8 @@ pub enum ManagedConfigError { ZeroEventCapacity, #[error("managed live-generation limit must be greater than zero")] ZeroLiveGenerationLimit, + #[error("managed final drain {duration:?} cannot be scheduled from the current instant")] + UnschedulableFinalDrain { duration: Duration }, #[error("invalid managed target {id}: {source}")] InvalidTarget { id: TargetId, diff --git a/crates/irtt-client/src/managed/task.rs b/crates/irtt-client/src/managed/task.rs index 30052cc..ec3e8be 100644 --- a/crates/irtt-client/src/managed/task.rs +++ b/crates/irtt-client/src/managed/task.rs @@ -215,6 +215,7 @@ struct TargetRuntime { counters: TargetCounters, latest_committed_timeout: Option, send_waiting: bool, + active: bool, state: TargetState, } @@ -318,6 +319,7 @@ pub struct ManagedClientTask { burst_remaining: usize, stagger_remaining: usize, send_gate: Option, + last_stagger_send: Option, stop_observed: bool, final_outcome: Option>, _next_generation: u64, @@ -413,7 +415,13 @@ 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; + self.targets[index].active = active; self.targets[index].state = state; + if membership_changed { + self.recompute_stagger_gate(Instant::now()); + } self.replace_status(); self.publish_event(ManagedEvent::TargetStateChanged { target: self.targets[index].instance.clone(), @@ -557,9 +565,12 @@ impl ManagedClientTask { mut open, } => { if self.state == DriverState::Stopping { - self.targets[index].counters.packets_sent = client.packets_sent(); - self.finish_target(index, ManagedTargetEndReason::Stopped, None); - return false; + if !open.has_in_flight_work() { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.finish_target(index, ManagedTargetEndReason::Stopped, None); + return false; + } + open.request_stop_after_current_attempt(); } match client.poll_open(&mut open, cx) { Poll::Pending => { @@ -567,17 +578,42 @@ impl ManagedClientTask { false } Poll::Ready(Ok(OpenOutcome::Started { event, .. })) => { - self.install_target_state(index, TargetState::Active { client }); self.publish_client_events(index, vec![event]); - true + if self.state == DriverState::Stopping { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.begin_drain(index, client, ManagedTargetEndReason::Stopped, now) + } else { + self.install_target_state(index, TargetState::Active { client }); + true + } } Poll::Ready(Ok(OpenOutcome::NoTestCompleted { event, .. })) => { self.publish_client_events(index, vec![event]); - self.finish_target(index, ManagedTargetEndReason::NoTestComplete, None); + let end_reason = if self.state == DriverState::Stopping { + ManagedTargetEndReason::Stopped + } else { + ManagedTargetEndReason::NoTestComplete + }; + self.finish_target(index, end_reason, None); false } Poll::Ready(Err(error)) => { - self.fail_target(index, ManagedTargetFailurePhase::Opening, error); + if self.state == DriverState::Stopping { + let cleanup_failure = + (!open.stopped_after_current_attempt()).then(|| { + classify_client_error( + ManagedTargetFailurePhase::Opening, + &error, + ) + }); + self.finish_target( + index, + ManagedTargetEndReason::Stopped, + cleanup_failure, + ); + } else { + self.fail_target(index, ManagedTargetFailurePhase::Opening, error); + } false } } @@ -588,6 +624,19 @@ impl ManagedClientTask { self.targets[index].counters.packets_sent = client.packets_sent(); return self.begin_drain(index, client, ManagedTargetEndReason::Stopped, now); } + if client + .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) => { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.fail_target(index, ManagedTargetFailurePhase::Timing, error); + return false; + } + } + } let received = match client.poll_recv(cx) { Poll::Pending => false, Poll::Ready(Ok(events)) => { @@ -605,19 +654,6 @@ impl ManagedClientTask { self.finish_target(index, ManagedTargetEndReason::PeerClosed, None); return false; } - if client - .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) => { - self.targets[index].counters.packets_sent = client.packets_sent(); - self.fail_target(index, ManagedTargetFailurePhase::Timing, error); - return false; - } - } - } if client.is_run_complete() { self.targets[index].counters.packets_sent = client.packets_sent(); return self.begin_drain( @@ -656,6 +692,27 @@ impl ManagedClientTask { let injected_receive_failure = self.take_drain_receive_failure(); #[cfg(not(test))] let injected_receive_failure = None; + 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), + 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)) => { @@ -679,30 +736,9 @@ impl ManagedClientTask { } if client.is_peer_closed() { self.targets[index].counters.packets_sent = client.packets_sent(); - self.finish_target(index, primary_end, cleanup_failure); + self.finish_target(index, ManagedTargetEndReason::PeerClosed, cleanup_failure); return 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), - 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, - ); - } - } - } if now >= deadline { if received && post_deadline_receives_remaining > 1 { post_deadline_receives_remaining -= 1; @@ -798,12 +834,8 @@ impl ManagedClientTask { .unwrap_or(now) .max(now); let Some(deadline) = latest.checked_add(self.config.final_drain) else { - self.fail_target( - index, - ManagedTargetFailurePhase::Timing, - ClientError::DurationOverflow, - ); - return false; + let cleanup_failure = Some(duration_overflow_failure()); + return self.begin_close(index, client, primary_end, cleanup_failure, now); }; self.install_target_state( index, @@ -823,10 +855,13 @@ impl ManagedClientTask { index: usize, client: AsyncClient, primary_end: ManagedTargetEndReason, - cleanup_failure: Option, + mut cleanup_failure: Option, now: Instant, ) -> bool { - let deadline = now.checked_add(self.config.final_drain).unwrap_or(now); + let deadline = now.checked_add(self.config.final_drain).unwrap_or_else(|| { + cleanup_failure.get_or_insert_with(duration_overflow_failure); + now + }); self.install_target_state( index, TargetState::Closing { @@ -860,10 +895,7 @@ impl ManagedClientTask { } fn active_count(&self) -> usize { - self.targets - .iter() - .filter(|target| matches!(target.state, TargetState::Active { .. })) - .count() + self.targets.iter().filter(|target| target.active).count() } fn active_stagger_spacing(&self) -> Option<(usize, Duration)> { @@ -882,6 +914,14 @@ impl ManagedClientTask { minimum.map(|interval| (active, stagger_spacing(interval, active))) } + fn recompute_stagger_gate(&mut self, now: Instant) { + self.send_gate = self + .last_stagger_send + .zip(self.active_stagger_spacing()) + .and_then(|(last, (_, spacing))| last.checked_add(spacing)) + .filter(|gate| *gate > now); + } + 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 { @@ -952,6 +992,7 @@ 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); if result == SendResult::NotAttempted { @@ -959,13 +1000,15 @@ impl ManagedClientTask { } self.stagger_remaining = 0; if result.accepted() { + #[cfg(test)] 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); } + 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 { .. }); } @@ -1231,6 +1274,11 @@ fn build_task( if config.max_live_target_generations == 0 { return Err(ManagedConfigError::ZeroLiveGenerationLimit); } + if Instant::now().checked_add(config.final_drain).is_none() { + return Err(ManagedConfigError::UnschedulableFinalDrain { + duration: config.final_drain, + }); + } if targets.len() > config.max_live_target_generations { return Err(ManagedConfigError::TooManyTargets { configured: targets.len(), @@ -1273,6 +1321,7 @@ fn build_task( counters: TargetCounters::default(), latest_committed_timeout: None, send_waiting: false, + active: false, state: TargetState::Pending { client_config }, }); } @@ -1337,6 +1386,7 @@ fn build_task( burst_remaining: 0, stagger_remaining: 0, send_gate: None, + last_stagger_send: None, stop_observed: false, final_outcome: None, _next_generation: next_generation, @@ -1377,6 +1427,13 @@ fn stagger_spacing(interval: Duration, active_targets: usize) -> Duration { Duration::new(seconds, subsec) } +fn duration_overflow_failure() -> ManagedTargetFailure { + classify_client_error( + ManagedTargetFailurePhase::Timing, + &ClientError::DurationOverflow, + ) +} + fn close_timeout_failure() -> ManagedTargetFailure { ManagedTargetFailure { phase: ManagedTargetFailurePhase::Closing, diff --git a/crates/irtt-client/src/managed/tests.rs b/crates/irtt-client/src/managed/tests.rs index e349972..23e11f1 100644 --- a/crates/irtt-client/src/managed/tests.rs +++ b/crates/irtt-client/src/managed/tests.rs @@ -1,5 +1,5 @@ use std::{ - future::Future, + future::{poll_fn, Future}, net::{SocketAddr, UdpSocket}, pin::Pin, sync::{Arc, Condvar, Mutex}, @@ -45,6 +45,7 @@ struct TestServer { addr: SocketAddr, records: Arc>>, probe_seen: Arc<(Mutex, Condvar)>, + reply_sent: Arc<(Mutex, Condvar)>, thread: JoinHandle<()>, } @@ -55,14 +56,47 @@ impl TestServer { } } +#[derive(Default)] +struct PacketGate { + seen: Mutex, + seen_ready: Condvar, + released: Mutex, + release_ready: Condvar, +} + +impl PacketGate { + fn arrive_and_wait(&self) { + *self.seen.lock().unwrap() = true; + self.seen_ready.notify_all(); + let released = self.released.lock().unwrap(); + let (released, timeout) = self + .release_ready + .wait_timeout_while(released, Duration::from_secs(2), |released| !*released) + .unwrap(); + assert!(*released && !timeout.timed_out()); + } + fn release(&self) { + *self.released.lock().unwrap() = true; + self.release_ready.notify_all(); + } +} 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 { + start_server_with_gates(behavior, key, interval, None, None) +} + +fn start_server_with_gates( + behavior: ServerBehavior, + key: Option>, + interval: Option, + open_gate: Option>, + reply_gate: Option>, ) -> TestServer { let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); socket @@ -73,6 +107,8 @@ fn start_server_negotiating( let thread_records = Arc::clone(&records); let probe_seen = Arc::new((Mutex::new(false), Condvar::new())); let thread_probe_seen = Arc::clone(&probe_seen); + let reply_sent = Arc::new((Mutex::new(false), Condvar::new())); + let thread_reply_sent = Arc::clone(&reply_sent); let thread = thread::spawn(move || { let mut buffer = [0_u8; 8192]; let (open_len, peer) = socket.recv_from(&mut buffer).unwrap(); @@ -81,6 +117,9 @@ fn start_server_negotiating( at: Instant::now(), }); let request = decode_open_request(&buffer[..open_len], key.as_deref()).unwrap(); + if let Some(gate) = &open_gate { + gate.arrive_and_wait(); + } let mut negotiated = request.params.clone(); if let Some(interval) = interval { negotiated.interval_ns = i64::try_from(interval.as_nanos()).unwrap(); @@ -148,6 +187,9 @@ fn start_server_negotiating( } continue; } + if let Some(gate) = &reply_gate { + gate.arrive_and_wait(); + } let peer_close = matches!(behavior, ServerBehavior::PeerClose); let reply = encode_echo_reply( &EchoReply { @@ -164,6 +206,9 @@ fn start_server_negotiating( ) .unwrap(); socket.send_to(&reply, packet_peer).unwrap(); + let (sent, ready) = &*thread_reply_sent; + *sent.lock().unwrap() = true; + ready.notify_all(); if peer_close { return; } @@ -180,6 +225,7 @@ fn start_server_negotiating( addr, records, probe_seen, + reply_sent, thread, } } @@ -216,6 +262,46 @@ fn target(id: &str, addr: SocketAddr) -> ManagedTargetConfig { fn poll_once(future: Pin<&mut F>) -> Poll { Future::poll(future, &mut Context::from_waker(Waker::noop())) } +fn drive_task_until( + runtime: &Runtime, + task: &mut Pin>, + condition: impl Fn(&ManagedClientTask) -> bool, +) { + runtime.block_on(poll_fn(|cx| { + assert!( + task.as_mut().poll(cx).is_pending(), + "task completed before condition" + ); + if condition(task.as_ref().get_ref()) { + Poll::Ready(()) + } else { + Poll::Pending + } + })); +} + +fn wait_flag(flag: &Arc<(Mutex, Condvar)>) { + let (flag, ready) = &**flag; + let value = flag.lock().unwrap(); + let (value, timeout) = ready + .wait_timeout_while(value, Duration::from_secs(2), |value| !*value) + .unwrap(); + assert!(*value && !timeout.timed_out()); +} + +fn client_events( + observations: &Mutex)>>, +) -> Vec { + observations + .lock() + .unwrap() + .iter() + .filter_map(|(event, _)| match event { + ManagedEvent::Client { event, .. } => Some(event.clone()), + _ => None, + }) + .collect() +} fn probes(records: &[PacketRecord]) -> Vec { records @@ -490,6 +576,114 @@ fn authenticated_peer_close_outcome() { assert!(!has_close(&records)); } +#[test] +fn peer_close_during_stop_drain_remains_authoritative() { + let key = b"managed-drain-peer-close".to_vec(); + let server = start_server(ServerBehavior::PeerClose, Some(key.clone())); + let mut configured = target("peer", server.addr); + configured.auth = Some(ClientAuthConfig { + hmac_key: Some(key), + }); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + managed.client.interval = Duration::from_secs(1); + let (task, handle) = ManagedClient::task(managed, vec![configured]).unwrap(); + let runtime = runtime(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| { + task.targets[0].counters.packets_sent == 1 + }); + drop(handle.stop()); + wait_flag(&server.reply_sent); + let outcome = runtime.block_on(task); + let records = server.finish(); + assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); + assert_eq!(outcome.peer_closed_target_outcomes, 1); + assert!(matches!( + outcome.recent_target_outcomes[0].end_reason, + ManagedTargetEndReason::PeerClosed + )); + assert!(!has_close(&records)); +} + +#[test] +fn stop_during_opening_finishes_in_flight_work() { + let open_gate = Arc::new(PacketGate::default()); + let server = start_server_with_gates( + ServerBehavior::Echo, + None, + None, + Some(Arc::clone(&open_gate)), + None, + ); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + managed.client.open_timeouts = vec![Duration::from_secs(1), Duration::from_secs(1)]; + let (task, handle) = + ManagedClient::task(managed, vec![target("success", server.addr)]).unwrap(); + let runtime = runtime(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| match &task.targets[0].state { + TargetState::Opening { open, .. } => open.has_in_flight_work(), + _ => false, + }); + drop(handle.stop()); + open_gate.release(); + let outcome = runtime.block_on(task); + let records = server.finish(); + assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); + assert!(matches!( + outcome.recent_target_outcomes[0].end_reason, + ManagedTargetEndReason::Stopped + )); + assert_eq!( + records + .iter() + .filter(|record| matches!(record.kind, PacketKind::Open)) + .count(), + 1 + ); + assert!(probes(&records).is_empty() && has_close(&records)); + let key = b"managed-retained-open-cleanup".to_vec(); + let cleanup_server = start_server(ServerBehavior::Echo, Some(key.clone())); + let mut configured = target("cleanup", cleanup_server.addr); + configured.auth = Some(ClientAuthConfig { + hmac_key: Some(key), + }); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + let (task, handle) = ManagedClient::task(managed, vec![configured]).unwrap(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| { + matches!(task.targets[0].state, TargetState::Opening { .. }) + }); + let TargetState::Opening { client, .. } = &task.as_ref().get_ref().targets[0].state else { + unreachable!() + }; + client.retain_open_cleanup_once(); + drive_task_until(&runtime, &mut task, |task| match &task.targets[0].state { + TargetState::Opening { open, .. } => open.has_retained_cleanup(), + _ => false, + }); + drop(handle.stop()); + let outcome = runtime.block_on(task); + let records = cleanup_server.finish(); + let target = &outcome.recent_target_outcomes[0]; + assert!(matches!(target.end_reason, ManagedTargetEndReason::Stopped)); + assert_eq!(outcome.failed_target_outcomes, 0); + assert!(matches!( + target.cleanup_failure, + Some(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Opening, + .. + }) + )); + assert!(probes(&records).is_empty() && has_close(&records)); +} + #[test] fn static_multi_target_completion() { let first = start_server(ServerBehavior::Echo, None); @@ -597,6 +791,121 @@ fn staggered_pacing_uses_all_active_negotiated_intervals() { .all(|(_, spacing)| *spacing == Duration::from_millis(30))); } +#[test] +fn stagger_gate_tracks_active_membership() { + let first = start_server_negotiating(ServerBehavior::Echo, None, Some(Duration::from_secs(1))); + let second_gate = Arc::new(PacketGate::default()); + let second = start_server_with_gates( + ServerBehavior::Echo, + None, + Some(Duration::from_secs(1)), + Some(Arc::clone(&second_gate)), + None, + ); + let mut managed = config(ManagedPacing::Staggered); + managed.client.duration = Some(Duration::from_secs(2)); + managed.client.interval = Duration::from_secs(1); + managed.client.negotiation_policy = NegotiationPolicy::Loose; + managed.client.open_timeouts = vec![Duration::from_secs(2)]; + let (task, handle) = ManagedClient::task( + managed, + vec![target("first", first.addr), target("second", second.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 one_target_last = task.last_stagger_send.unwrap(); + assert_eq!( + task.send_gate, + one_target_last.checked_add(Duration::from_secs(1)) + ); + 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()); + drive_task_until(&runtime, &mut task, |task| { + task.targets[1].counters.packets_sent == 1 + }); + drop(handle.stop()); + let outcome = runtime.block_on(task); + assert_eq!(outcome.failed_target_outcomes, 0); + first.finish(); + second.finish(); +} + +#[test] +fn overdue_queued_reply_is_lost_then_late() { + 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(30); + let (mut task, handle) = + ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + task.event_observations = Some(Arc::clone(&observations)); + let runtime = runtime(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| { + task.targets[0].counters.packets_sent == 1 + }); + let timeout = match &task.targets[0].state { + TargetState::Active { client } => client.next_probe_timeout_deadline().unwrap(), + _ => panic!("probe sender did not remain active"), + }; + thread::sleep(timeout.saturating_duration_since(Instant::now())); + while Instant::now() <= timeout { + thread::yield_now(); + } + reply_gate.release(); + wait_flag(&server.reply_sent); + runtime.block_on(poll_fn(|cx| { + let ready = match &task.as_ref().get_ref().targets[0].state { + TargetState::Active { client } => client.poll_recv_ready_for_test(cx), + _ => panic!("probe sender did not remain active"), + }; + ready.map(|result| result.unwrap()) + })); + runtime.block_on(poll_fn(|cx| { + assert!(task.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + })); + let classified = client_events(&observations) + .into_iter() + .filter(|event| { + matches!( + event, + ClientEvent::EchoLoss { .. } + | ClientEvent::LateReply { .. } + | ClientEvent::EchoReply { .. } + ) + }) + .collect::>(); + assert!(matches!( + classified.as_slice(), + [ClientEvent::EchoLoss { .. }, ClientEvent::LateReply { .. }] + )); + drop(handle.stop()); + let outcome = runtime.block_on(task); + let target = &outcome.recent_target_outcomes[0]; + assert_eq!(target.replies_received, 0); + assert_eq!(target.late, 1); + server.finish(); +} + #[test] fn committed_probe_graceful_drain() { let server = start_server(ServerBehavior::DelayedEcho(Duration::from_millis(25)), None); @@ -694,6 +1003,46 @@ fn drain_failures_preserve_primary_outcome_and_first_cleanup_failure() { stopped_server.finish(); } +#[test] +fn final_drain_overflow_is_rejected_and_preserves_success() { + let mut invalid = config(ManagedPacing::Staggered); + invalid.final_drain = Duration::MAX; + assert!(matches!( + ManagedClient::task( + invalid, + vec![target("invalid", "127.0.0.1:9".parse().unwrap())] + ), + Err(ManagedConfigError::UnschedulableFinalDrain { + duration: Duration::MAX + }) + )); + + let server = start_server(ServerBehavior::Echo, None); + let (mut task, _) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("completed", server.addr)], + ) + .unwrap(); + task.config.final_drain = Duration::MAX; + let outcome = runtime().block_on(task); + let target = &outcome.recent_target_outcomes[0]; + assert!(matches!( + target.end_reason, + ManagedTargetEndReason::TestComplete + )); + assert_eq!(outcome.successful_target_outcomes, 1); + assert_eq!(outcome.failed_target_outcomes, 0); + assert!(matches!( + target.cleanup_failure, + Some(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Timing, + kind: ManagedTargetFailureKind::ResourceExhausted, + .. + }) + )); + assert!(has_close(&server.finish())); +} + #[test] fn post_deadline_receive_sweep_drains_queued_packets_and_is_bounded() { let drained = run_deferred_drain_burst(3); @@ -702,9 +1051,9 @@ fn post_deadline_receive_sweep_drains_queued_packets_and_is_bounded() { drained_target.end_reason, ManagedTargetEndReason::Stopped )); - assert_eq!(drained_target.replies_received, 1); + assert_eq!(drained_target.replies_received, 0); assert_eq!(drained_target.duplicates, 1); - assert_eq!(drained_target.late, 1); + assert_eq!(drained_target.late, 2); let bounded = run_deferred_drain_burst(POST_DEADLINE_RECEIVE_BUDGET + 64); let bounded_target = &bounded.recent_target_outcomes[0];