Files
irtt-rs/crates/irtt-client/src/managed/task.rs
T
Nils Andreas Svee e7de8bdb30 refactor(client): split managed task into coherent internal modules
Group control and construction, target lifecycle, target-set transactions,
send pacing and commitment, and bounded timeout processing into private
children of the existing task owner. Keep durable status, outcome history,
top-level polling, sealing, and abandonment in task.rs.

Preserve phase-owned resources, stop/update synchronization, publication
ordering, pacing, and timeout fairness. Move the existing private gate and
timeout tests with their implementation without adding coverage or hooks.
2026-10-05 00:24:15 +07:00

648 lines
22 KiB
Rust

//! Task ownership, durable status/history, and top-level polling and sealing.
//!
//! Private children implement control, target lifecycle, updates, pacing, and
//! bounded timeout work on this same task; they do not introduce other owners.
mod control;
mod pacing;
mod target;
mod timeouts;
mod update;
pub use control::{ManagedClient, ManagedClientHandle, ManagedCommandReceipt, ManagedStopReceipt};
use std::{
collections::VecDeque,
future::Future,
pin::Pin,
sync::Arc,
task::{Context, Poll},
time::Instant,
};
use tokio::{
sync::{broadcast, mpsc, watch},
time::{self, Sleep},
};
use super::{
ManagedClientConfig, ManagedCommandApplyError, ManagedCompletionPolicy, ManagedDriverFailure,
ManagedEndReason, ManagedEvent, ManagedLifecycle, ManagedOutcome, ManagedPacing, ManagedStatus,
ManagedTargetEndReason, ManagedTargetLifecycle, ManagedTargetOutcome, ManagedTargetStatus,
};
use control::{ManagedCommand, StopSignal};
use pacing::gated_send_deadline;
use target::{TargetPhase, TargetRuntime};
use timeouts::TimeoutBacklogEntry;
const TARGET_WORK_BUDGET: usize = 128;
type WakeFuture = Pin<Box<dyn Future<Output = Option<watch::Receiver<()>>> + Send + 'static>>;
fn arm_wake(mut receiver: watch::Receiver<()>) -> WakeFuture {
Box::pin(async move {
receiver.changed().await.ok()?;
Some(receiver)
})
}
#[derive(Default)]
struct OutcomeHistory {
limit: usize,
recent: VecDeque<ManagedTargetOutcome>,
total: u64,
successful: u64,
failed: u64,
peer_closed: u64,
discarded: u64,
}
impl OutcomeHistory {
fn new(limit: usize) -> Self {
Self {
limit,
..Self::default()
}
}
fn record(&mut self, outcome: ManagedTargetOutcome) {
self.total = self.total.saturating_add(1);
match &outcome.end_reason {
ManagedTargetEndReason::Failed(_) => self.failed = self.failed.saturating_add(1),
ManagedTargetEndReason::PeerClosed => {
self.successful = self.successful.saturating_add(1);
self.peer_closed = self.peer_closed.saturating_add(1);
}
ManagedTargetEndReason::TestComplete
| ManagedTargetEndReason::NoTestComplete
| ManagedTargetEndReason::Removed
| ManagedTargetEndReason::Replaced
| ManagedTargetEndReason::Stopped => {
self.successful = self.successful.saturating_add(1);
}
}
if self.limit == 0 {
self.discarded = self.discarded.saturating_add(1);
return;
}
if self.recent.len() == self.limit {
self.recent.pop_front();
self.discarded = self.discarded.saturating_add(1);
}
self.recent.push_back(outcome);
}
fn recent(&self) -> Arc<[ManagedTargetOutcome]> {
Arc::from(
self.recent
.iter()
.cloned()
.collect::<Vec<_>>()
.into_boxed_slice(),
)
}
fn outcome(
&self,
end_reason: ManagedEndReason,
applied_command_sequence: u64,
) -> ManagedOutcome {
ManagedOutcome {
end_reason,
applied_command_sequence,
total_target_outcomes: self.total,
successful_target_outcomes: self.successful,
failed_target_outcomes: self.failed,
peer_closed_target_outcomes: self.peer_closed,
discarded_target_outcomes: self.discarded,
recent_target_outcomes: self.recent(),
}
}
}
struct TaskResources {
status: watch::Sender<Arc<ManagedStatus>>,
events: Option<broadcast::Sender<ManagedEvent>>,
stop: Arc<StopSignal>,
}
enum TaskPhase {
NotStarted,
Running,
Stopping,
Completed { outcome: Arc<ManagedOutcome> },
Failed { outcome: Arc<ManagedOutcome> },
Abandoned,
}
impl TaskPhase {
fn lifecycle(&self) -> ManagedLifecycle {
match self {
Self::NotStarted => ManagedLifecycle::NotStarted,
Self::Running => ManagedLifecycle::Running,
Self::Stopping => ManagedLifecycle::Stopping,
Self::Completed { .. } => ManagedLifecycle::Completed,
Self::Failed { .. } => ManagedLifecycle::Failed,
Self::Abandoned => ManagedLifecycle::Abandoned,
}
}
fn final_outcome(&self) -> Option<Arc<ManagedOutcome>> {
match self {
Self::Completed { outcome } | Self::Failed { outcome } => Some(Arc::clone(outcome)),
_ => None,
}
}
fn is_terminal(&self) -> bool {
matches!(
self,
Self::Completed { .. } | Self::Failed { .. } | Self::Abandoned
)
}
fn is_running(&self) -> bool {
matches!(self, Self::Running)
}
fn is_stopping(&self) -> bool {
matches!(self, Self::Stopping)
}
}
/// The sole authoritative zero-or-many target driver.
#[must_use = "ManagedClientTask must be awaited or deliberately dropped"]
pub struct ManagedClientTask {
#[cfg(test)]
timeout_inspections: std::cell::RefCell<Vec<usize>>,
phase: TaskPhase,
config: ManagedClientConfig,
targets: Vec<TargetRuntime>,
commands: mpsc::Receiver<ManagedCommand>,
history: OutcomeHistory,
resources: Option<TaskResources>,
wake: Option<WakeFuture>,
timer: Option<Pin<Box<Sleep>>>,
cursor: usize,
timeout_cursor: usize,
timeout_backlog: VecDeque<TimeoutBacklogEntry>,
scan_remaining: usize,
send_cursor: usize,
burst_remaining: usize,
stagger_remaining: usize,
send_gate: Option<Instant>,
last_stagger_send: Option<Instant>,
// Explicit stop observation is independent of natural stopping or terminality.
stop_observed: bool,
next_generation: u64,
applied_command_sequence: u64,
}
impl ManagedClientTask {
fn resources(&self) -> &TaskResources {
self.resources
.as_ref()
.expect("managed task resources exist before terminal sealing")
}
fn publish_event(&self, event: ManagedEvent) {
if let Some(events) = &self.resources().events {
let _ = events.send(event);
}
}
fn snapshot(&self) -> Arc<ManagedStatus> {
let mut connecting = 0;
let mut opening = 0;
let mut active = 0;
let mut draining = 0;
let mut closing = 0;
let mut terminal = 0;
let targets = self
.targets
.iter()
.map(|target| {
let lifecycle = target.phase().lifecycle();
match lifecycle {
ManagedTargetLifecycle::Pending => {}
ManagedTargetLifecycle::Connecting => connecting += 1,
ManagedTargetLifecycle::Opening => opening += 1,
ManagedTargetLifecycle::Active => active += 1,
ManagedTargetLifecycle::Draining => draining += 1,
ManagedTargetLifecycle::Closing => closing += 1,
ManagedTargetLifecycle::Terminal => terminal += 1,
}
ManagedTargetStatus {
target: target.instance.clone(),
desired: target.membership.is_desired(),
lifecycle,
outcome: match target.phase() {
TargetPhase::Terminal { outcome } => Some(Arc::clone(outcome)),
_ => None,
},
server_addr: Arc::clone(&target.server_addr),
remote: target.remote,
}
})
.collect::<Vec<_>>();
Arc::new(ManagedStatus {
lifecycle: self.phase.lifecycle(),
stop_requested: self.stop_observed,
applied_command_sequence: self.applied_command_sequence,
desired_target_count: self
.targets
.iter()
.filter(|target| target.membership.is_desired())
.count(),
connecting_target_count: connecting,
opening_target_count: opening,
active_target_count: active,
draining_target_count: draining,
closing_target_count: closing,
terminal_target_count: terminal,
total_target_outcomes: self.history.total,
successful_target_outcomes: self.history.successful,
failed_target_outcomes: self.history.failed,
peer_closed_target_outcomes: self.history.peer_closed,
discarded_target_outcomes: self.history.discarded,
targets: Arc::from(targets.into_boxed_slice()),
recent_target_outcomes: self.history.recent(),
final_outcome: self.phase.final_outcome(),
})
}
fn replace_status(&self) {
self.resources().status.send_replace(self.snapshot());
}
fn begin_running(&mut self) -> Result<(), ManagedDriverFailure> {
tokio::runtime::Handle::try_current().map_err(|_| ManagedDriverFailure::NoTokioRuntime)?;
self.phase = TaskPhase::Running;
self.replace_status();
self.publish_event(ManagedEvent::Started);
Ok(())
}
fn observe_stop(&mut self) {
if self.stop_observed {
return;
}
self.stop_observed = true;
self.resources().stop.begin_stopping();
let lifecycle_transition = !self.phase.is_stopping();
self.phase = TaskPhase::Stopping;
self.replace_status();
if lifecycle_transition {
self.publish_event(ManagedEvent::Stopping);
}
self.scan_remaining = self.targets.len();
self.burst_remaining = 0;
}
fn poll_target_pass(&mut self, cx: &mut Context<'_>, now: Instant) -> bool {
if self.targets.is_empty() {
return false;
}
if self.scan_remaining == 0 {
self.scan_remaining = self.targets.len();
}
let work = self.scan_remaining.min(TARGET_WORK_BUDGET);
let mut immediate = false;
for _ in 0..work {
let index = self.cursor % self.targets.len();
self.cursor = (self.cursor + 1) % self.targets.len();
self.scan_remaining -= 1;
immediate |= self.poll_target(index, cx, now);
}
immediate || self.scan_remaining > 0
}
fn all_targets_terminal(&self) -> bool {
self.targets
.iter()
.all(|target| matches!(target.phase(), TargetPhase::Terminal { .. }))
}
fn next_deadline(&self) -> Option<Instant> {
let mut non_send_deadline = None;
let mut send_deadline = None;
for target in &self.targets {
match target.phase() {
TargetPhase::Active {
client,
schedule,
send_waiting,
} => {
non_send_deadline = non_send_deadline
.into_iter()
.chain(client.next_probe_timeout_deadline())
.chain(schedule.end_deadline())
.min();
if target.membership.is_desired() && !send_waiting {
send_deadline = send_deadline
.into_iter()
.chain(schedule.next_send_deadline())
.min();
}
}
TargetPhase::Draining {
client, deadline, ..
} => {
non_send_deadline = non_send_deadline
.into_iter()
.chain(client.next_probe_timeout_deadline())
.chain(Some(*deadline))
.min();
}
TargetPhase::Closing { deadline, .. } => {
non_send_deadline = non_send_deadline.into_iter().chain(Some(*deadline)).min();
}
_ => {}
}
}
if self.phase.is_running()
&& self.config.pacing == ManagedPacing::Staggered
&& self.active_count() > 0
{
send_deadline = gated_send_deadline(send_deadline, self.send_gate);
}
non_send_deadline.into_iter().chain(send_deadline).min()
}
fn register_timer(&mut self, cx: &mut Context<'_>) -> bool {
let Some(deadline) = self.next_deadline() else {
self.timer = None;
return false;
};
let timer = self
.timer
.get_or_insert_with(|| Box::pin(time::sleep_until(deadline.into())));
timer.as_mut().reset(deadline.into());
timer.as_mut().poll(cx).is_ready()
}
fn register_stop_wake(&mut self, cx: &mut Context<'_>) -> bool {
let Some(mut wake) = self.wake.take() else {
return false;
};
match wake.as_mut().poll(cx) {
Poll::Pending => {
self.wake = Some(wake);
false
}
Poll::Ready(Some(receiver)) => {
self.wake = Some(arm_wake(receiver));
self.resources().stop.is_requested()
}
Poll::Ready(None) => false,
}
}
fn seal(&mut self, mut end_reason: ManagedEndReason, failed: bool) -> Poll<ManagedOutcome> {
// Sealing is the final stop-observation point. A stop that reaches the
// latch before this terminal linearization becomes durable even when
// quiescence had already entered `Stopping`.
if self.resources().stop.is_requested() {
self.observe_stop();
}
if self.stop_observed && matches!(end_reason, ManagedEndReason::TargetsComplete) {
end_reason = ManagedEndReason::StopRequested;
}
// Closing admission before closing and draining the receiver makes every
// accepted command task-owned. Tokio still permits buffered commands
// to be drained after `Receiver::close`, but rejects later sends.
self.resources().stop.close_updates();
self.commands.close();
while let Ok(ManagedCommand::UpdateTargets {
acknowledgement, ..
}) = self.commands.try_recv()
{
let error = match &end_reason {
ManagedEndReason::DriverFailed(failure) => {
ManagedCommandApplyError::DriverFailed(failure.clone())
}
_ => ManagedCommandApplyError::Stopping,
};
let _ = acknowledgement.send(Err(error));
}
let outcome = Arc::new(
self.history
.outcome(end_reason, self.applied_command_sequence),
);
self.phase = if failed {
TaskPhase::Failed {
outcome: Arc::clone(&outcome),
}
} else {
TaskPhase::Completed {
outcome: Arc::clone(&outcome),
}
};
self.replace_status();
self.publish_event(if failed {
ManagedEvent::Failed {
outcome: Arc::clone(&outcome),
}
} else {
ManagedEvent::Completed {
outcome: Arc::clone(&outcome),
}
});
if let Some(resources) = self.resources.as_mut() {
resources.events.take();
}
self.wake = None;
self.timer = None;
self.resources.take();
Poll::Ready((*outcome).clone())
}
fn fail_driver(&mut self, failure: ManagedDriverFailure) -> Poll<ManagedOutcome> {
self.seal(ManagedEndReason::DriverFailed(failure), true)
}
}
impl Future for ManagedClientTask {
type Output = ManagedOutcome;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.as_mut().get_mut();
if this.phase.is_terminal() {
panic!("ManagedClientTask polled after completion");
}
if matches!(this.phase, TaskPhase::NotStarted) {
if this.resources().stop.is_requested() {
this.observe_stop();
} else if let Err(failure) = this.begin_running() {
return this.fail_driver(failure);
}
}
if this.resources().stop.is_requested() {
this.observe_stop();
}
let mut immediate = this.process_commands(cx);
if this.resources().stop.is_requested() {
this.observe_stop();
}
let now = Instant::now();
immediate |= this.poll_timeout_pass(now);
immediate |= this.poll_target_pass(cx, now);
this.prune_undesired_terminal();
if this.phase.is_running() {
immediate |= match this.config.pacing {
ManagedPacing::Staggered => this.poll_staggered_send(cx, now),
ManagedPacing::Burst => this.poll_burst_sends(cx, now),
};
}
if this.resources().stop.is_requested() {
this.observe_stop();
}
if this.all_targets_terminal() {
match this.phase {
TaskPhase::Stopping => {
return this.seal(
if this.stop_observed {
ManagedEndReason::StopRequested
} else {
ManagedEndReason::TargetsComplete
},
false,
);
}
TaskPhase::Running
if this.config.completion == ManagedCompletionPolicy::FinishWhenQuiescent =>
{
this.phase = TaskPhase::Stopping;
this.resources().stop.begin_stopping();
this.replace_status();
this.publish_event(ManagedEvent::Stopping);
return this.seal(ManagedEndReason::TargetsComplete, false);
}
_ => {}
}
}
immediate |= this.register_stop_wake(cx);
if !immediate {
immediate |= this.register_timer(cx);
}
if immediate {
cx.waker().wake_by_ref();
}
Poll::Pending
}
}
impl Drop for ManagedClientTask {
fn drop(&mut self) {
if self.phase.is_terminal() || self.resources.is_none() {
return;
}
self.phase = TaskPhase::Abandoned;
self.resources().stop.close_updates();
self.commands.close();
self.replace_status();
self.publish_event(ManagedEvent::Abandoned);
if let Some(resources) = self.resources.as_mut() {
resources.events.take();
}
self.resources.take();
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
async_client::AsyncClient,
managed::{schedule::ProbeSchedule, ManagedTargetConfig},
ClientConfig,
};
use std::time::Duration;
// Real EAGAIN needs privileged queue shaping. Inspect only the owning
// deadline selection here, using a production socket without sending.
#[test]
fn finite_lifetime_deadline_survives_send_waiting_and_stagger_gating() {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
runtime.block_on(async {
let start = Instant::now();
let end = start + Duration::from_millis(500);
let accepted = crate::AcceptedSessionParameters {
duration: Some(Duration::from_millis(500)),
interval: Duration::from_millis(1),
length: 0,
received_stats: irtt_proto::ReceivedStats::None,
stamp_at: irtt_proto::StampAt::None,
clock: irtt_proto::Clock::Both,
dscp: 0,
server_fill: None,
};
for pacing in [ManagedPacing::Burst, ManagedPacing::Staggered] {
let (mut task, _handle) = ManagedClient::task(
ManagedClientConfig {
pacing,
..ManagedClientConfig::default()
},
vec![ManagedTargetConfig::new("target", "127.0.0.1:9")],
)
.unwrap();
let client = AsyncClient::connect("127.0.0.1:9", ClientConfig::default())
.await
.unwrap();
task.phase = TaskPhase::Running;
task.install_target_phase(
0,
TargetPhase::Active {
client,
schedule: ProbeSchedule::new(start, &accepted).unwrap(),
send_waiting: false,
},
);
task.send_gate = Some(end + Duration::from_secs(1));
for send_waiting in [false, true] {
let TargetPhase::Active {
schedule,
send_waiting: waiting,
..
} = task.targets[0].phase_mut()
else {
unreachable!()
};
*waiting = send_waiting;
*schedule = ProbeSchedule::new(start, &accepted).unwrap();
let expected = if pacing == ManagedPacing::Burst && !send_waiting {
start
} else {
end
};
assert_eq!(task.next_deadline(), Some(expected));
let TargetPhase::Active { schedule, .. } = task.targets[0].phase_mut() else {
unreachable!()
};
schedule.permit_probe_at(end);
assert_eq!(task.next_deadline(), None);
}
let mut continuous = accepted.clone();
continuous.duration = None;
let TargetPhase::Active { schedule, .. } = task.targets[0].phase_mut() else {
unreachable!()
};
*schedule = ProbeSchedule::new(start, &continuous).unwrap();
assert_eq!(task.next_deadline(), None);
}
});
}
}