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