Lochnair-Labs/irtt-rs
2
0 Fork 0
Code Issues 1 Pull Requests 4 Actions 4 Packages Projects Releases 6 Wiki Activity

feat(managed): add unified Tokio managed task #49

Merged
Lochnair merged 5 commits from feat/tokio-managed-task into main 2026-08-04 17:19:18 +00:00
Conversation 4 Commits 5 Files Changed 11
+317 -83
2 changed files with 317 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 99b6ddfead - Show all commits

No files matched your search

crates/irtt-client/src/managed/task.rs
+162 -48
View File
8
 @@ -171,6 +171,7 @@ enum TargetState {
}, },
Draining { Draining {
client: AsyncClient, client: AsyncClient,
drain_started_at: Instant,
deadline: Instant, deadline: Instant,
primary_end: ManagedTargetEndReason, primary_end: ManagedTargetEndReason,
cleanup_failure: Option<ManagedTargetFailure>, cleanup_failure: Option<ManagedTargetFailure>,
@@ -213,7 +214,6 @@ struct TargetRuntime {
server_addr: Arc<str>, server_addr: Arc<str>,
remote: Option<std::net::SocketAddr>, remote: Option<std::net::SocketAddr>,
counters: TargetCounters, counters: TargetCounters,
latest_committed_timeout: Option<Instant>,
send_waiting: bool, send_waiting: bool,
active: bool, active: bool,
state: TargetState, state: TargetState,
@@ -302,6 +302,12 @@ enum DriverState {
Finished, Finished,
} }
#[derive(Clone, Copy)]
enum OpenSessionFailureCleanup {
Drain,
Close,
}
/// The sole authoritative zero-or-many target driver. /// The sole authoritative zero-or-many target driver.
#[must_use = "ManagedClientTask must be awaited or deliberately dropped"] #[must_use = "ManagedClientTask must be awaited or deliberately dropped"]
pub struct ManagedClientTask { pub struct ManagedClientTask {
@@ -416,11 +422,14 @@ impl ManagedClientTask {
fn install_target_state(&mut self, index: usize, state: TargetState) { fn install_target_state(&mut self, index: usize, state: TargetState) {
let lifecycle = state.lifecycle(); let lifecycle = state.lifecycle();
let active = matches!(state, TargetState::Active { .. }); 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].active = active;
self.targets[index].state = state; self.targets[index].state = state;
if membership_changed { let now = Instant::now();
self.recompute_stagger_gate(Instant::now()); match (was_active, active) {
(false, true) => self.stagger_target_added(now),
(true, false) => self.stagger_target_removed(now),
_ => {}
} }
self.replace_status(); self.replace_status();
self.publish_event(ManagedEvent::TargetStateChanged { self.publish_event(ManagedEvent::TargetStateChanged {
@@ -493,6 +502,26 @@ impl ManagedClientTask {
self.finish_target(index, ManagedTargetEndReason::Failed(failure), None); 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> { fn begin_running(&mut self) -> Result<(), ManagedDriverFailure> {
tokio::runtime::Handle::try_current().map_err(|_| ManagedDriverFailure::NoTokioRuntime)?; tokio::runtime::Handle::try_current().map_err(|_| ManagedDriverFailure::NoTokioRuntime)?;
self.state = DriverState::Running; self.state = DriverState::Running;
@@ -631,9 +660,14 @@ impl ManagedClientTask {
match client.poll_timeouts_at(now) { match client.poll_timeouts_at(now) {
Ok(events) => self.publish_client_events(index, events), Ok(events) => self.publish_client_events(index, events),
Err(error) => { Err(error) => {
self.targets[index].counters.packets_sent = client.packets_sent(); return self.begin_open_session_failure(
self.fail_target(index, ManagedTargetFailurePhase::Timing, error); index,
return false; client,
ManagedTargetFailurePhase::Timing,
error,
now,
OpenSessionFailureCleanup::Close,
);
} }
} }
} }
1
 @@ -644,9 +678,14 @@ impl ManagedClientTask {
true true
} }
Poll::Ready(Err(error)) => { Poll::Ready(Err(error)) => {
self.targets[index].counters.packets_sent = client.packets_sent(); return self.begin_open_session_failure(
self.fail_target(index, ManagedTargetFailurePhase::Receiving, error); index,
return false; client,
ManagedTargetFailurePhase::Receiving,
error,
now,
OpenSessionFailureCleanup::Close,
);
} }
}; };
if client.is_peer_closed() { if client.is_peer_closed() {
@@ -668,6 +707,7 @@ impl ManagedClientTask {
} }
TargetState::Draining { TargetState::Draining {
mut client, mut client,
drain_started_at,
deadline, deadline,
primary_end, primary_end,
mut cleanup_failure, mut cleanup_failure,
@@ -680,6 +720,7 @@ impl ManagedClientTask {
if defer_work { if defer_work {
self.targets[index].state = TargetState::Draining { self.targets[index].state = TargetState::Draining {
client, client,
drain_started_at,
deadline, deadline,
primary_end, primary_end,
cleanup_failure, cleanup_failure,
@@ -692,12 +733,16 @@ impl ManagedClientTask {
let injected_receive_failure = self.take_drain_receive_failure(); let injected_receive_failure = self.take_drain_receive_failure();
#[cfg(not(test))] #[cfg(not(test))]
let injected_receive_failure = None; let injected_receive_failure = None;
let mut retained_state_changed = false;
if client if client
.next_probe_timeout_deadline() .next_probe_timeout_deadline()
.is_some_and(|timeout| timeout <= now) .is_some_and(|timeout| timeout <= now)
{ {
match client.poll_timeouts_at(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) => { Err(error) => {
self.targets[index].counters.packets_sent = client.packets_sent(); self.targets[index].counters.packets_sent = client.packets_sent();
cleanup_failure.get_or_insert_with(|| { cleanup_failure.get_or_insert_with(|| {
@@ -717,6 +762,7 @@ impl ManagedClientTask {
Poll::Pending => false, Poll::Pending => false,
Poll::Ready(Ok(events)) => { Poll::Ready(Ok(events)) => {
self.publish_client_events(index, events); self.publish_client_events(index, events);
retained_state_changed = true;
true true
} }
Poll::Ready(Err(error)) => { Poll::Ready(Err(error)) => {
@@ -739,11 +785,29 @@ impl ManagedClientTask {
self.finish_target(index, ManagedTargetEndReason::PeerClosed, cleanup_failure); self.finish_target(index, ManagedTargetEndReason::PeerClosed, cleanup_failure);
return false; 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 now >= deadline {
if received && post_deadline_receives_remaining > 1 { if received && post_deadline_receives_remaining > 1 {
post_deadline_receives_remaining -= 1; post_deadline_receives_remaining -= 1;
self.targets[index].state = TargetState::Draining { self.targets[index].state = TargetState::Draining {
client, client,
drain_started_at,
deadline, deadline,
primary_end, primary_end,
cleanup_failure, cleanup_failure,
@@ -756,6 +820,7 @@ impl ManagedClientTask {
} else { } else {
self.targets[index].state = TargetState::Draining { self.targets[index].state = TargetState::Draining {
client, client,
drain_started_at,
deadline, deadline,
primary_end, primary_end,
cleanup_failure, cleanup_failure,
@@ -824,23 +889,24 @@ impl ManagedClientTask {
index: usize, index: usize,
client: AsyncClient, client: AsyncClient,
primary_end: ManagedTargetEndReason, primary_end: ManagedTargetEndReason,
now: Instant, _now: Instant,
) -> bool { ) -> bool {
let latest = self.targets[index] let drain_started_at = Instant::now();
.latest_committed_timeout let Some(deadline) = self.drain_deadline(&client, drain_started_at) else {
.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 cleanup_failure = Some(duration_overflow_failure()); 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( self.install_target_state(
index, index,
TargetState::Draining { TargetState::Draining {
client, client,
drain_started_at,
deadline, deadline,
primary_end, primary_end,
cleanup_failure: None, cleanup_failure: None,
@@ -850,6 +916,14 @@ impl ManagedClientTask {
true true
} }
fn drain_deadline(&self, client: &AsyncClient, drain_started_at: Instant) -> Option<Instant> {
client
.latest_probe_timeout_deadline()
.unwrap_or(drain_started_at)
.max(drain_started_at)
.checked_add(self.config.final_drain)
}
fn begin_close( fn begin_close(
&mut self, &mut self,
index: usize, index: usize,
@@ -914,15 +988,50 @@ impl ManagedClientTask {
minimum.map(|interval| (active, stagger_spacing(interval, active))) minimum.map(|interval| (active, stagger_spacing(interval, active)))
} }
fn recompute_stagger_gate(&mut self, now: Instant) { fn stagger_target_added(&mut self, now: Instant) {
self.send_gate = self let Some(existing) = self.send_gate.filter(|gate| *gate > now) else {
self.send_gate = None;
return;
};
let candidate = self
.last_stagger_send .last_stagger_send
.zip(self.active_stagger_spacing()) .zip(self.active_stagger_spacing())
.and_then(|(last, (_, spacing))| last.checked_add(spacing)) .and_then(|(last, (_, spacing))| last.checked_add(spacing))
.filter(|gate| *gate > now); .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 state = mem::replace(&mut self.targets[index].state, TargetState::Terminal);
let TargetState::Active { mut client } = state else { let TargetState::Active { mut client } = state else {
self.targets[index].state = state; self.targets[index].state = state;
@@ -940,18 +1049,17 @@ impl ManagedClientTask {
let result = client.poll_send_probe(cx); let result = client.poll_send_probe(cx);
let after = client.packets_sent(); let after = client.packets_sent();
let accepted = after > before; 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 let schedule_error = accepted
.then(|| client.skip_missed_probe_slots_at(Instant::now())) .then(|| client.skip_missed_probe_slots_at(Instant::now()))
.transpose() .transpose()
.err(); .err();
self.targets[index].counters.packets_sent = after; 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 { match result {
Poll::Pending => { Poll::Pending => {
self.targets[index].send_waiting = true; self.targets[index].send_waiting = true;
@@ -963,7 +1071,19 @@ impl ManagedClientTask {
self.targets[index].state = TargetState::Active { client }; self.targets[index].state = TargetState::Active { client };
self.publish_client_events(index, events); self.publish_client_events(index, events);
if let Some(error) = schedule_error { 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 } SendResult::Failed { accepted }
} else { } else {
SendResult::Ready { accepted } SendResult::Ready { accepted }
@@ -971,7 +1091,14 @@ impl ManagedClientTask {
} }
Poll::Ready(Err(error)) => { Poll::Ready(Err(error)) => {
self.targets[index].send_waiting = false; 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 } SendResult::Failed { accepted }
chatgpt-codex-connector[bot] commented 2026-08-04 16:36:30 +00:00 (Migrated from github.com)
Review
Copy link
Copy source

P2 Badge Close active sessions after send failures

When an open target encounters a send-side error, this branch immediately marks it terminal and drops the AsyncClient without sending its prepared close packet. This occurs through ordinary configuration rather than only transport failure—for example, with max_pending_probes = 1 and an interval shorter than probe_timeout, the second send returns PendingLimitExceeded while the first probe remains outstanding—so the server retains an allocated session until its own timeout and the outstanding probe is never drained. Preserve the classified failure as the primary end reason while driving the target through drain/close cleanup.

Useful? React with 👍 / 👎.

**<sub><sub>![P2 Badge](https://img.shields.io/badge/P2-yellow?style=flat)</sub></sub> Close active sessions after send failures** When an open target encounters a send-side error, this branch immediately marks it terminal and drops the `AsyncClient` without sending its prepared close packet. This occurs through ordinary configuration rather than only transport failure—for example, with `max_pending_probes = 1` and an interval shorter than `probe_timeout`, the second send returns `PendingLimitExceeded` while the first probe remains outstanding—so the server retains an allocated session until its own timeout and the outstanding probe is never drained. Preserve the classified failure as the primary end reason while driving the target through drain/close cleanup. Useful? React with 👍 / 👎.
} }
} }
@@ -992,24 +1119,12 @@ impl ManagedClientTask {
if !matches!(self.targets[index].state, TargetState::Active { .. }) { if !matches!(self.targets[index].state, TargetState::Active { .. }) {
continue; continue;
} }
#[cfg(test)]
let stagger_spacing = self.active_stagger_spacing(); 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 { if result == SendResult::NotAttempted {
continue; continue;
} }
self.stagger_remaining = 0; 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 { .. }); return matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. });
} }
self.stagger_remaining > 0 self.stagger_remaining > 0
@@ -1028,7 +1143,7 @@ impl ManagedClientTask {
let index = self.send_cursor % self.targets.len(); let index = self.send_cursor % self.targets.len();
self.send_cursor = (self.send_cursor + 1) % self.targets.len(); self.send_cursor = (self.send_cursor + 1) % self.targets.len();
self.burst_remaining -= 1; 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 |= matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. });
} }
immediate || self.burst_remaining > 0 immediate || self.burst_remaining > 0
@@ -1319,7 +1434,6 @@ fn build_task(
server_addr: Arc::from(target.server_addr), server_addr: Arc::from(target.server_addr),
remote: None, remote: None,
counters: TargetCounters::default(), counters: TargetCounters::default(),
latest_committed_timeout: None,
send_waiting: false, send_waiting: false,
active: false, active: false,
state: TargetState::Pending { client_config }, state: TargetState::Pending { client_config },
crates/irtt-client/src/managed/tests.rs
+155 -35
View File
@@ -1,5 +1,6 @@
use std::{ use std::{
future::{poll_fn, Future}, future::{poll_fn, Future},
mem,
net::{SocketAddr, UdpSocket}, net::{SocketAddr, UdpSocket},
pin::Pin, pin::Pin,
sync::{Arc, Condvar, Mutex}, sync::{Arc, Condvar, Mutex},
@@ -524,14 +525,19 @@ fn pre_poll_stop() {
} }
#[test] #[test]
fn finite_one_target_lifecycle() { fn finite_immediate_replies_use_retained_drain_deadline() {
let server = start_server(ServerBehavior::Echo, None); let server = start_server(ServerBehavior::Echo, None);
let (task, _) = ManagedClient::task( let mut managed = config(ManagedPacing::Staggered);
config(ManagedPacing::Staggered), managed.client.probe_timeout = Duration::from_secs(3);
vec![target("one", server.addr)], managed.final_drain = Duration::from_millis(30);
) let (task, _) = ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap();
.unwrap(); let started_at = Instant::now();
let outcome = runtime().block_on(task); 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(); let records = server.finish();
assert_eq!(outcome.end_reason, ManagedEndReason::TargetsComplete); assert_eq!(outcome.end_reason, ManagedEndReason::TargetsComplete);
assert!(matches!( assert!(matches!(
@@ -793,7 +799,8 @@ fn staggered_pacing_uses_all_active_negotiated_intervals() {
#[test] #[test]
fn stagger_gate_tracks_active_membership() { 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_gate = Arc::new(PacketGate::default());
let second = start_server_with_gates( let second = start_server_with_gates(
ServerBehavior::Echo, ServerBehavior::Echo,
@@ -803,7 +810,7 @@ fn stagger_gate_tracks_active_membership() {
None, None,
); );
let mut managed = config(ManagedPacing::Staggered); 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.interval = Duration::from_secs(1);
managed.client.negotiation_policy = NegotiationPolicy::Loose; managed.client.negotiation_policy = NegotiationPolicy::Loose;
managed.client.open_timeouts = vec![Duration::from_secs(2)]; 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 task.targets[0].counters.packets_sent == 1
}); });
let one_target_last = task.last_stagger_send.unwrap(); let one_target_last = task.last_stagger_send.unwrap();
let one_target_gate = task.send_gate.unwrap();
assert_eq!( assert_eq!(
task.send_gate, task.send_gate,
one_target_last.checked_add(Duration::from_secs(1)) one_target_last.checked_add(Duration::from_millis(100))
); );
second_gate.release(); second_gate.release();
drive_task_until(&runtime, &mut task, |task| task.active_count() == 2); drive_task_until(&runtime, &mut task, |task| task.active_count() == 2);
let last = task.last_stagger_send.unwrap(); let last = task.last_stagger_send.unwrap();
let gate = task.send_gate.unwrap(); let gate = task.send_gate.unwrap();
assert_eq!(gate, last.checked_add(Duration::from_millis(500)).unwrap()); assert_eq!(gate, last.checked_add(Duration::from_millis(50)).unwrap());
assert!(gate > last && gate < last.checked_add(Duration::from_secs(1)).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| { drive_task_until(&runtime, &mut task, |task| {
task.targets[1].counters.packets_sent == 1 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()); drop(handle.stop());
let outcome = runtime.block_on(task); let outcome = runtime.block_on(task);
assert_eq!(outcome.failed_target_outcomes, 0); assert_eq!(outcome.failed_target_outcomes, 0);
@@ -907,37 +939,125 @@ fn overdue_queued_reply_is_lost_then_late() {
} }
#[test] #[test]
fn committed_probe_graceful_drain() { fn retained_probe_deadline_shortens_after_late_reply() {
let server = start_server(ServerBehavior::DelayedEcho(Duration::from_millis(25)), None); let reply_gate = Arc::new(PacketGate::default());
let mut config = config(ManagedPacing::Staggered); let server = start_server_with_gates(
config.completion = ManagedCompletionPolicy::ExplicitStop; ServerBehavior::Echo,
config.client.duration = None; None,
config.client.interval = Duration::from_millis(100); None,
config.client.probe_timeout = Duration::from_millis(80); None,
config.final_drain = Duration::from_millis(5); Some(Arc::clone(&reply_gate)),
let (task, handle) = ManagedClient::task(config, vec![target("one", server.addr)]).unwrap(); );
let stop_handle = handle.clone(); let mut managed = config(ManagedPacing::Staggered);
let seen = Arc::clone(&server.probe_seen); managed.completion = ManagedCompletionPolicy::ExplicitStop;
let stopper = thread::spawn(move || { managed.client.duration = None;
let (flag, ready) = &*seen; managed.client.interval = Duration::from_secs(1);
let guard = flag.lock().unwrap(); managed.client.probe_timeout = Duration::from_millis(100);
let (guard, timeout) = ready managed.final_drain = Duration::from_millis(200);
.wait_timeout_while(guard, Duration::from_secs(2), |seen| !*seen) let final_drain = managed.final_drain;
.unwrap(); let (task, handle) = ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap();
let probe_seen = *guard; let runtime = runtime();
drop(guard); let mut task = Box::pin(task);
drop(stop_handle.stop()); drive_task_until(&runtime, &mut task, |task| {
assert!(probe_seen && !timeout.timed_out()); task.targets[0].counters.packets_sent == 1
}); });
let outcome = runtime().block_on(task); drop(handle.stop());
stopper.join().unwrap(); 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 records = server.finish();
let target = &outcome.recent_target_outcomes[0]; let target = &outcome.recent_target_outcomes[0];
assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested);
assert!(matches!(target.end_reason, ManagedTargetEndReason::Stopped)); assert!(matches!(target.end_reason, ManagedTargetEndReason::Stopped));
assert_eq!(target.packets_sent, 1); 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!(target.replies_received, 1);
assert_eq!(probes(&records).len(), 1); assert_eq!(probes(&records).len(), 1);
assert!(has_close(&records));
} }
#[test] #[test]