From fd3c29e63a889e9040aa83e939729d95ff9d3b36 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 20:38:52 +0700 Subject: [PATCH] feat(managed): add unified Tokio managed task --- crates/irtt-client/src/async_client.rs | 35 + crates/irtt-client/src/lib.rs | 7 +- crates/irtt-client/src/managed/group.rs | 57 +- crates/irtt-client/src/managed/mod.rs | 40 +- crates/irtt-client/src/managed/model.rs | 329 +++++ crates/irtt-client/src/managed/target_id.rs | 47 + crates/irtt-client/src/managed/task.rs | 1220 +++++++++++++++++++ crates/irtt-client/src/probe.rs | 15 + crates/irtt-client/src/session/machine.rs | 46 +- 9 files changed, 1720 insertions(+), 76 deletions(-) create mode 100644 crates/irtt-client/src/managed/model.rs create mode 100644 crates/irtt-client/src/managed/target_id.rs create mode 100644 crates/irtt-client/src/managed/task.rs diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs index d3616d8..4c561c3 100644 --- a/crates/irtt-client/src/async_client.rs +++ b/crates/irtt-client/src/async_client.rs @@ -682,6 +682,41 @@ impl AsyncClient { self.machine.is_peer_closed() } + pub(crate) fn remote_addr(&self) -> SocketAddr { + self.remote + } + + pub(crate) fn packets_sent(&self) -> u64 { + self.machine.packets_sent() + } + + pub(crate) fn next_probe_timeout_deadline(&self) -> Option { + self.machine.next_probe_timeout_deadline() + } + + pub(crate) fn latest_probe_timeout_deadline(&self) -> Option { + self.machine.latest_probe_timeout_deadline() + } + + pub(crate) fn discard_prepared_probe(&mut self) { + self.prepared_probe = None; + } + + pub(crate) fn skip_missed_probe_slots_at(&mut self, now: Instant) -> Result<(), ClientError> { + let Some(schedule) = self.schedule.as_mut() else { + return Ok(()); + }; + let Some(deadline) = schedule.next_send_deadline() else { + return Ok(()); + }; + if deadline > now { + return Ok(()); + } + let commit = schedule.preflight_managed_commit(deadline, now)?; + schedule.commit(commit); + Ok(()) + } + fn prepare_async_open( &mut self, machine: PreparedOpenAcceptance, diff --git a/crates/irtt-client/src/lib.rs b/crates/irtt-client/src/lib.rs index e0971c3..743b5f9 100644 --- a/crates/irtt-client/src/lib.rs +++ b/crates/irtt-client/src/lib.rs @@ -59,7 +59,7 @@ mod client; mod config; mod error; mod event; -mod managed; +pub mod managed; mod metadata; mod probe; mod receive; @@ -80,14 +80,15 @@ pub use event::{ ClientEvent, OneWayDelaySample, OpenOutcome, PacketMeta, ReceivedStatsSample, RttSample, ServerTiming, SignedDuration, WarningKind, }; -pub use managed::{ +pub use managed::legacy::{ CancellationToken, EventHub, EventSubscription, ManagedClient, ManagedClientGroup, ManagedClientGroupConfig, ManagedClientGroupSession, ManagedClientSession, ManagedGroupCompletionPolicy, ManagedGroupEndReason, ManagedGroupEvent, ManagedGroupOutcome, ManagedGroupPacing, ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetFailure, ManagedTargetFailureKind, ManagedTargetOutcome, SessionEndReason, SessionOutcome, - SubscriberConfig, SubscriberOverflow, TargetEvent, TargetEventSubscription, TargetId, + SubscriberConfig, SubscriberOverflow, TargetEvent, TargetEventSubscription, MANAGED_GROUP_OUTCOME_HISTORY_LIMIT, }; +pub use managed::TargetId; pub use session::{NegotiatedParams, NegotiationRestriction}; pub use timing::ClientTimestamp; diff --git a/crates/irtt-client/src/managed/group.rs b/crates/irtt-client/src/managed/group.rs index 85e42ca..6003792 100644 --- a/crates/irtt-client/src/managed/group.rs +++ b/crates/irtt-client/src/managed/group.rs @@ -1,7 +1,6 @@ use std::{ collections::{HashMap, HashSet, VecDeque}, fmt, - hash::{Hash, Hasher}, net::{SocketAddr, UdpSocket}, sync::{ atomic::{AtomicU64, Ordering}, @@ -37,6 +36,7 @@ use crate::client::ProbeSendTimestamps; use super::{ cancellation::CancellationToken, hub::{EventHub, EventSubscription, SubscriberConfig}, + TargetId, }; const GROUP_RECV_TIMEOUT: Duration = Duration::from_millis(20); @@ -49,59 +49,6 @@ const RECV_BUFFER_SIZE: usize = 65_536; /// Maximum number of recent target outcomes retained in [`ManagedGroupOutcome`]. pub const MANAGED_GROUP_OUTCOME_HISTORY_LIMIT: usize = 256; -/// Caller-owned target identity for managed multi-target probing. -#[derive(Clone, Eq)] -pub struct TargetId(Arc); - -impl TargetId { - /// Borrow the identifier as a string slice. - pub fn as_str(&self) -> &str { - &self.0 - } -} - -impl fmt::Debug for TargetId { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_tuple("TargetId").field(&self.0).finish() - } -} - -impl fmt::Display for TargetId { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(&self.0) - } -} - -impl PartialEq for TargetId { - fn eq(&self, other: &Self) -> bool { - self.0 == other.0 - } -} - -impl Hash for TargetId { - fn hash(&self, state: &mut H) { - self.0.hash(state); - } -} - -impl From<&str> for TargetId { - fn from(value: &str) -> Self { - Self(Arc::from(value)) - } -} - -impl From for TargetId { - fn from(value: String) -> Self { - Self(Arc::from(value)) - } -} - -impl AsRef for TargetId { - fn as_ref(&self) -> &str { - self.as_str() - } -} - /// Per-target configuration for [`ManagedClientGroup`]. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ManagedTargetConfig { @@ -2283,7 +2230,7 @@ fn divide_duration(duration: Duration, divisor: usize) -> Duration { mod tests { use super::*; use crate::{ - config::NegotiationPolicy, managed::SubscriberOverflow, EventSubscriptionError, WarningKind, + config::NegotiationPolicy, EventSubscriptionError, SubscriberOverflow, WarningKind, }; use irtt_proto::{ compute_hmac_in_place, echo_packet_len, diff --git a/crates/irtt-client/src/managed/mod.rs b/crates/irtt-client/src/managed/mod.rs index db04c9e..69869f1 100644 --- a/crates/irtt-client/src/managed/mod.rs +++ b/crates/irtt-client/src/managed/mod.rs @@ -2,14 +2,34 @@ mod cancellation; mod group; mod hub; mod runner; +mod target_id; -pub use cancellation::CancellationToken; -pub use group::{ - ManagedClientGroup, ManagedClientGroupConfig, ManagedClientGroupSession, - ManagedGroupCompletionPolicy, ManagedGroupEndReason, ManagedGroupEvent, ManagedGroupOutcome, - ManagedGroupPacing, ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetFailure, - ManagedTargetFailureKind, ManagedTargetOutcome, TargetEvent, TargetEventSubscription, TargetId, - MANAGED_GROUP_OUTCOME_HISTORY_LIMIT, -}; -pub use hub::{EventHub, EventSubscription, SubscriberConfig, SubscriberOverflow}; -pub use runner::{ManagedClient, ManagedClientSession, SessionEndReason, SessionOutcome}; +pub use target_id::TargetId; + +/// Temporary exports for the obsolete thread-based managed implementations. +#[doc(hidden)] +pub mod legacy { + pub use super::cancellation::CancellationToken; + pub use super::group::{ + ManagedClientGroup, ManagedClientGroupConfig, ManagedClientGroupSession, + ManagedGroupCompletionPolicy, ManagedGroupEndReason, ManagedGroupEvent, + ManagedGroupOutcome, ManagedGroupPacing, ManagedTargetConfig, ManagedTargetEndReason, + ManagedTargetFailure, ManagedTargetFailureKind, ManagedTargetOutcome, TargetEvent, + TargetEventSubscription, MANAGED_GROUP_OUTCOME_HISTORY_LIMIT, + }; + pub use super::hub::{EventHub, EventSubscription, SubscriberConfig, SubscriberOverflow}; + pub use super::runner::{ + ManagedClient, ManagedClientSession, SessionEndReason, SessionOutcome, + }; + pub use super::TargetId; +} + +#[cfg(feature = "tokio")] +mod model; +#[cfg(feature = "tokio")] +mod task; + +#[cfg(feature = "tokio")] +pub use model::*; +#[cfg(feature = "tokio")] +pub use task::*; diff --git a/crates/irtt-client/src/managed/model.rs b/crates/irtt-client/src/managed/model.rs new file mode 100644 index 0000000..80bf437 --- /dev/null +++ b/crates/irtt-client/src/managed/model.rs @@ -0,0 +1,329 @@ +use std::{net::SocketAddr, sync::Arc, time::Duration}; + +use thiserror::Error; + +use crate::{ClientAuthConfig, ClientConfig, ClientError, ClientEvent}; + +use super::TargetId; + +/// Default capacity of the lossy managed event channel. +pub const DEFAULT_MANAGED_EVENT_CAPACITY: usize = 256; +/// Default number of recent target outcomes retained in status and output. +pub const DEFAULT_MANAGED_OUTCOME_HISTORY_LIMIT: usize = 256; +/// Default maximum number of simultaneously live target generations. +pub const DEFAULT_MANAGED_MAX_LIVE_TARGET_GENERATIONS: usize = 256; +/// Default time retained for late replies after the final committed timeout. +pub const DEFAULT_MANAGED_FINAL_DRAIN: Duration = Duration::from_millis(100); + +/// One globally allocated incarnation of a target within a managed task. +/// +/// Generations increase globally across the task, not independently per ID. +#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub struct TargetInstance { + pub id: TargetId, + pub generation: u64, +} + +/// Endpoint and optional authentication override for one managed target. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ManagedTargetConfig { + pub id: TargetId, + pub server_addr: String, + pub auth: Option, +} + +impl ManagedTargetConfig { + pub fn new(id: impl Into, server_addr: impl Into) -> Self { + Self { + id: id.into(), + server_addr: server_addr.into(), + auth: None, + } + } +} + +/// Coordination mode for sends across active targets. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum ManagedPacing { + #[default] + Staggered, + Burst, +} + +/// Policy controlling top-level natural completion. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum ManagedCompletionPolicy { + #[default] + FinishWhenQuiescent, + ExplicitStop, +} + +/// Shared configuration for one managed task. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ManagedClientConfig { + pub client: ClientConfig, + pub pacing: ManagedPacing, + pub completion: ManagedCompletionPolicy, + pub event_capacity: usize, + pub outcome_history_limit: usize, + pub max_live_target_generations: usize, + pub final_drain: Duration, +} + +impl Default for ManagedClientConfig { + fn default() -> Self { + Self { + client: ClientConfig::default(), + pacing: ManagedPacing::default(), + completion: ManagedCompletionPolicy::default(), + event_capacity: DEFAULT_MANAGED_EVENT_CAPACITY, + outcome_history_limit: DEFAULT_MANAGED_OUTCOME_HISTORY_LIMIT, + max_live_target_generations: DEFAULT_MANAGED_MAX_LIVE_TARGET_GENERATIONS, + final_drain: DEFAULT_MANAGED_FINAL_DRAIN, + } + } +} + +/// Durable top-level lifecycle. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ManagedLifecycle { + NotStarted, + Running, + Stopping, + Completed, + Failed, + Abandoned, +} + +/// Durable lifecycle for one target generation. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ManagedTargetLifecycle { + Pending, + Connecting, + Opening, + Active, + Draining, + Closing, + Terminal, +} + +/// Status for one target generation. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ManagedTargetStatus { + pub target: TargetInstance, + pub lifecycle: ManagedTargetLifecycle, + pub server_addr: Arc, + pub remote: Option, +} + +/// Authoritative immutable managed status snapshot. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ManagedStatus { + pub lifecycle: ManagedLifecycle, + pub stop_requested: bool, + pub desired_target_count: usize, + pub connecting_target_count: usize, + pub opening_target_count: usize, + pub active_target_count: usize, + pub draining_target_count: usize, + pub closing_target_count: usize, + pub terminal_target_count: usize, + pub total_target_outcomes: u64, + pub successful_target_outcomes: u64, + pub failed_target_outcomes: u64, + pub peer_closed_target_outcomes: u64, + pub discarded_target_outcomes: u64, + pub targets: Arc<[ManagedTargetStatus]>, + pub recent_target_outcomes: Arc<[ManagedTargetOutcome]>, + pub final_outcome: Option>, +} + +/// Lossy presentation event emitted by the managed task. +#[allow(clippy::large_enum_variant)] +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum ManagedEvent { + Started, + TargetStateChanged { + target: TargetInstance, + lifecycle: ManagedTargetLifecycle, + }, + Client { + target: TargetInstance, + event: ClientEvent, + }, + TargetFinished { + outcome: Arc, + }, + Stopping, + Completed { + outcome: Arc, + }, + Failed { + outcome: Arc, + }, + Abandoned, +} + +/// Receiving half of the lossy managed event stream. +pub type ManagedEventSubscription = tokio::sync::broadcast::Receiver; + +/// Final authoritative outcome returned by the task. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ManagedOutcome { + pub end_reason: ManagedEndReason, + pub total_target_outcomes: u64, + pub successful_target_outcomes: u64, + pub failed_target_outcomes: u64, + pub peer_closed_target_outcomes: u64, + pub discarded_target_outcomes: u64, + pub recent_target_outcomes: Arc<[ManagedTargetOutcome]>, +} + +/// Why the managed task ended. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum ManagedEndReason { + TargetsComplete, + StopRequested, + DriverFailed(ManagedDriverFailure), +} + +/// Durable outcome for one target generation. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ManagedTargetOutcome { + pub target: TargetInstance, + pub server_addr: Arc, + pub remote: Option, + pub end_reason: ManagedTargetEndReason, + pub packets_sent: u64, + pub replies_received: u64, + pub duplicates: u64, + pub late: u64, + pub warning_events: u64, + pub cleanup_failure: Option, +} + +/// Why one target generation ended. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum ManagedTargetEndReason { + TestComplete, + NoTestComplete, + PeerClosed, + Removed, + Replaced, + Stopped, + Failed(ManagedTargetFailure), +} + +/// Phase in which a target-local failure occurred. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ManagedTargetFailurePhase { + Connecting, + Opening, + Sending, + Receiving, + Timing, + Closing, +} + +/// Stable category for a target-local failure. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ManagedTargetFailureKind { + Resolve, + Socket, + SocketOption, + Protocol, + Timeout, + Configuration, + ResourceExhausted, + InvalidState, + Other, +} + +/// Durable target-local failure details. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ManagedTargetFailure { + pub phase: ManagedTargetFailurePhase, + pub kind: ManagedTargetFailureKind, + pub message: Arc, +} + +/// Configuration failure detected before any runtime work begins. +#[derive(Debug, Error)] +pub enum ManagedConfigError { + #[error("FinishWhenQuiescent requires at least one initial target")] + EmptyInitialTargets, + #[error("duplicate managed target id {id}")] + DuplicateTargetId { id: TargetId }, + #[error("configured {configured} targets exceeds live-generation limit {limit}")] + TooManyTargets { configured: usize, limit: usize }, + #[error("managed event capacity must be greater than zero")] + ZeroEventCapacity, + #[error("managed live-generation limit must be greater than zero")] + ZeroLiveGenerationLimit, + #[error("invalid managed target {id}: {source}")] + InvalidTarget { + id: TargetId, + #[source] + source: ClientError, + }, + #[error("managed target generation space is exhausted")] + GenerationExhausted, +} + +/// Failure to create a new event subscription. +#[derive(Clone, Debug, Eq, Error, PartialEq)] +pub enum ManagedSubscribeError { + #[error("managed event stream is closed")] + Closed, +} + +/// Whole-driver failure. +#[derive(Clone, Debug, Eq, Error, PartialEq)] +pub enum ManagedDriverFailure { + #[error("ManagedClientTask requires a current Tokio runtime")] + NoTokioRuntime, + #[error("managed driver exhausted resources for {operation}")] + ResourceExhausted { operation: &'static str }, + #[error("managed driver invariant failed: {message}")] + Internal { message: Arc }, +} + +pub(crate) fn classify_client_error( + phase: ManagedTargetFailurePhase, + error: &ClientError, +) -> ManagedTargetFailure { + let kind = match error { + ClientError::Resolve { .. } => ManagedTargetFailureKind::Resolve, + ClientError::Socket(_) => ManagedTargetFailureKind::Socket, + ClientError::SocketOption { .. } | ClientError::ReadTimeoutRestore { .. } => { + ManagedTargetFailureKind::SocketOption + } + ClientError::Protocol(_) + | ClientError::ProtocolVersionMismatch { .. } + | ClientError::ZeroToken + | ClientError::UnexpectedNoTestReply + | ClientError::NonZeroNoTestToken { .. } + | ClientError::ServerRejected + | ClientError::NegotiationRejected { .. } => ManagedTargetFailureKind::Protocol, + ClientError::OpenTimeout => ManagedTargetFailureKind::Timeout, + ClientError::InvalidConfig { .. } + | ClientError::OpenTimeoutTooSmall { .. } + | ClientError::NoOpenTimeouts => ManagedTargetFailureKind::Configuration, + ClientError::AllocationFailed { .. } + | ClientError::CounterOverflow { .. } + | ClientError::DurationOverflow + | ClientError::PendingLimitExceeded { .. } => ManagedTargetFailureKind::ResourceExhausted, + ClientError::NotOpen + | ClientError::AlreadyOpen + | ClientError::AlreadyCompleted + | ClientError::AlreadyClosed + | ClientError::StalePreparedProbe { .. } + | ClientError::PendingSequenceCollision { .. } => ManagedTargetFailureKind::InvalidState, + _ => ManagedTargetFailureKind::Other, + }; + ManagedTargetFailure { + phase, + kind, + message: Arc::from(error.to_string()), + } +} diff --git a/crates/irtt-client/src/managed/target_id.rs b/crates/irtt-client/src/managed/target_id.rs new file mode 100644 index 0000000..9c73a4e --- /dev/null +++ b/crates/irtt-client/src/managed/target_id.rs @@ -0,0 +1,47 @@ +use std::{fmt, sync::Arc}; + +/// Caller-owned stable target identity. +#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub struct TargetId(Arc); + +impl TargetId { + /// Construct a target identifier. + pub fn new(value: impl Into>) -> Self { + Self(value.into()) + } + + /// Borrow the identifier as a string slice. + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl fmt::Debug for TargetId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("TargetId").field(&self.0).finish() + } +} + +impl fmt::Display for TargetId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl From<&str> for TargetId { + fn from(value: &str) -> Self { + Self::new(value) + } +} + +impl From for TargetId { + fn from(value: String) -> Self { + Self::new(value) + } +} + +impl AsRef for TargetId { + fn as_ref(&self) -> &str { + self.as_str() + } +} diff --git a/crates/irtt-client/src/managed/task.rs b/crates/irtt-client/src/managed/task.rs new file mode 100644 index 0000000..1e7e680 --- /dev/null +++ b/crates/irtt-client/src/managed/task.rs @@ -0,0 +1,1220 @@ +use std::{ + collections::{HashSet, VecDeque}, + fmt, + future::Future, + mem, + pin::Pin, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + task::{Context, Poll}, + time::{Duration, Instant}, +}; + +use tokio::{ + sync::{broadcast, watch}, + time::{self, Sleep}, +}; + +use crate::{ + async_client::{AsyncClient, AsyncOpenState}, + session::machine::SessionMachine, + socket::validate_open_timeouts, + socket_options::validate_ttl, + ClientConfig, ClientError, ClientEvent, OpenOutcome, +}; + +use super::{ + classify_client_error, ManagedClientConfig, ManagedCompletionPolicy, ManagedConfigError, + ManagedDriverFailure, ManagedEndReason, ManagedEvent, ManagedEventSubscription, + ManagedLifecycle, ManagedOutcome, ManagedPacing, ManagedStatus, ManagedSubscribeError, + ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetFailure, ManagedTargetFailureKind, + ManagedTargetFailurePhase, ManagedTargetLifecycle, ManagedTargetOutcome, ManagedTargetStatus, + TargetInstance, +}; + +const TARGET_WORK_BUDGET: usize = 128; + +/// Entry point for constructing a unified Tokio managed task. +#[derive(Debug, Default)] +pub struct ManagedClient; + +impl ManagedClient { + /// Validate configuration and construct a runtime-independent task/handle pair. + pub fn task( + config: ManagedClientConfig, + targets: Vec, + ) -> Result<(ManagedClientTask, ManagedClientHandle), ManagedConfigError> { + build_task(config, targets) + } +} + +#[derive(Debug)] +struct StopSignal { + requested: AtomicBool, + wake: watch::Sender<()>, + acknowledged: AtomicBool, + acknowledgement: watch::Sender<()>, +} + +impl StopSignal { + fn request(&self) { + if !self.requested.swap(true, Ordering::AcqRel) { + self.wake.send_replace(()); + } + } + + fn is_requested(&self) -> bool { + self.requested.load(Ordering::Acquire) + } + + fn acknowledge(&self) { + if !self.acknowledged.swap(true, Ordering::AcqRel) { + self.acknowledgement.send_replace(()); + } + } +} + +/// Cloneable control and observation capability for a managed task. +#[derive(Clone)] +pub struct ManagedClientHandle { + stop: Arc, + status: watch::Receiver>, + events: broadcast::WeakSender, +} + +impl fmt::Debug for ManagedClientHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ManagedClientHandle") + .field("status", &self.status()) + .finish_non_exhaustive() + } +} + +impl ManagedClientHandle { + /// Return the latest durable immutable status snapshot. + pub fn status(&self) -> Arc { + Arc::clone(&self.status.borrow()) + } + + /// Subscribe to future lossy presentation events. + pub fn subscribe(&self) -> Result { + self.events + .upgrade() + .map(|sender| sender.subscribe()) + .ok_or(ManagedSubscribeError::Closed) + } + + /// Request idempotent graceful stop and return a durable-status receipt. + pub fn stop(&self) -> ManagedStopReceipt { + self.stop.request(); + ManagedStopReceipt::new(Arc::clone(&self.stop)) + } +} + +/// Receipt resolving once stop is durably observed or the task is terminal. +#[must_use = "await the receipt to observe durable stop acknowledgement"] +pub struct ManagedStopReceipt { + future: Pin + Send>>, +} + +impl ManagedStopReceipt { + fn new(stop: Arc) -> Self { + let mut acknowledgement = stop.acknowledgement.subscribe(); + let future = Box::pin(async move { + while !stop.acknowledged.load(Ordering::Acquire) { + let _ = acknowledgement.changed().await; + } + }); + Self { future } + } +} + +impl Future for ManagedStopReceipt { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.future.as_mut().poll(cx) + } +} + +type ConnectFuture = + Pin> + Send + 'static>>; +type WakeFuture = Pin>> + Send + 'static>>; + +fn arm_wake(mut receiver: watch::Receiver<()>) -> WakeFuture { + Box::pin(async move { + receiver.changed().await.ok()?; + Some(receiver) + }) +} + +enum TargetState { + Pending { + client_config: ClientConfig, + }, + Connecting { + future: ConnectFuture, + }, + Opening { + client: AsyncClient, + open: Box, + }, + Active { + client: AsyncClient, + }, + Draining { + client: AsyncClient, + deadline: Instant, + primary_end: ManagedTargetEndReason, + }, + Closing { + client: AsyncClient, + deadline: Instant, + primary_end: ManagedTargetEndReason, + }, + Terminal, +} + +impl TargetState { + fn lifecycle(&self) -> ManagedTargetLifecycle { + match self { + Self::Pending { .. } => ManagedTargetLifecycle::Pending, + Self::Connecting { .. } => ManagedTargetLifecycle::Connecting, + Self::Opening { .. } => ManagedTargetLifecycle::Opening, + Self::Active { .. } => ManagedTargetLifecycle::Active, + Self::Draining { .. } => ManagedTargetLifecycle::Draining, + Self::Closing { .. } => ManagedTargetLifecycle::Closing, + Self::Terminal => ManagedTargetLifecycle::Terminal, + } + } +} + +#[derive(Default)] +struct TargetCounters { + packets_sent: u64, + replies_received: u64, + duplicates: u64, + late: u64, + warning_events: u64, +} + +struct TargetRuntime { + instance: TargetInstance, + server_addr: Arc, + remote: Option, + counters: TargetCounters, + latest_committed_timeout: Option, + send_waiting: bool, + state: TargetState, +} + +#[derive(Default)] +struct OutcomeHistory { + limit: usize, + recent: VecDeque, + 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::>() + .into_boxed_slice(), + ) + } + + fn outcome(&self, end_reason: ManagedEndReason) -> ManagedOutcome { + ManagedOutcome { + end_reason, + 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>, + events: Option>, + stop: Arc, +} + +#[derive(Clone, Copy, Eq, PartialEq)] +enum DriverState { + NotStarted, + Running, + Stopping, + Finished, +} + +/// The sole authoritative zero-or-many target driver. +#[must_use = "ManagedClientTask must be awaited or deliberately dropped"] +pub struct ManagedClientTask { + state: DriverState, + lifecycle: ManagedLifecycle, + config: ManagedClientConfig, + targets: Vec, + history: OutcomeHistory, + resources: Option, + wake: Option, + timer: Option>>, + cursor: usize, + scan_remaining: usize, + send_cursor: usize, + burst_remaining: usize, + stagger_remaining: usize, + send_gate: Option, + stop_observed: bool, + final_outcome: Option>, + _next_generation: 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 { + 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.state.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(), + lifecycle, + server_addr: Arc::clone(&target.server_addr), + remote: target.remote, + } + }) + .collect::>(); + Arc::new(ManagedStatus { + lifecycle: self.lifecycle, + stop_requested: self.stop_observed, + desired_target_count: self.targets.len(), + 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.final_outcome.clone(), + }) + } + + fn replace_status(&self) { + self.resources().status.send_replace(self.snapshot()); + } + + fn install_target_state(&mut self, index: usize, state: TargetState) { + let lifecycle = state.lifecycle(); + self.targets[index].state = state; + self.replace_status(); + self.publish_event(ManagedEvent::TargetStateChanged { + target: self.targets[index].instance.clone(), + lifecycle, + }); + } + + fn publish_client_events(&mut self, index: usize, events: Vec) { + for event in events { + match event { + ClientEvent::EchoReply { .. } => { + self.targets[index].counters.replies_received = self.targets[index] + .counters + .replies_received + .saturating_add(1); + } + ClientEvent::DuplicateReply { .. } => { + self.targets[index].counters.duplicates = + self.targets[index].counters.duplicates.saturating_add(1); + } + ClientEvent::LateReply { .. } => { + self.targets[index].counters.late = + self.targets[index].counters.late.saturating_add(1); + } + ClientEvent::Warning { .. } => { + self.targets[index].counters.warning_events = self.targets[index] + .counters + .warning_events + .saturating_add(1); + } + _ => {} + } + self.publish_event(ManagedEvent::Client { + target: self.targets[index].instance.clone(), + event, + }); + } + } + + fn finish_target( + &mut self, + index: usize, + end_reason: ManagedTargetEndReason, + cleanup_failure: Option, + ) { + self.install_target_state(index, TargetState::Terminal); + let target = &self.targets[index]; + let outcome = ManagedTargetOutcome { + target: target.instance.clone(), + server_addr: Arc::clone(&target.server_addr), + remote: target.remote, + end_reason, + packets_sent: target.counters.packets_sent, + replies_received: target.counters.replies_received, + duplicates: target.counters.duplicates, + late: target.counters.late, + warning_events: target.counters.warning_events, + cleanup_failure, + }; + self.history.record(outcome.clone()); + self.replace_status(); + self.publish_event(ManagedEvent::TargetFinished { + outcome: Arc::new(outcome), + }); + } + + fn fail_target(&mut self, index: usize, phase: ManagedTargetFailurePhase, error: ClientError) { + let failure = classify_client_error(phase, &error); + self.finish_target(index, ManagedTargetEndReason::Failed(failure), None); + } + + fn begin_running(&mut self) -> Result<(), ManagedDriverFailure> { + tokio::runtime::Handle::try_current().map_err(|_| ManagedDriverFailure::NoTokioRuntime)?; + self.state = DriverState::Running; + self.lifecycle = ManagedLifecycle::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.state = DriverState::Stopping; + self.lifecycle = ManagedLifecycle::Stopping; + self.replace_status(); + self.publish_event(ManagedEvent::Stopping); + self.resources().stop.acknowledge(); + self.scan_remaining = self.targets.len(); + self.burst_remaining = 0; + } + + fn start_connecting(&mut self, index: usize, config: ClientConfig) { + let future = Box::pin(AsyncClient::connect(config)); + self.install_target_state(index, TargetState::Connecting { future }); + } + + fn poll_target(&mut self, index: usize, cx: &mut Context<'_>, now: Instant) -> bool { + let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal); + match state { + TargetState::Pending { client_config } => { + if self.state == DriverState::Stopping { + self.finish_target(index, ManagedTargetEndReason::Stopped, None); + false + } else { + self.start_connecting(index, client_config); + true + } + } + TargetState::Connecting { mut future } => { + if self.state == DriverState::Stopping { + self.finish_target(index, ManagedTargetEndReason::Stopped, None); + return false; + } + match future.as_mut().poll(cx) { + Poll::Pending => { + self.targets[index].state = TargetState::Connecting { future }; + false + } + Poll::Ready(Ok(client)) => { + self.targets[index].remote = Some(client.remote_addr()); + self.install_target_state( + index, + TargetState::Opening { + client, + open: Box::new(AsyncOpenState::new()), + }, + ); + true + } + Poll::Ready(Err(error)) => { + self.fail_target(index, ManagedTargetFailurePhase::Connecting, error); + false + } + } + } + TargetState::Opening { + mut client, + mut open, + } => { + if self.state == DriverState::Stopping { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.finish_target(index, ManagedTargetEndReason::Stopped, None); + return false; + } + 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]); + true + } + Poll::Ready(Ok(OpenOutcome::NoTestCompleted { event, .. })) => { + self.publish_client_events(index, vec![event]); + self.finish_target(index, ManagedTargetEndReason::NoTestComplete, None); + false + } + Poll::Ready(Err(error)) => { + self.fail_target(index, ManagedTargetFailurePhase::Opening, error); + false + } + } + } + TargetState::Active { mut client } => { + if self.state == DriverState::Stopping { + client.discard_prepared_probe(); + self.targets[index].counters.packets_sent = client.packets_sent(); + return self.begin_drain(index, client, ManagedTargetEndReason::Stopped, now); + } + let received = match client.poll_recv(cx) { + Poll::Pending => false, + Poll::Ready(Ok(events)) => { + self.publish_client_events(index, events); + true + } + Poll::Ready(Err(error)) => { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.fail_target(index, ManagedTargetFailurePhase::Receiving, error); + return false; + } + }; + if client.is_peer_closed() { + self.targets[index].counters.packets_sent = client.packets_sent(); + 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( + index, + client, + ManagedTargetEndReason::TestComplete, + now, + ); + } + self.targets[index].state = TargetState::Active { client }; + received + } + TargetState::Draining { + mut client, + deadline, + primary_end, + } => { + let received = match client.poll_recv(cx) { + Poll::Pending => false, + Poll::Ready(Ok(events)) => { + self.publish_client_events(index, events); + true + } + Poll::Ready(Err(error)) => { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.fail_target(index, ManagedTargetFailurePhase::Receiving, error); + return false; + } + }; + if client.is_peer_closed() { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.finish_target(index, ManagedTargetEndReason::PeerClosed, None); + 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(); + self.fail_target(index, ManagedTargetFailurePhase::Timing, error); + return false; + } + } + } + if now >= deadline { + let close_deadline = now.checked_add(self.config.final_drain).unwrap_or(now); + self.install_target_state( + index, + TargetState::Closing { + client, + deadline: close_deadline, + primary_end, + }, + ); + true + } else { + self.targets[index].state = TargetState::Draining { + client, + deadline, + primary_end, + }; + received + } + } + TargetState::Closing { + mut client, + deadline, + primary_end, + } => match client.poll_close(cx) { + Poll::Pending => { + if now >= deadline { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.finish_target(index, primary_end, Some(close_timeout_failure())); + false + } else { + self.targets[index].state = TargetState::Closing { + client, + deadline, + primary_end, + }; + false + } + } + Poll::Ready(Ok(events)) => { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.publish_client_events(index, events); + self.finish_target(index, primary_end, None); + false + } + Poll::Ready(Err(error)) => { + self.targets[index].counters.packets_sent = client.packets_sent(); + let cleanup = classify_client_error(ManagedTargetFailurePhase::Closing, &error); + self.finish_target(index, primary_end, Some(cleanup)); + false + } + }, + TargetState::Terminal => { + self.targets[index].state = TargetState::Terminal; + false + } + } + } + + fn begin_drain( + &mut self, + index: usize, + client: AsyncClient, + primary_end: ManagedTargetEndReason, + 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 { + self.fail_target( + index, + ManagedTargetFailurePhase::Timing, + ClientError::DurationOverflow, + ); + return false; + }; + self.install_target_state( + index, + TargetState::Draining { + client, + deadline, + primary_end, + }, + ); + true + } + + fn active_count(&self) -> usize { + self.targets + .iter() + .filter(|target| matches!(target.state, TargetState::Active { .. })) + .count() + } + + 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 { + self.targets[index].state = state; + return SendResult::NotAttempted; + }; + if client + .next_send_deadline() + .is_none_or(|deadline| deadline > now) + { + self.targets[index].send_waiting = false; + self.targets[index].state = TargetState::Active { client }; + return SendResult::NotAttempted; + } + let before = client.packets_sent(); + let result = client.poll_send_probe(cx); + let after = client.packets_sent(); + let accepted = after > before; + 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; + self.targets[index].state = TargetState::Active { client }; + SendResult::Pending + } + Poll::Ready(Ok(events)) => { + self.targets[index].send_waiting = false; + 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); + SendResult::Failed { accepted } + } else { + SendResult::Ready { accepted } + } + } + Poll::Ready(Err(error)) => { + self.targets[index].send_waiting = false; + self.fail_target(index, ManagedTargetFailurePhase::Sending, error); + SendResult::Failed { accepted } + } + } + } + + fn poll_staggered_send(&mut self, cx: &mut Context<'_>, now: Instant) -> bool { + let active = self.active_count(); + if active == 0 || self.send_gate.is_some_and(|gate| gate > now) { + return false; + } + if self.stagger_remaining == 0 { + self.stagger_remaining = self.targets.len(); + } + let work = self.stagger_remaining.min(TARGET_WORK_BUDGET); + for _ in 0..work { + let index = self.send_cursor % self.targets.len(); + self.send_cursor = (self.send_cursor + 1) % self.targets.len(); + self.stagger_remaining -= 1; + if !matches!(self.targets[index].state, TargetState::Active { .. }) { + continue; + } + let result = self.poll_one_send(index, cx, now); + if result == SendResult::NotAttempted { + continue; + } + self.stagger_remaining = 0; + if result.accepted() { + self.send_gate = Instant::now() + .checked_add(stagger_spacing(self.config.client.interval, active)); + } + return matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. }); + } + self.stagger_remaining > 0 + } + + fn poll_burst_sends(&mut self, cx: &mut Context<'_>, now: Instant) -> bool { + if self.targets.is_empty() { + return false; + } + if self.burst_remaining == 0 { + self.burst_remaining = self.targets.len(); + } + let work = self.burst_remaining.min(TARGET_WORK_BUDGET); + let mut immediate = false; + for _ in 0..work { + 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); + immediate |= matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. }); + } + immediate || 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.state, TargetState::Terminal)) + } + + fn next_deadline(&self) -> Option { + let mut non_send_deadline = None; + let mut send_deadline = None; + for target in &self.targets { + match &target.state { + TargetState::Active { client } => { + non_send_deadline = non_send_deadline + .into_iter() + .chain(client.next_probe_timeout_deadline()) + .min(); + if !target.send_waiting { + send_deadline = send_deadline + .into_iter() + .chain(client.next_send_deadline()) + .min(); + } + } + TargetState::Draining { + client, deadline, .. + } => { + non_send_deadline = non_send_deadline + .into_iter() + .chain(client.next_probe_timeout_deadline()) + .chain(Some(*deadline)) + .min(); + } + TargetState::Closing { deadline, .. } => { + non_send_deadline = non_send_deadline.into_iter().chain(Some(*deadline)).min(); + } + _ => {} + } + } + if self.state == DriverState::Running + && self.config.pacing == ManagedPacing::Staggered + && self.active_count() > 0 + { + send_deadline = match (send_deadline, self.send_gate) { + (Some(target), Some(gate)) => Some(target.max(gate)), + (None, gate) => gate, + (target, None) => target, + }; + } + 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, end_reason: ManagedEndReason, failed: bool) -> Poll { + let outcome = Arc::new(self.history.outcome(end_reason)); + self.final_outcome = Some(Arc::clone(&outcome)); + self.lifecycle = if failed { + ManagedLifecycle::Failed + } else { + ManagedLifecycle::Completed + }; + self.replace_status(); + self.publish_event(if failed { + ManagedEvent::Failed { + outcome: Arc::clone(&outcome), + } + } else { + ManagedEvent::Completed { + outcome: Arc::clone(&outcome), + } + }); + self.resources().stop.acknowledge(); + if let Some(resources) = self.resources.as_mut() { + resources.events.take(); + } + self.state = DriverState::Finished; + self.wake = None; + self.timer = None; + self.resources.take(); + Poll::Ready((*outcome).clone()) + } + + fn fail_driver(&mut self, failure: ManagedDriverFailure) -> Poll { + self.seal(ManagedEndReason::DriverFailed(failure), true) + } +} + +#[derive(Clone, Copy, Eq, PartialEq)] +enum SendResult { + NotAttempted, + Pending, + Ready { accepted: bool }, + Failed { accepted: bool }, +} + +impl SendResult { + fn accepted(self) -> bool { + matches!( + self, + Self::Ready { accepted: true } | Self::Failed { accepted: true } + ) + } +} + +impl Future for ManagedClientTask { + type Output = ManagedOutcome; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.as_mut().get_mut(); + if this.state == DriverState::Finished { + panic!("ManagedClientTask polled after completion"); + } + + if this.state == DriverState::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.state == DriverState::Running && this.resources().stop.is_requested() { + this.observe_stop(); + } + + let now = Instant::now(); + let mut immediate = this.poll_target_pass(cx, now); + if this.state == DriverState::Running { + immediate |= match this.config.pacing { + ManagedPacing::Staggered => this.poll_staggered_send(cx, now), + ManagedPacing::Burst => this.poll_burst_sends(cx, now), + }; + } + + if this.all_targets_terminal() { + match this.state { + DriverState::Stopping => { + return this.seal(ManagedEndReason::StopRequested, false); + } + DriverState::Running + if this.config.completion == ManagedCompletionPolicy::FinishWhenQuiescent => + { + this.state = DriverState::Stopping; + this.lifecycle = ManagedLifecycle::Stopping; + this.replace_status(); + this.publish_event(ManagedEvent::Stopping); + return this.seal(ManagedEndReason::TargetsComplete, false); + } + _ => {} + } + } + + immediate |= this.register_stop_wake(cx); + immediate |= this.register_timer(cx); + if immediate { + cx.waker().wake_by_ref(); + } + Poll::Pending + } +} + +impl Drop for ManagedClientTask { + fn drop(&mut self) { + if self.state == DriverState::Finished || self.resources.is_none() { + return; + } + self.lifecycle = ManagedLifecycle::Abandoned; + self.final_outcome = None; + self.replace_status(); + self.publish_event(ManagedEvent::Abandoned); + self.resources().stop.acknowledge(); + if let Some(resources) = self.resources.as_mut() { + resources.events.take(); + } + self.resources.take(); + } +} + +fn build_task( + config: ManagedClientConfig, + targets: Vec, +) -> Result<(ManagedClientTask, ManagedClientHandle), ManagedConfigError> { + if config.event_capacity == 0 { + return Err(ManagedConfigError::ZeroEventCapacity); + } + if config.max_live_target_generations == 0 { + return Err(ManagedConfigError::ZeroLiveGenerationLimit); + } + if targets.len() > config.max_live_target_generations { + return Err(ManagedConfigError::TooManyTargets { + configured: targets.len(), + limit: config.max_live_target_generations, + }); + } + if targets.is_empty() && config.completion == ManagedCompletionPolicy::FinishWhenQuiescent { + return Err(ManagedConfigError::EmptyInitialTargets); + } + + let mut ids = HashSet::with_capacity(targets.len()); + let mut runtimes = Vec::with_capacity(targets.len()); + let mut next_generation = 1_u64; + for target in targets { + if !ids.insert(target.id.clone()) { + return Err(ManagedConfigError::DuplicateTargetId { id: target.id }); + } + let generation = next_generation; + next_generation = next_generation + .checked_add(1) + .ok_or(ManagedConfigError::GenerationExhausted)?; + let mut client_config = config.client.clone(); + client_config.server_addr.clone_from(&target.server_addr); + if let Some(auth) = &target.auth { + client_config.hmac_key.clone_from(&auth.hmac_key); + } + validate_target_config(&client_config).map_err(|source| { + ManagedConfigError::InvalidTarget { + id: target.id.clone(), + source, + } + })?; + runtimes.push(TargetRuntime { + instance: TargetInstance { + id: target.id, + generation, + }, + server_addr: Arc::from(target.server_addr), + remote: None, + counters: TargetCounters::default(), + latest_committed_timeout: None, + send_waiting: false, + state: TargetState::Pending { client_config }, + }); + } + + let (event_sender, _) = broadcast::channel(config.event_capacity); + let weak_events = event_sender.downgrade(); + let (wake_sender, wake_receiver) = watch::channel(()); + let (acknowledgement_sender, _) = watch::channel(()); + let stop = Arc::new(StopSignal { + requested: AtomicBool::new(false), + wake: wake_sender, + acknowledged: AtomicBool::new(false), + acknowledgement: acknowledgement_sender, + }); + let history = OutcomeHistory::new(config.outcome_history_limit); + let initial_targets = runtimes + .iter() + .map(|target| ManagedTargetStatus { + target: target.instance.clone(), + lifecycle: ManagedTargetLifecycle::Pending, + server_addr: Arc::clone(&target.server_addr), + remote: None, + }) + .collect::>(); + let initial = Arc::new(ManagedStatus { + lifecycle: ManagedLifecycle::NotStarted, + stop_requested: false, + desired_target_count: runtimes.len(), + connecting_target_count: 0, + opening_target_count: 0, + active_target_count: 0, + draining_target_count: 0, + closing_target_count: 0, + terminal_target_count: 0, + total_target_outcomes: 0, + successful_target_outcomes: 0, + failed_target_outcomes: 0, + peer_closed_target_outcomes: 0, + discarded_target_outcomes: 0, + targets: Arc::from(initial_targets.into_boxed_slice()), + recent_target_outcomes: history.recent(), + final_outcome: None, + }); + let (status_sender, status_receiver) = watch::channel(initial); + let resources = TaskResources { + status: status_sender, + events: Some(event_sender), + stop: Arc::clone(&stop), + }; + let task = ManagedClientTask { + state: DriverState::NotStarted, + lifecycle: ManagedLifecycle::NotStarted, + config, + targets: runtimes, + history, + resources: Some(resources), + wake: Some(arm_wake(wake_receiver)), + timer: None, + cursor: 0, + scan_remaining: 0, + send_cursor: 0, + burst_remaining: 0, + stagger_remaining: 0, + send_gate: None, + stop_observed: false, + final_outcome: None, + _next_generation: next_generation, + }; + let handle = ManagedClientHandle { + stop, + status: status_receiver, + events: weak_events, + }; + Ok((task, handle)) +} + +fn validate_target_config(config: &ClientConfig) -> Result<(), ClientError> { + validate_open_timeouts(&config.open_timeouts)?; + SessionMachine::validate_config(config)?; + if config.socket_config.ipv4_only && config.socket_config.ipv6_only { + return Err(ClientError::InvalidConfig { + reason: "ipv4_only and ipv6_only cannot both be true".to_owned(), + }); + } + if let Some(ttl) = config.socket_config.ttl { + validate_ttl(ttl)?; + } + Ok(()) +} + +fn stagger_spacing(interval: Duration, active_targets: usize) -> Duration { + let divisor = u128::try_from(active_targets.max(1)).unwrap_or(u128::MAX); + let nanos = (interval.as_nanos() / divisor).max(1); + let seconds = u64::try_from(nanos / 1_000_000_000).unwrap_or(u64::MAX); + let subsec = u32::try_from(nanos % 1_000_000_000).unwrap_or(999_999_999); + Duration::new(seconds, subsec) +} + +fn close_timeout_failure() -> ManagedTargetFailure { + ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Closing, + kind: ManagedTargetFailureKind::Timeout, + message: Arc::from("best-effort close did not become writable before its deadline"), + } +} diff --git a/crates/irtt-client/src/probe.rs b/crates/irtt-client/src/probe.rs index c14916b..a1f4af9 100644 --- a/crates/irtt-client/src/probe.rs +++ b/crates/irtt-client/src/probe.rs @@ -71,6 +71,16 @@ impl PendingMap { self.map.len() } + #[cfg(feature = "tokio")] + pub fn next_timeout_deadline(&self) -> Option { + self.map.values().map(|probe| probe.timeout_at).min() + } + + #[cfg(feature = "tokio")] + pub fn latest_timeout_deadline(&self) -> Option { + self.map.values().map(|probe| probe.timeout_at).max() + } + #[cfg(test)] pub fn contains(&self, wire_seq: u32) -> bool { self.map.contains_key(&wire_seq) @@ -139,6 +149,11 @@ impl TimedOutMap { self.map.len() } + #[cfg(feature = "tokio")] + pub fn latest_timeout_deadline(&self) -> Option { + self.map.values().map(|probe| probe.timeout_at).max() + } + #[cfg(test)] pub fn contains(&self, wire_seq: u32) -> bool { self.map.contains_key(&wire_seq) diff --git a/crates/irtt-client/src/session/machine.rs b/crates/irtt-client/src/session/machine.rs index 5072ced..1bae69d 100644 --- a/crates/irtt-client/src/session/machine.rs +++ b/crates/irtt-client/src/session/machine.rs @@ -169,6 +169,18 @@ impl SessionMachine { config: ClientConfig, remote: std::net::SocketAddr, ) -> Result { + Self::validate_config(&config)?; + let requested = params_from_config(&config)?; + + Ok(Self { + config, + remote, + requested, + state: MachineState::Connected, + }) + } + + pub(crate) fn validate_config(config: &ClientConfig) -> Result<(), ClientError> { if config.max_pending_probes == 0 { return Err(ClientError::InvalidConfig { reason: "max_pending_probes must be greater than zero".to_owned(), @@ -179,14 +191,7 @@ impl SessionMachine { reason: "probe_timeout must be greater than zero".to_owned(), }); } - let requested = params_from_config(&config)?; - - Ok(Self { - config, - remote, - requested, - state: MachineState::Connected, - }) + params_from_config(config).map(|_| ()) } pub(crate) fn config(&self) -> &ClientConfig { @@ -490,6 +495,31 @@ impl SessionMachine { } } + #[cfg(feature = "tokio")] + pub(crate) fn next_probe_timeout_deadline(&self) -> Option { + match &self.state { + MachineState::Open(session) => session.pending.next_timeout_deadline(), + MachineState::Connected + | MachineState::NoTestCompleted + | MachineState::Closed { .. } => None, + } + } + + #[cfg(feature = "tokio")] + pub(crate) fn latest_probe_timeout_deadline(&self) -> Option { + match &self.state { + MachineState::Open(session) => session + .pending + .latest_timeout_deadline() + .into_iter() + .chain(session.timed_out.latest_timeout_deadline()) + .max(), + MachineState::Connected + | MachineState::NoTestCompleted + | MachineState::Closed { .. } => None, + } + } + pub(crate) fn is_open(&self) -> bool { matches!(self.state, MachineState::Open(_)) }