fix(managed): harden stop and deadline ordering

crates/irtt-client/src/async_client.rs

@@ -86,6 +86,8 @@ pub(crate) struct AsyncOpenState {
     deadline: Option<Instant>,
     deadline_timer: Option<Pin<Box<time::Sleep>>>,
     request_submitted: bool,
+    stop_after_current_attempt: bool,
+    stopped_after_current_attempt: bool,
     buffer: [u8; MAX_OPEN_PACKET_SIZE],
     cleanup: Option<AsyncOpenCleanup>,
 }
@@ -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<bool>,
     fail_peer_close_dscp: Cell<bool>,
     fail_dscp_restore: Cell<bool>,
+    pause_open_cleanup_before_send: Cell<bool>,
     pause_open_before_writable: Cell<bool>,
     pause_open_before_readable: Cell<bool>,
     pause_probe_before_writable: Cell<bool>,
@@ -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<io::Result<()>> {
+        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()));

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,

crates/irtt-client/src/managed/task.rs

@@ -215,6 +215,7 @@ struct TargetRuntime {
     counters: TargetCounters,
     latest_committed_timeout: Option<Instant>,
     send_waiting: bool,
+    active: bool,
     state: TargetState,
 }
 
@@ -318,6 +319,7 @@ pub struct ManagedClientTask {
     burst_remaining: usize,
     stagger_remaining: usize,
     send_gate: Option<Instant>,
+    last_stagger_send: Option<Instant>,
     stop_observed: bool,
     final_outcome: Option<Arc<ManagedOutcome>>,
     _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,27 +565,55 @@ impl ManagedClientTask {
                 mut open,
             } => {
                 if self.state == DriverState::Stopping {
+                    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 => {
                         self.targets[index].state = TargetState::Opening { client, open };
                         false
                     }
                     Poll::Ready(Ok(OpenOutcome::Started { event, .. })) => {
-                        self.install_target_state(index, TargetState::Active { client });
                         self.publish_client_events(index, vec![event]);
+                        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)) => {
+                        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(
+            let cleanup_failure = Some(duration_overflow_failure());
-                index,
+            return self.begin_close(index, client, primary_end, cleanup_failure, now);
-                ManagedTargetFailurePhase::Timing,
-                ClientError::DurationOverflow,
-            );
-            return false;
         };
         self.install_target_state(
             index,
@@ -823,10 +855,13 @@ impl ManagedClientTask {
         index: usize,
         client: AsyncClient,
         primary_end: ManagedTargetEndReason,
-        cleanup_failure: Option<ManagedTargetFailure>,
+        mut cleanup_failure: Option<ManagedTargetFailure>,
         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
+        self.targets.iter().filter(|target| target.active).count()
-            .iter()
-            .filter(|target| matches!(target.state, TargetState::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() {
-                if let Some((_active, spacing)) = stagger_spacing {
                 #[cfg(test)]
+                if let Some((_active, spacing)) = stagger_spacing {
                     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,

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<Mutex<Vec<PacketRecord>>>,
     probe_seen: Arc<(Mutex<bool>, Condvar)>,
+    reply_sent: Arc<(Mutex<bool>, Condvar)>,
     thread: JoinHandle<()>,
 }
 
@@ -55,14 +56,47 @@ impl TestServer {
     }
 }
 
+#[derive(Default)]
+struct PacketGate {
+    seen: Mutex<bool>,
+    seen_ready: Condvar,
+    released: Mutex<bool>,
+    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<Vec<u8>>) -> TestServer {
     start_server_negotiating(behavior, key, None)
 }
-
 fn start_server_negotiating(
     behavior: ServerBehavior,
     key: Option<Vec<u8>>,
     interval: Option<Duration>,
+) -> TestServer {
+    start_server_with_gates(behavior, key, interval, None, None)
+}
+
+fn start_server_with_gates(
+    behavior: ServerBehavior,
+    key: Option<Vec<u8>>,
+    interval: Option<Duration>,
+    open_gate: Option<Arc<PacketGate>>,
+    reply_gate: Option<Arc<PacketGate>>,
 ) -> 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<F: Future>(future: Pin<&mut F>) -> Poll<F::Output> {
     Future::poll(future, &mut Context::from_waker(Waker::noop()))
 }
+fn drive_task_until(
+    runtime: &Runtime,
+    task: &mut Pin<Box<ManagedClientTask>>,
+    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<bool>, 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<(ManagedEvent, Arc<ManagedStatus>)>>,
+) -> Vec<ClientEvent> {
+    observations
+        .lock()
+        .unwrap()
+        .iter()
+        .filter_map(|(event, _)| match event {
+            ManagedEvent::Client { event, .. } => Some(event.clone()),
+            _ => None,
+        })
+        .collect()
+}
 
 fn probes(records: &[PacketRecord]) -> Vec<Instant> {
     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::<Vec<_>>();
+    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];