diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs index d3616d8..8523746 100644 --- a/crates/irtt-client/src/async_client.rs +++ b/crates/irtt-client/src/async_client.rs @@ -86,6 +86,8 @@ pub(crate) struct AsyncOpenState { deadline: Option, deadline_timer: Option>>, request_submitted: bool, + stop_after_current_attempt: bool, + stopped_after_current_attempt: bool, buffer: [u8; MAX_OPEN_PACKET_SIZE], cleanup: Option, } @@ -97,6 +99,8 @@ impl AsyncOpenState { deadline: None, deadline_timer: None, request_submitted: false, + stop_after_current_attempt: false, + stopped_after_current_attempt: false, buffer: [0_u8; MAX_OPEN_PACKET_SIZE], cleanup: None, } @@ -117,6 +121,24 @@ impl AsyncOpenState { self.deadline = None; self.deadline_timer = None; self.request_submitted = false; + self.stopped_after_current_attempt = self.stop_after_current_attempt; + } + + pub(crate) fn has_in_flight_work(&self) -> bool { + self.request_submitted || self.cleanup.is_some() + } + + pub(crate) fn request_stop_after_current_attempt(&mut self) { + self.stop_after_current_attempt = true; + } + + pub(crate) fn stopped_after_current_attempt(&self) -> bool { + self.stopped_after_current_attempt + } + + #[cfg(test)] + pub(crate) fn has_retained_cleanup(&self) -> bool { + self.cleanup.is_some() } fn deadline(&self) -> Instant { @@ -149,6 +171,7 @@ struct AsyncClientTestHooks { fail_cleanup_send: Cell, fail_peer_close_dscp: Cell, fail_dscp_restore: Cell, + pause_open_cleanup_before_send: Cell, pause_open_before_writable: Cell, pause_open_before_readable: Cell, pause_probe_before_writable: Cell, @@ -301,6 +324,10 @@ impl AsyncClient { if state.attempt >= attempt_count { return Poll::Ready(Err(ClientError::OpenTimeout)); } + if state.stop_after_current_attempt && state.deadline.is_none() { + state.stopped_after_current_attempt = true; + return Poll::Ready(Err(ClientError::OpenTimeout)); + } if state.deadline.is_none() { let timeout = self.machine.config().open_timeouts[state.attempt]; if let Err(error) = state.start_attempt(timeout) { @@ -333,10 +360,10 @@ impl AsyncClient { let send_result = self.socket.try_send(&request.bytes); match send_result { Ok(bytes) => { + state.request_submitted = true; if let Err(error) = validate_datagram_length(request.bytes.len(), bytes) { return Poll::Ready(Err(error)); } - state.request_submitted = true; } Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, Err(error) => return Poll::Ready(Err(ClientError::Socket(error))), @@ -352,6 +379,9 @@ impl AsyncClient { Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), Poll::Ready(Ok(false)) => { state.finish_attempt(); + if state.stopped_after_current_attempt() { + return Poll::Ready(Err(ClientError::OpenTimeout)); + } continue; } Poll::Ready(Ok(true)) => {} @@ -364,6 +394,9 @@ impl AsyncClient { }; if datagram.received_at.mono > state.deadline() { state.finish_attempt(); + if state.stopped_after_current_attempt() { + return Poll::Ready(Err(ClientError::OpenTimeout)); + } continue; } @@ -663,6 +696,13 @@ impl AsyncClient { self.schedule.as_ref()?.next_send_deadline() } + pub(crate) fn probe_interval(&self) -> Option { + if !self.machine.is_open() { + return None; + } + self.schedule.as_ref().map(ProbeSchedule::interval) + } + /// Return the configured local probe timeout. pub fn probe_timeout(&self) -> Duration { self.machine.probe_timeout() @@ -682,6 +722,52 @@ 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; + } + + #[cfg(test)] + pub(crate) fn retain_open_cleanup_once(&self) { + self.test_hooks.fail_open_dscp.set(true); + self.test_hooks.pause_open_cleanup_before_send.set(true); + } + + #[cfg(test)] + pub(crate) fn poll_recv_ready_for_test(&self, cx: &mut Context<'_>) -> Poll> { + self.socket.poll_recv_ready(cx) + } + + pub(crate) fn skip_missed_probe_slots_at(&mut self, now: Instant) -> Result<(), ClientError> { + let Some(schedule) = self.schedule.as_mut() else { + return Ok(()); + }; + 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, @@ -779,6 +865,15 @@ impl AsyncClient { return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); }; + #[cfg(test)] + if self + .test_hooks + .pause_open_cleanup_before_send + .replace(false) + { + return Poll::Pending; + } + #[cfg(test)] if self.test_hooks.fail_cleanup_send.replace(false) { return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); diff --git a/crates/irtt-client/src/client/schedule.rs b/crates/irtt-client/src/client/schedule.rs index 19fb7b3..8779b98 100644 --- a/crates/irtt-client/src/client/schedule.rs +++ b/crates/irtt-client/src/client/schedule.rs @@ -50,6 +50,11 @@ impl ProbeSchedule { self.next_send_at } + #[cfg(feature = "tokio")] + pub(crate) fn interval(&self) -> Duration { + self.interval + } + pub(crate) fn permit_probe_at(&mut self, now: Instant) -> bool { if self.next_send_at.is_none() { return false; 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..230559e --- /dev/null +++ b/crates/irtt-client/src/managed/model.rs @@ -0,0 +1,331 @@ +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("managed final drain {duration:?} cannot be scheduled from the current instant")] + UnschedulableFinalDrain { duration: Duration }, + #[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..a0bf27d --- /dev/null +++ b/crates/irtt-client/src/managed/task.rs @@ -0,0 +1,1561 @@ +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; +const POST_DEADLINE_RECEIVE_BUDGET: usize = TARGET_WORK_BUDGET; + +/// 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>>; +#[cfg(test)] +type EventObservations = Arc)>>>; +#[cfg(test)] +type StaggerObservations = Arc>>; + +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, + drain_started_at: Instant, + deadline: Instant, + primary_end: ManagedTargetEndReason, + cleanup_failure: Option, + post_deadline_receives_remaining: usize, + }, + Closing { + client: AsyncClient, + deadline: Instant, + primary_end: ManagedTargetEndReason, + cleanup_failure: Option, + }, + 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, + send_waiting: bool, + active: 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, +} + +#[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 { + 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, + last_stagger_send: Option, + stop_observed: bool, + final_outcome: Option>, + _next_generation: u64, + #[cfg(test)] + event_observations: Option, + #[cfg(test)] + stagger_observations: Option, + #[cfg(test)] + drain_test_hook: DrainTestHook, +} + +#[cfg(test)] +#[derive(Default)] +struct DrainTestHook { + defer_work_until_deadline: bool, + fail_receive: bool, + fail_close: bool, +} + +impl ManagedClientTask { + fn resources(&self) -> &TaskResources { + self.resources + .as_ref() + .expect("managed task resources exist before terminal sealing") + } + + fn publish_event(&self, event: ManagedEvent) { + #[cfg(test)] + if let Some(observations) = &self.event_observations { + observations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push((event.clone(), Arc::clone(&self.resources().status.borrow()))); + } + 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(); + let active = matches!(state, TargetState::Active { .. }); + let was_active = self.targets[index].active; + self.targets[index].active = active; + self.targets[index].state = state; + 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 { + 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_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; + 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 { + if !open.has_in_flight_work() { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.finish_target(index, ManagedTargetEndReason::Stopped, None); + return false; + } + open.request_stop_after_current_attempt(); + } + match client.poll_open(&mut open, cx) { + Poll::Pending => { + self.targets[index].state = TargetState::Opening { client, open }; + false + } + Poll::Ready(Ok(OpenOutcome::Started { event, .. })) => { + self.publish_client_events(index, vec![event]); + if self.state == DriverState::Stopping { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.begin_drain(index, client, ManagedTargetEndReason::Stopped, now) + } else { + self.install_target_state(index, TargetState::Active { client }); + true + } + } + Poll::Ready(Ok(OpenOutcome::NoTestCompleted { event, .. })) => { + self.publish_client_events(index, vec![event]); + let end_reason = if self.state == DriverState::Stopping { + ManagedTargetEndReason::Stopped + } else { + ManagedTargetEndReason::NoTestComplete + }; + self.finish_target(index, end_reason, None); + false + } + Poll::Ready(Err(error)) => { + if self.state == DriverState::Stopping { + let cleanup_failure = + (!open.stopped_after_current_attempt()).then(|| { + classify_client_error( + ManagedTargetFailurePhase::Opening, + &error, + ) + }); + self.finish_target( + index, + ManagedTargetEndReason::Stopped, + cleanup_failure, + ); + } else { + self.fail_target(index, ManagedTargetFailurePhase::Opening, error); + } + false + } + } + } + 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); + } + 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) => { + return self.begin_open_session_failure( + index, + client, + ManagedTargetFailurePhase::Timing, + error, + now, + OpenSessionFailureCleanup::Close, + ); + } + } + } + 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)) => { + return self.begin_open_session_failure( + index, + client, + ManagedTargetFailurePhase::Receiving, + error, + now, + OpenSessionFailureCleanup::Close, + ); + } + }; + 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.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, + drain_started_at, + deadline, + primary_end, + mut cleanup_failure, + mut post_deadline_receives_remaining, + } => { + #[cfg(test)] + let defer_work = self.drain_test_hook.defer_work_until_deadline && now < deadline; + #[cfg(not(test))] + let defer_work = false; + if defer_work { + self.targets[index].state = TargetState::Draining { + client, + drain_started_at, + deadline, + primary_end, + cleanup_failure, + post_deadline_receives_remaining, + }; + return false; + } + + #[cfg(test)] + 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); + retained_state_changed = true; + } + Err(error) => { + self.targets[index].counters.packets_sent = client.packets_sent(); + cleanup_failure.get_or_insert_with(|| { + classify_client_error(ManagedTargetFailurePhase::Timing, &error) + }); + return self.begin_close( + index, + client, + primary_end, + cleanup_failure, + now, + ); + } + } + } + let received = match client.poll_recv(cx) { + Poll::Pending => false, + Poll::Ready(Ok(events)) => { + self.publish_client_events(index, events); + retained_state_changed = true; + true + } + Poll::Ready(Err(error)) => { + self.targets[index].counters.packets_sent = client.packets_sent(); + cleanup_failure.get_or_insert_with(|| { + classify_client_error(ManagedTargetFailurePhase::Receiving, &error) + }); + return self.begin_close(index, client, primary_end, cleanup_failure, now); + } + }; + if let Some(error) = injected_receive_failure { + self.targets[index].counters.packets_sent = client.packets_sent(); + cleanup_failure.get_or_insert_with(|| { + classify_client_error(ManagedTargetFailurePhase::Receiving, &error) + }); + return self.begin_close(index, client, primary_end, cleanup_failure, now); + } + if client.is_peer_closed() { + self.targets[index].counters.packets_sent = client.packets_sent(); + 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, + post_deadline_receives_remaining, + }; + true + } else { + self.begin_close(index, client, primary_end, cleanup_failure, now) + } + } else { + self.targets[index].state = TargetState::Draining { + client, + drain_started_at, + deadline, + primary_end, + cleanup_failure, + post_deadline_receives_remaining, + }; + received + } + } + TargetState::Closing { + mut client, + deadline, + primary_end, + cleanup_failure, + } => 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, + cleanup_failure.or_else(|| Some(close_timeout_failure())), + ); + false + } else { + self.targets[index].state = TargetState::Closing { + client, + deadline, + primary_end, + cleanup_failure, + }; + false + } + } + Poll::Ready(Ok(events)) => { + self.targets[index].counters.packets_sent = client.packets_sent(); + self.publish_client_events(index, events); + #[cfg(test)] + let injected_close_failure = self.take_drain_close_failure().map(|error| { + classify_client_error(ManagedTargetFailurePhase::Closing, &error) + }); + #[cfg(not(test))] + let injected_close_failure = None; + self.finish_target( + index, + primary_end, + cleanup_failure.or(injected_close_failure), + ); + 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, cleanup_failure.or(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 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, + drain_started_at, + ); + }; + self.install_target_state( + index, + TargetState::Draining { + client, + drain_started_at, + deadline, + primary_end, + cleanup_failure: None, + post_deadline_receives_remaining: POST_DEADLINE_RECEIVE_BUDGET, + }, + ); + true + } + + fn drain_deadline(&self, client: &AsyncClient, drain_started_at: Instant) -> Option { + 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, + client: AsyncClient, + primary_end: ManagedTargetEndReason, + mut cleanup_failure: Option, + now: Instant, + ) -> bool { + let deadline = now.checked_add(self.config.final_drain).unwrap_or_else(|| { + cleanup_failure.get_or_insert_with(duration_overflow_failure); + now + }); + self.install_target_state( + index, + TargetState::Closing { + client, + deadline, + primary_end, + cleanup_failure, + }, + ); + true + } + + #[cfg(test)] + fn take_drain_receive_failure(&mut self) -> Option { + if !mem::take(&mut self.drain_test_hook.fail_receive) { + return None; + } + Some(ClientError::Socket(std::io::Error::other( + "injected drain receive failure", + ))) + } + + #[cfg(test)] + fn take_drain_close_failure(&mut self) -> Option { + if !mem::take(&mut self.drain_test_hook.fail_close) { + return None; + } + Some(ClientError::Socket(std::io::Error::other( + "injected drain close failure", + ))) + } + + fn active_count(&self) -> usize { + self.targets.iter().filter(|target| target.active).count() + } + + fn active_stagger_spacing(&self) -> Option<(usize, Duration)> { + let mut active = 0; + let mut minimum: Option = None; + for target in &self.targets { + let TargetState::Active { client } = &target.state else { + continue; + }; + active += 1; + let interval = client + .probe_interval() + .expect("active managed targets have a committed probe schedule"); + minimum = minimum.into_iter().chain(Some(interval)).min(); + } + minimum.map(|interval| (active, stagger_spacing(interval, active))) + } + + 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 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; + 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 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; + 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 { + 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 } + } + } + Poll::Ready(Err(error)) => { + self.targets[index].send_waiting = false; + self.begin_open_session_failure( + index, + client, + ManagedTargetFailurePhase::Sending, + error, + now, + OpenSessionFailureCleanup::Drain, + ); + SendResult::Failed { accepted } + } + } + } + + fn poll_staggered_send(&mut self, cx: &mut Context<'_>, now: Instant) -> bool { + if self.active_count() == 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 stagger_spacing = self.active_stagger_spacing(); + let result = self.poll_one_send(index, cx, now, stagger_spacing); + if result == SendResult::NotAttempted { + continue; + } + self.stagger_remaining = 0; + 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, None); + 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, .. + } => { + #[cfg(test)] + let defer_work = self.drain_test_hook.defer_work_until_deadline; + #[cfg(not(test))] + let defer_work = false; + non_send_deadline = non_send_deadline + .into_iter() + .chain( + (!defer_work) + .then(|| client.next_probe_timeout_deadline()) + .flatten(), + ) + .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 Instant::now().checked_add(config.final_drain).is_none() { + return Err(ManagedConfigError::UnschedulableFinalDrain { + duration: config.final_drain, + }); + } + if targets.len() > config.max_live_target_generations { + return Err(ManagedConfigError::TooManyTargets { + configured: targets.len(), + 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(), + send_waiting: false, + active: 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, + last_stagger_send: None, + stop_observed: false, + final_outcome: None, + _next_generation: next_generation, + #[cfg(test)] + event_observations: None, + #[cfg(test)] + stagger_observations: None, + #[cfg(test)] + drain_test_hook: DrainTestHook::default(), + }; + 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 duration_overflow_failure() -> ManagedTargetFailure { + classify_client_error( + ManagedTargetFailurePhase::Timing, + &ClientError::DurationOverflow, + ) +} + +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"), + } +} + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; diff --git a/crates/irtt-client/src/managed/tests.rs b/crates/irtt-client/src/managed/tests.rs new file mode 100644 index 0000000..bbc8e35 --- /dev/null +++ b/crates/irtt-client/src/managed/tests.rs @@ -0,0 +1,1334 @@ +use std::{ + future::{poll_fn, Future}, + mem, + net::{SocketAddr, UdpSocket}, + pin::Pin, + sync::{Arc, Condvar, Mutex}, + task::{Context, Poll, Waker}, + thread::{self, JoinHandle}, + time::{Duration, Instant}, +}; + +use irtt_proto::{ + decode_close_request, decode_echo_request, decode_open_request, encode_echo_reply, + encode_open_reply, flags, EchoReply, OpenReply, ReceivedStats, StampAt, TimestampFields, +}; +use tokio::runtime::{Builder, Runtime}; + +use super::*; +use crate::{socket::resolution_call_counts, ClientAuthConfig, NegotiationPolicy, RunMode}; + +const TOKEN: u64 = 0x1234_5678_90ab_cdef; + +#[derive(Clone, Copy)] +enum ServerBehavior { + Echo, + NoTest, + PeerClose, + DelayedEcho(Duration), + DeferredBurst(usize), +} + +#[derive(Clone, Debug)] +enum PacketKind { + Open, + Probe, + Close, +} + +#[derive(Clone, Debug)] +struct PacketRecord { + kind: PacketKind, + at: Instant, +} + +struct TestServer { + addr: SocketAddr, + records: Arc>>, + probe_seen: Arc<(Mutex, Condvar)>, + reply_sent: Arc<(Mutex, Condvar)>, + thread: JoinHandle<()>, +} + +impl TestServer { + fn finish(self) -> Vec { + self.thread.join().unwrap(); + self.records.lock().unwrap().clone() + } +} + +#[derive(Default)] +struct PacketGate { + seen: Mutex, + seen_ready: Condvar, + released: Mutex, + release_ready: Condvar, +} + +impl PacketGate { + fn arrive_and_wait(&self) { + *self.seen.lock().unwrap() = true; + self.seen_ready.notify_all(); + let released = self.released.lock().unwrap(); + let (released, timeout) = self + .release_ready + .wait_timeout_while(released, Duration::from_secs(2), |released| !*released) + .unwrap(); + assert!(*released && !timeout.timed_out()); + } + fn release(&self) { + *self.released.lock().unwrap() = true; + self.release_ready.notify_all(); + } +} +fn start_server(behavior: ServerBehavior, key: Option>) -> TestServer { + start_server_negotiating(behavior, key, None) +} +fn start_server_negotiating( + behavior: ServerBehavior, + key: Option>, + interval: Option, +) -> TestServer { + start_server_with_gates(behavior, key, interval, None, None) +} + +fn start_server_with_gates( + behavior: ServerBehavior, + key: Option>, + interval: Option, + open_gate: Option>, + reply_gate: Option>, +) -> TestServer { + let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + socket + .set_read_timeout(Some(Duration::from_secs(3))) + .unwrap(); + let addr = socket.local_addr().unwrap(); + let records = Arc::new(Mutex::new(Vec::new())); + let thread_records = Arc::clone(&records); + let probe_seen = Arc::new((Mutex::new(false), Condvar::new())); + let thread_probe_seen = Arc::clone(&probe_seen); + let reply_sent = Arc::new((Mutex::new(false), Condvar::new())); + let thread_reply_sent = Arc::clone(&reply_sent); + let thread = thread::spawn(move || { + let mut buffer = [0_u8; 8192]; + let (open_len, peer) = socket.recv_from(&mut buffer).unwrap(); + thread_records.lock().unwrap().push(PacketRecord { + kind: PacketKind::Open, + at: Instant::now(), + }); + let request = decode_open_request(&buffer[..open_len], key.as_deref()).unwrap(); + if let Some(gate) = &open_gate { + gate.arrive_and_wait(); + } + let mut negotiated = request.params.clone(); + if let Some(interval) = interval { + negotiated.interval_ns = i64::try_from(interval.as_nanos()).unwrap(); + } + let no_test = matches!(behavior, ServerBehavior::NoTest); + let open_reply = encode_open_reply( + &OpenReply { + flags: flags::FLAG_OPEN + | flags::FLAG_REPLY + | if no_test { flags::FLAG_CLOSE } else { 0 }, + token: if no_test { 0 } else { TOKEN }, + params: negotiated.clone(), + }, + key.as_deref(), + ) + .unwrap(); + socket.send_to(&open_reply, peer).unwrap(); + if no_test { + return; + } + + let mut deferred_burst_sent = false; + loop { + let Ok((len, packet_peer)) = socket.recv_from(&mut buffer) else { + return; + }; + let packet = &buffer[..len]; + if let Ok(probe) = decode_echo_request(packet, &negotiated, key.as_deref()) { + thread_records.lock().unwrap().push(PacketRecord { + kind: PacketKind::Probe, + at: Instant::now(), + }); + let (seen, ready) = &*thread_probe_seen; + *seen.lock().unwrap() = true; + ready.notify_all(); + if let ServerBehavior::DelayedEcho(delay) = behavior { + thread::sleep(delay); + } + if let ServerBehavior::DeferredBurst(count) = behavior { + if deferred_burst_sent { + continue; + } + deferred_burst_sent = true; + thread::sleep(Duration::from_millis(30)); + for reply_index in 0..count { + let reply = encode_echo_reply( + &EchoReply { + flags: flags::FLAG_REPLY, + token: TOKEN, + sequence: if reply_index < 2 { + probe.sequence + } else { + u32::MAX + }, + recv_count: None, + recv_window: None, + timestamps: TimestampFields::default(), + payload: Vec::new(), + }, + &negotiated, + key.as_deref(), + ) + .unwrap(); + socket.send_to(&reply, packet_peer).unwrap(); + } + continue; + } + if let Some(gate) = &reply_gate { + gate.arrive_and_wait(); + } + let peer_close = matches!(behavior, ServerBehavior::PeerClose); + let reply = encode_echo_reply( + &EchoReply { + flags: flags::FLAG_REPLY | if peer_close { flags::FLAG_CLOSE } else { 0 }, + token: TOKEN, + sequence: probe.sequence, + recv_count: None, + recv_window: None, + timestamps: TimestampFields::default(), + payload: Vec::new(), + }, + &negotiated, + key.as_deref(), + ) + .unwrap(); + socket.send_to(&reply, packet_peer).unwrap(); + let (sent, ready) = &*thread_reply_sent; + *sent.lock().unwrap() = true; + ready.notify_all(); + if peer_close { + return; + } + } else if decode_close_request(packet, key.as_deref()).is_ok() { + thread_records.lock().unwrap().push(PacketRecord { + kind: PacketKind::Close, + at: Instant::now(), + }); + return; + } + } + }); + TestServer { + addr, + records, + probe_seen, + reply_sent, + thread, + } +} + +fn runtime() -> Runtime { + Builder::new_current_thread() + .enable_io() + .enable_time() + .build() + .unwrap() +} + +fn config(pacing: ManagedPacing) -> ManagedClientConfig { + ManagedClientConfig { + client: ClientConfig { + duration: Some(Duration::from_millis(130)), + interval: Duration::from_millis(60), + probe_timeout: Duration::from_millis(35), + open_timeouts: vec![Duration::from_millis(200)], + received_stats: ReceivedStats::None, + stamp_at: StampAt::None, + ..ClientConfig::default() + }, + pacing, + final_drain: Duration::from_millis(5), + ..ManagedClientConfig::default() + } +} + +fn target(id: &str, addr: SocketAddr) -> ManagedTargetConfig { + ManagedTargetConfig::new(id, addr.to_string()) +} + +fn poll_once(future: Pin<&mut F>) -> Poll { + Future::poll(future, &mut Context::from_waker(Waker::noop())) +} +fn drive_task_until( + runtime: &Runtime, + task: &mut Pin>, + condition: impl Fn(&ManagedClientTask) -> bool, +) { + runtime.block_on(poll_fn(|cx| { + assert!( + task.as_mut().poll(cx).is_pending(), + "task completed before condition" + ); + if condition(task.as_ref().get_ref()) { + Poll::Ready(()) + } else { + Poll::Pending + } + })); +} + +fn wait_flag(flag: &Arc<(Mutex, Condvar)>) { + let (flag, ready) = &**flag; + let value = flag.lock().unwrap(); + let (value, timeout) = ready + .wait_timeout_while(value, Duration::from_secs(2), |value| !*value) + .unwrap(); + assert!(*value && !timeout.timed_out()); +} + +fn client_events( + observations: &Mutex)>>, +) -> Vec { + observations + .lock() + .unwrap() + .iter() + .filter_map(|(event, _)| match event { + ManagedEvent::Client { event, .. } => Some(event.clone()), + _ => None, + }) + .collect() +} + +fn probes(records: &[PacketRecord]) -> Vec { + records + .iter() + .filter_map(|record| matches!(record.kind, PacketKind::Probe).then_some(record.at)) + .collect() +} + +fn has_close(records: &[PacketRecord]) -> bool { + records + .iter() + .any(|record| matches!(record.kind, PacketKind::Close)) +} + +fn echo_sends( + observations: &Mutex)>>, +) -> Vec<(TargetInstance, u32, Instant, Instant)> { + observations + .lock() + .unwrap() + .iter() + .filter_map(|(event, _)| match event { + ManagedEvent::Client { + target, + event: + ClientEvent::EchoSent { + seq, + scheduled_at, + sent_at, + .. + }, + } => Some((target.clone(), *seq, *scheduled_at, sent_at.mono)), + _ => None, + }) + .collect() +} + +fn assert_rotated_without_catchup(sends: &[(TargetInstance, u32, Instant, Instant)]) { + assert!(sends.len() >= 4); + assert_ne!(sends[0].0, sends[1].0); + for target in [&sends[0].0, &sends[1].0] { + let target_sends = sends + .iter() + .filter(|(instance, ..)| instance == target) + .collect::>(); + assert!(target_sends.len() >= 2); + assert!(target_sends[1].2 > target_sends[0].3); + } +} + +fn run_negotiated_stagger_case( + requested: Duration, + negotiated: &[Duration], +) -> (Vec<(usize, Duration)>, Vec) { + let servers = negotiated + .iter() + .map(|interval| start_server_negotiating(ServerBehavior::Echo, None, Some(*interval))) + .collect::>(); + let mut managed = config(ManagedPacing::Staggered); + managed.client.duration = Some(Duration::from_millis(190)); + managed.client.interval = requested; + managed.client.negotiation_policy = NegotiationPolicy::Loose; + let targets = servers + .iter() + .enumerate() + .map(|(index, server)| target(&format!("target-{index}"), server.addr)) + .collect(); + let (mut task, _) = ManagedClient::task(managed, targets).unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + task.stagger_observations = Some(Arc::clone(&observations)); + let outcome = runtime().block_on(task); + let packets_sent = (0..servers.len()) + .map(|index| { + outcome + .recent_target_outcomes + .iter() + .find(|target| target.target.id.as_ref() == format!("target-{index}")) + .unwrap() + .packets_sent + }) + .collect(); + for server in servers { + server.finish(); + } + let spacings = observations.lock().unwrap().clone(); + (spacings, packets_sent) +} + +fn run_deferred_drain_burst(packet_count: usize) -> ManagedOutcome { + let server = start_server(ServerBehavior::DeferredBurst(packet_count), None); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + managed.client.interval = Duration::from_millis(100); + managed.client.probe_timeout = Duration::from_millis(60); + managed.final_drain = Duration::from_millis(20); + let (mut task, handle) = + ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap(); + task.drain_test_hook.defer_work_until_deadline = true; + 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(); + assert!(*guard && !timeout.timed_out()); + drop(handle.stop()); + }); + let outcome = runtime().block_on(async { + tokio::time::timeout(Duration::from_secs(2), task) + .await + .unwrap() + }); + stopper.join().unwrap(); + server.finish(); + outcome +} + +#[test] +fn construction_is_runtime_and_io_free() { + let before = resolution_call_counts(); + let (task, handle) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![ManagedTargetConfig::new("dns", "no-such-host.invalid")], + ) + .unwrap(); + assert_eq!(resolution_call_counts(), before); + assert_eq!(handle.status().lifecycle, ManagedLifecycle::NotStarted); + drop(task); +} + +#[test] +fn quiescent_rejects_empty_initial_targets() { + assert!(matches!( + ManagedClient::task(config(ManagedPacing::Staggered), vec![]), + Err(ManagedConfigError::EmptyInitialTargets) + )); +} + +#[test] +fn explicit_empty_waits_for_stop() { + let mut config = config(ManagedPacing::Staggered); + config.completion = ManagedCompletionPolicy::ExplicitStop; + let (task, handle) = ManagedClient::task(config, vec![]).unwrap(); + let outcome = runtime().block_on(async move { + let mut task = Box::pin(task); + assert!( + tokio::time::timeout(Duration::from_millis(20), task.as_mut()) + .await + .is_err() + ); + assert_eq!(handle.status().lifecycle, ManagedLifecycle::Running); + let receipt = handle.stop(); + let outcome = task.await; + receipt.await; + outcome + }); + assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); +} + +#[test] +fn duplicate_and_limit_rejections_are_transactional() { + let mut limited = config(ManagedPacing::Staggered); + limited.max_live_target_generations = 1; + assert!(matches!( + ManagedClient::task( + limited, + vec![ + target("a", "127.0.0.1:1".parse().unwrap()), + target("b", "127.0.0.1:2".parse().unwrap()) + ] + ), + Err(ManagedConfigError::TooManyTargets { .. }) + )); + let duplicate = vec![ + target("same", "127.0.0.1:1".parse().unwrap()), + target("same", "127.0.0.1:2".parse().unwrap()), + ]; + assert!(matches!( + ManagedClient::task(config(ManagedPacing::Staggered), duplicate), + Err(ManagedConfigError::DuplicateTargetId { .. }) + )); +} + +#[test] +fn first_poll_without_runtime_fails_durably() { + let (task, handle) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("one", "127.0.0.1:9".parse().unwrap())], + ) + .unwrap(); + let mut task = Box::pin(task); + let Poll::Ready(outcome) = poll_once(task.as_mut()) else { + panic!("task did not fail on its first poll") + }; + assert_eq!( + outcome.end_reason, + ManagedEndReason::DriverFailed(ManagedDriverFailure::NoTokioRuntime) + ); + assert_eq!(handle.status().lifecycle, ManagedLifecycle::Failed); +} + +#[test] +fn pre_poll_stop() { + let before = resolution_call_counts(); + let (task, handle) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![ManagedTargetConfig::new("one", "no-such-host.invalid")], + ) + .unwrap(); + let mut receipt = Box::pin(handle.stop()); + let mut task = Box::pin(task); + let Poll::Ready(outcome) = poll_once(task.as_mut()) else { + panic!("pre-stopped task did not complete") + }; + assert!(poll_once(receipt.as_mut()).is_ready()); + assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); + assert_eq!(resolution_call_counts(), before); +} + +#[test] +fn finite_immediate_replies_use_retained_drain_deadline() { + let server = start_server(ServerBehavior::Echo, None); + 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!( + outcome.recent_target_outcomes[0].end_reason, + ManagedTargetEndReason::TestComplete + )); + assert!(!probes(&records).is_empty() && has_close(&records)); +} + +#[test] +fn no_test_target_completion() { + let server = start_server(ServerBehavior::NoTest, None); + let mut config = config(ManagedPacing::Staggered); + config.client.run_mode = RunMode::NoTest; + let (task, _) = ManagedClient::task(config, vec![target("no-test", server.addr)]).unwrap(); + let outcome = runtime().block_on(task); + let records = server.finish(); + assert!(matches!( + outcome.recent_target_outcomes[0].end_reason, + ManagedTargetEndReason::NoTestComplete + )); + assert!(probes(&records).is_empty() && !has_close(&records)); +} + +#[test] +fn authenticated_peer_close_outcome() { + let key = b"managed-peer-close".to_vec(); + let server = start_server(ServerBehavior::PeerClose, Some(key.clone())); + let mut configured = target("peer", server.addr); + configured.auth = Some(ClientAuthConfig { + hmac_key: Some(key), + }); + let (task, _) = + ManagedClient::task(config(ManagedPacing::Staggered), vec![configured]).unwrap(); + let outcome = runtime().block_on(task); + let records = server.finish(); + assert_eq!(outcome.peer_closed_target_outcomes, 1); + assert!(matches!( + outcome.recent_target_outcomes[0].end_reason, + ManagedTargetEndReason::PeerClosed + )); + assert!(!has_close(&records)); +} + +#[test] +fn peer_close_during_stop_drain_remains_authoritative() { + let key = b"managed-drain-peer-close".to_vec(); + let server = start_server(ServerBehavior::PeerClose, Some(key.clone())); + let mut configured = target("peer", server.addr); + configured.auth = Some(ClientAuthConfig { + hmac_key: Some(key), + }); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + managed.client.interval = Duration::from_secs(1); + let (task, handle) = ManagedClient::task(managed, vec![configured]).unwrap(); + let runtime = runtime(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| { + task.targets[0].counters.packets_sent == 1 + }); + drop(handle.stop()); + wait_flag(&server.reply_sent); + let outcome = runtime.block_on(task); + let records = server.finish(); + assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); + assert_eq!(outcome.peer_closed_target_outcomes, 1); + assert!(matches!( + outcome.recent_target_outcomes[0].end_reason, + ManagedTargetEndReason::PeerClosed + )); + assert!(!has_close(&records)); +} + +#[test] +fn stop_during_opening_finishes_in_flight_work() { + let open_gate = Arc::new(PacketGate::default()); + let server = start_server_with_gates( + ServerBehavior::Echo, + None, + None, + Some(Arc::clone(&open_gate)), + None, + ); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + managed.client.open_timeouts = vec![Duration::from_secs(1), Duration::from_secs(1)]; + let (task, handle) = + ManagedClient::task(managed, vec![target("success", server.addr)]).unwrap(); + let runtime = runtime(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| match &task.targets[0].state { + TargetState::Opening { open, .. } => open.has_in_flight_work(), + _ => false, + }); + drop(handle.stop()); + open_gate.release(); + let outcome = runtime.block_on(task); + let records = server.finish(); + assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); + assert!(matches!( + outcome.recent_target_outcomes[0].end_reason, + ManagedTargetEndReason::Stopped + )); + assert_eq!( + records + .iter() + .filter(|record| matches!(record.kind, PacketKind::Open)) + .count(), + 1 + ); + assert!(probes(&records).is_empty() && has_close(&records)); + let key = b"managed-retained-open-cleanup".to_vec(); + let cleanup_server = start_server(ServerBehavior::Echo, Some(key.clone())); + let mut configured = target("cleanup", cleanup_server.addr); + configured.auth = Some(ClientAuthConfig { + hmac_key: Some(key), + }); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + let (task, handle) = ManagedClient::task(managed, vec![configured]).unwrap(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| { + matches!(task.targets[0].state, TargetState::Opening { .. }) + }); + let TargetState::Opening { client, .. } = &task.as_ref().get_ref().targets[0].state else { + unreachable!() + }; + client.retain_open_cleanup_once(); + drive_task_until(&runtime, &mut task, |task| match &task.targets[0].state { + TargetState::Opening { open, .. } => open.has_retained_cleanup(), + _ => false, + }); + drop(handle.stop()); + let outcome = runtime.block_on(task); + let records = cleanup_server.finish(); + let target = &outcome.recent_target_outcomes[0]; + assert!(matches!(target.end_reason, ManagedTargetEndReason::Stopped)); + assert_eq!(outcome.failed_target_outcomes, 0); + assert!(matches!( + target.cleanup_failure, + Some(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Opening, + .. + }) + )); + assert!(probes(&records).is_empty() && has_close(&records)); +} + +#[test] +fn static_multi_target_completion() { + let first = start_server(ServerBehavior::Echo, None); + let second = start_server(ServerBehavior::Echo, None); + let (task, _) = ManagedClient::task( + config(ManagedPacing::Burst), + vec![target("first", first.addr), target("second", second.addr)], + ) + .unwrap(); + let outcome = runtime().block_on(task); + assert_eq!(outcome.total_target_outcomes, 2); + assert_eq!(outcome.successful_target_outcomes, 2); + assert!(!probes(&first.finish()).is_empty()); + assert!(!probes(&second.finish()).is_empty()); +} + +#[test] +fn sibling_failure_isolation() { + let good = start_server(ServerBehavior::Echo, None); + let unused = UdpSocket::bind("127.0.0.1:0").unwrap(); + let bad_addr = unused.local_addr().unwrap(); + drop(unused); + let (task, _) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("good", good.addr), target("bad", bad_addr)], + ) + .unwrap(); + let outcome = runtime().block_on(task); + assert_eq!(outcome.total_target_outcomes, 2); + assert_eq!(outcome.successful_target_outcomes, 1); + assert_eq!(outcome.failed_target_outcomes, 1); + assert!(!probes(&good.finish()).is_empty()); +} + +#[test] +fn burst_pacing_fairness() { + let first = start_server(ServerBehavior::Echo, None); + let second = start_server(ServerBehavior::Echo, None); + let (mut task, _) = ManagedClient::task( + config(ManagedPacing::Burst), + vec![target("first", first.addr), target("second", second.addr)], + ) + .unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + task.event_observations = Some(Arc::clone(&observations)); + runtime().block_on(task); + assert_rotated_without_catchup(&echo_sends(&observations)); + first.finish(); + second.finish(); +} + +#[test] +fn staggered_pacing_fairness() { + let first = start_server(ServerBehavior::Echo, None); + let second = start_server(ServerBehavior::Echo, None); + let (mut task, _) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("first", first.addr), target("second", second.addr)], + ) + .unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + task.event_observations = Some(Arc::clone(&observations)); + runtime().block_on(task); + assert_rotated_without_catchup(&echo_sends(&observations)); + first.finish(); + second.finish(); +} + +#[test] +fn staggered_pacing_uses_all_active_negotiated_intervals() { + let requested = Duration::from_millis(120); + + let (single, single_packets) = + run_negotiated_stagger_case(requested, &[Duration::from_millis(40)]); + assert!(single + .iter() + .all(|entry| *entry == (1, Duration::from_millis(40)))); + assert!(single_packets[0] >= 4); + + let (different, different_packets) = run_negotiated_stagger_case( + requested, + &[Duration::from_millis(30), Duration::from_millis(90)], + ); + let while_both_active = different + .iter() + .filter(|(active, _)| *active == 2) + .collect::>(); + assert!(!while_both_active.is_empty()); + assert!(while_both_active + .iter() + .all(|(_, spacing)| *spacing == Duration::from_millis(15))); + assert!(different_packets[0] > different_packets[1]); + + let (equal, _) = run_negotiated_stagger_case( + requested, + &[Duration::from_millis(60), Duration::from_millis(60)], + ); + let while_both_active = equal + .iter() + .filter(|(active, _)| *active == 2) + .collect::>(); + assert!(!while_both_active.is_empty()); + assert!(while_both_active + .iter() + .all(|(_, spacing)| *spacing == Duration::from_millis(30))); +} + +#[test] +fn stagger_gate_tracks_active_membership() { + let first = + start_server_negotiating(ServerBehavior::Echo, None, Some(Duration::from_millis(100))); + let second_gate = Arc::new(PacketGate::default()); + let second = start_server_with_gates( + ServerBehavior::Echo, + None, + Some(Duration::from_secs(1)), + Some(Arc::clone(&second_gate)), + None, + ); + let mut managed = config(ManagedPacing::Staggered); + managed.client.duration = Some(Duration::from_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)]; + let (task, handle) = ManagedClient::task( + managed, + vec![target("first", first.addr), target("second", second.addr)], + ) + .unwrap(); + let runtime = runtime(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| { + task.targets[0].counters.packets_sent == 1 + }); + let one_target_last = task.last_stagger_send.unwrap(); + let one_target_gate = task.send_gate.unwrap(); + assert_eq!( + task.send_gate, + 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(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); + first.finish(); + second.finish(); +} + +#[test] +fn overdue_queued_reply_is_lost_then_late() { + let reply_gate = Arc::new(PacketGate::default()); + let server = start_server_with_gates( + ServerBehavior::Echo, + None, + None, + None, + Some(Arc::clone(&reply_gate)), + ); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + managed.client.interval = Duration::from_secs(1); + managed.client.probe_timeout = Duration::from_millis(30); + let (mut task, handle) = + ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + task.event_observations = Some(Arc::clone(&observations)); + let runtime = runtime(); + let mut task = Box::pin(task); + drive_task_until(&runtime, &mut task, |task| { + task.targets[0].counters.packets_sent == 1 + }); + let timeout = match &task.targets[0].state { + TargetState::Active { client } => client.next_probe_timeout_deadline().unwrap(), + _ => panic!("probe sender did not remain active"), + }; + thread::sleep(timeout.saturating_duration_since(Instant::now())); + while Instant::now() <= timeout { + thread::yield_now(); + } + reply_gate.release(); + wait_flag(&server.reply_sent); + runtime.block_on(poll_fn(|cx| { + let ready = match &task.as_ref().get_ref().targets[0].state { + TargetState::Active { client } => client.poll_recv_ready_for_test(cx), + _ => panic!("probe sender did not remain active"), + }; + ready.map(|result| result.unwrap()) + })); + runtime.block_on(poll_fn(|cx| { + assert!(task.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + })); + let classified = client_events(&observations) + .into_iter() + .filter(|event| { + matches!( + event, + ClientEvent::EchoLoss { .. } + | ClientEvent::LateReply { .. } + | ClientEvent::EchoReply { .. } + ) + }) + .collect::>(); + assert!(matches!( + classified.as_slice(), + [ClientEvent::EchoLoss { .. }, ClientEvent::LateReply { .. }] + )); + drop(handle.stop()); + let outcome = runtime.block_on(task); + let target = &outcome.recent_target_outcomes[0]; + assert_eq!(target.replies_received, 0); + assert_eq!(target.late, 1); + server.finish(); +} + +#[test] +fn 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 + }); + 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] +fn drain_failures_preserve_primary_outcome_and_first_cleanup_failure() { + let completed_server = start_server(ServerBehavior::Echo, None); + let (mut completed_task, _) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("completed", completed_server.addr)], + ) + .unwrap(); + completed_task.drain_test_hook.fail_receive = true; + completed_task.drain_test_hook.fail_close = true; + let completed = runtime().block_on(completed_task); + let completed_target = &completed.recent_target_outcomes[0]; + assert!(matches!( + completed_target.end_reason, + ManagedTargetEndReason::TestComplete + )); + assert_eq!(completed.failed_target_outcomes, 0); + assert!(matches!( + completed_target.cleanup_failure, + Some(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Receiving, + kind: ManagedTargetFailureKind::Socket, + .. + }) + )); + completed_server.finish(); + + let stopped_server = start_server(ServerBehavior::Echo, None); + let mut managed = config(ManagedPacing::Staggered); + managed.completion = ManagedCompletionPolicy::ExplicitStop; + managed.client.duration = None; + let (mut stopped_task, handle) = + ManagedClient::task(managed, vec![target("stopped", stopped_server.addr)]).unwrap(); + stopped_task.drain_test_hook.fail_receive = true; + let seen = Arc::clone(&stopped_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(); + assert!(*guard && !timeout.timed_out()); + drop(handle.stop()); + }); + let stopped = runtime().block_on(stopped_task); + stopper.join().unwrap(); + let stopped_target = &stopped.recent_target_outcomes[0]; + assert!(matches!( + stopped_target.end_reason, + ManagedTargetEndReason::Stopped + )); + assert_eq!(stopped.failed_target_outcomes, 0); + assert!(matches!( + stopped_target.cleanup_failure, + Some(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Receiving, + kind: ManagedTargetFailureKind::Socket, + .. + }) + )); + stopped_server.finish(); +} + +#[test] +fn final_drain_overflow_is_rejected_and_preserves_success() { + let mut invalid = config(ManagedPacing::Staggered); + invalid.final_drain = Duration::MAX; + assert!(matches!( + ManagedClient::task( + invalid, + vec![target("invalid", "127.0.0.1:9".parse().unwrap())] + ), + Err(ManagedConfigError::UnschedulableFinalDrain { + duration: Duration::MAX + }) + )); + + let server = start_server(ServerBehavior::Echo, None); + let (mut task, _) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("completed", server.addr)], + ) + .unwrap(); + task.config.final_drain = Duration::MAX; + let outcome = runtime().block_on(task); + let target = &outcome.recent_target_outcomes[0]; + assert!(matches!( + target.end_reason, + ManagedTargetEndReason::TestComplete + )); + assert_eq!(outcome.successful_target_outcomes, 1); + assert_eq!(outcome.failed_target_outcomes, 0); + assert!(matches!( + target.cleanup_failure, + Some(ManagedTargetFailure { + phase: ManagedTargetFailurePhase::Timing, + kind: ManagedTargetFailureKind::ResourceExhausted, + .. + }) + )); + assert!(has_close(&server.finish())); +} + +#[test] +fn post_deadline_receive_sweep_drains_queued_packets_and_is_bounded() { + let drained = run_deferred_drain_burst(3); + let drained_target = &drained.recent_target_outcomes[0]; + assert!(matches!( + drained_target.end_reason, + ManagedTargetEndReason::Stopped + )); + assert_eq!(drained_target.replies_received, 0); + assert_eq!(drained_target.duplicates, 1); + assert_eq!(drained_target.late, 2); + + let bounded = run_deferred_drain_burst(POST_DEADLINE_RECEIVE_BUDGET + 64); + let bounded_target = &bounded.recent_target_outcomes[0]; + assert!(matches!( + bounded_target.end_reason, + ManagedTargetEndReason::Stopped + )); + assert_eq!( + bounded_target.replies_received + bounded_target.duplicates + bounded_target.late, + u64::try_from(POST_DEADLINE_RECEIVE_BUDGET).unwrap() + ); + assert_eq!(bounded.failed_target_outcomes, 0); +} + +#[test] +fn event_loss_independence() { + let server = start_server(ServerBehavior::Echo, None); + let mut config = config(ManagedPacing::Staggered); + config.event_capacity = 1; + let (task, handle) = ManagedClient::task(config, vec![target("one", server.addr)]).unwrap(); + let mut events = handle.subscribe().unwrap(); + let outcome = runtime().block_on(task); + assert!(matches!( + events.try_recv(), + Err(broadcast::error::TryRecvError::Lagged(_)) + )); + assert_eq!(outcome.successful_target_outcomes, 1); + assert_eq!(handle.status().lifecycle, ManagedLifecycle::Completed); + server.finish(); +} + +#[test] +fn outcome_history_eviction_preserves_aggregates() { + let servers = [ + start_server(ServerBehavior::NoTest, None), + start_server(ServerBehavior::NoTest, None), + start_server(ServerBehavior::NoTest, None), + ]; + let mut config = config(ManagedPacing::Staggered); + config.client.run_mode = RunMode::NoTest; + config.outcome_history_limit = 1; + let targets = servers + .iter() + .enumerate() + .map(|(index, server)| target(&format!("target-{index}"), server.addr)) + .collect(); + let (mut task, _) = ManagedClient::task(config, targets).unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + task.event_observations = Some(Arc::clone(&observations)); + let outcome = runtime().block_on(task); + assert_eq!(outcome.total_target_outcomes, 3); + assert_eq!(outcome.successful_target_outcomes, 3); + assert_eq!(outcome.discarded_target_outcomes, 2); + assert_eq!(outcome.recent_target_outcomes.len(), 1); + let last_finished = observations + .lock() + .unwrap() + .iter() + .filter_map(|(event, _)| match event { + ManagedEvent::TargetFinished { outcome } => Some(outcome.target.clone()), + _ => None, + }) + .next_back() + .unwrap(); + assert_eq!(outcome.recent_target_outcomes[0].target, last_finished); + for server in servers { + server.finish(); + } +} + +#[test] +fn status_precedes_corresponding_events() { + let server = start_server(ServerBehavior::NoTest, None); + let mut config = config(ManagedPacing::Staggered); + config.client.run_mode = RunMode::NoTest; + let (mut task, _) = ManagedClient::task(config, vec![target("one", server.addr)]).unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + task.event_observations = Some(Arc::clone(&observations)); + runtime().block_on(task); + let mut saw_started = false; + let mut saw_state = false; + let mut saw_finished = false; + let mut saw_stopping = false; + let mut saw_completed = false; + for (event, status) in observations.lock().unwrap().iter() { + match event { + ManagedEvent::Started => { + saw_started = true; + assert_eq!(status.lifecycle, ManagedLifecycle::Running); + } + ManagedEvent::TargetStateChanged { target, lifecycle } => { + saw_state = true; + assert!(status + .targets + .iter() + .any(|entry| entry.target == *target && entry.lifecycle == *lifecycle)); + } + ManagedEvent::TargetFinished { .. } => { + saw_finished = true; + assert!(status.total_target_outcomes >= 1); + } + ManagedEvent::Stopping => { + saw_stopping = true; + assert_eq!(status.lifecycle, ManagedLifecycle::Stopping); + } + ManagedEvent::Completed { .. } => { + saw_completed = true; + assert_eq!(status.lifecycle, ManagedLifecycle::Completed); + } + _ => {} + } + } + assert!(saw_started && saw_state && saw_finished && saw_stopping && saw_completed); + server.finish(); +} + +#[test] +fn dropping_all_handles_does_not_stop() { + let server = start_server(ServerBehavior::Echo, None); + let (task, handle) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("one", server.addr)], + ) + .unwrap(); + drop(handle); + let outcome = runtime().block_on(task); + assert_eq!(outcome.end_reason, ManagedEndReason::TargetsComplete); + server.finish(); +} + +#[test] +fn task_drop_abandonment() { + let (task, handle) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("one", "127.0.0.1:9".parse().unwrap())], + ) + .unwrap(); + let mut events = handle.subscribe().unwrap(); + drop(task); + assert_eq!(handle.status().lifecycle, ManagedLifecycle::Abandoned); + assert_eq!(handle.status().total_target_outcomes, 0); + assert!(matches!(events.try_recv(), Ok(ManagedEvent::Abandoned))); +} + +#[test] +fn terminal_subscription_is_closed() { + let (task, handle) = ManagedClient::task( + config(ManagedPacing::Staggered), + vec![target("one", "127.0.0.1:9".parse().unwrap())], + ) + .unwrap(); + let mut task = Box::pin(task); + assert!(poll_once(task.as_mut()).is_ready()); + assert!(matches!( + handle.subscribe(), + Err(ManagedSubscribeError::Closed) + )); +} 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(_)) }