From e7de8bdb309277669488d602218080f35ed30377 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Mon, 5 Oct 2026 00:24:15 +0700 Subject: [PATCH] refactor(client): split managed task into coherent internal modules Group control and construction, target lifecycle, target-set transactions, send pacing and commitment, and bounded timeout processing into private children of the existing task owner. Keep durable status, outcome history, top-level polling, sealing, and abandonment in task.rs. Preserve phase-owned resources, stop/update synchronization, publication ordering, pacing, and timeout fairness. Move the existing private gate and timeout tests with their implementation without adding coverage or hooks. --- crates/irtt-client/src/managed/task.rs | 2104 +---------------- .../irtt-client/src/managed/task/control.rs | 511 ++++ crates/irtt-client/src/managed/task/pacing.rs | 365 +++ crates/irtt-client/src/managed/task/target.rs | 780 ++++++ .../irtt-client/src/managed/task/timeouts.rs | 226 ++ crates/irtt-client/src/managed/task/update.rs | 308 +++ 6 files changed, 2220 insertions(+), 2074 deletions(-) create mode 100644 crates/irtt-client/src/managed/task/control.rs create mode 100644 crates/irtt-client/src/managed/task/pacing.rs create mode 100644 crates/irtt-client/src/managed/task/target.rs create mode 100644 crates/irtt-client/src/managed/task/timeouts.rs create mode 100644 crates/irtt-client/src/managed/task/update.rs diff --git a/crates/irtt-client/src/managed/task.rs b/crates/irtt-client/src/managed/task.rs index 2ad601b..ca2356d 100644 --- a/crates/irtt-client/src/managed/task.rs +++ b/crates/irtt-client/src/managed/task.rs @@ -1,356 +1,42 @@ +//! Task ownership, durable status/history, and top-level polling and sealing. +//! +//! Private children implement control, target lifecycle, updates, pacing, and +//! bounded timeout work on this same task; they do not introduce other owners. + +mod control; +mod pacing; +mod target; +mod timeouts; +mod update; + +pub use control::{ManagedClient, ManagedClientHandle, ManagedCommandReceipt, ManagedStopReceipt}; + use std::{ - collections::{HashSet, VecDeque}, - fmt, + collections::VecDeque, future::Future, pin::Pin, - sync::{ - atomic::{AtomicBool, AtomicU8, Ordering}, - Arc, - }, + sync::Arc, task::{Context, Poll}, - time::{Duration, Instant}, + time::Instant, }; use tokio::{ - sync::{broadcast, mpsc, oneshot, watch}, + sync::{broadcast, mpsc, watch}, time::{self, Sleep}, }; -use crate::{ - async_client::{AsyncClient, AsyncOpenState}, - session::machine::SessionMachine, - socket_options::validate_ttl, - ClientConfig, ClientError, ClientEvent, OpenOutcome, SendProbeError, -}; - -use super::schedule::{instant_abs_diff, ProbeSchedule}; - use super::{ - classify_client_error, ManagedClientConfig, ManagedCommandAcknowledgement, - ManagedCommandApplyError, ManagedCommandError, ManagedCompletionPolicy, ManagedConfigError, - ManagedDriverFailure, ManagedEndReason, ManagedEvent, ManagedEventSubscription, - ManagedLifecycle, ManagedOutcome, ManagedPacing, ManagedStatus, ManagedStatusSubscription, - ManagedSubscribeError, ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetFailure, - ManagedTargetFailureKind, ManagedTargetFailurePhase, ManagedTargetLifecycle, - ManagedTargetOutcome, ManagedTargetStatus, TargetInstance, + ManagedClientConfig, ManagedCommandApplyError, ManagedCompletionPolicy, ManagedDriverFailure, + ManagedEndReason, ManagedEvent, ManagedLifecycle, ManagedOutcome, ManagedPacing, ManagedStatus, + ManagedTargetEndReason, ManagedTargetLifecycle, ManagedTargetOutcome, ManagedTargetStatus, }; +use control::{ManagedCommand, StopSignal}; +use pacing::gated_send_deadline; +use target::{TargetPhase, TargetRuntime}; +use timeouts::TimeoutBacklogEntry; const TARGET_WORK_BUDGET: usize = 128; -const COMMAND_WORK_BUDGET: usize = 32; -const TIMEOUT_WORK_BUDGET: usize = 128; -const POST_DEADLINE_RECEIVE_BUDGET: usize = TARGET_WORK_BUDGET; -// Tokio 1.53 broadcast asserts this exact bound before rounding capacity to a -// power of two; Tokio does not expose it as a public constant. -const MAX_BROADCAST_CHANNEL_CAPACITY: usize = usize::MAX >> 1; -/// Entry point for constructing a unified Tokio managed task. -/// -/// # Example -/// -/// Construct the task and its separate control handle, subscribe before -/// starting the task, then run it under the caller's Tokio runtime. A reachable -/// server is required at runtime, so this example only type-checks. -/// -/// ```no_run -/// use irtt_client::{ClientConfig, ClientEvent}; -/// use irtt_client::managed::{ -/// ManagedClient, ManagedClientConfig, ManagedEvent, ManagedTargetConfig, -/// }; -/// -/// # async fn run() -> Result<(), Box> { -/// let config = ManagedClientConfig { -/// client: ClientConfig { -/// request: irtt_client::SessionRequest { -/// duration: Some(std::time::Duration::from_secs(10)), -/// ..irtt_client::SessionRequest::default() -/// }, -/// ..ClientConfig::default() -/// }, -/// ..ManagedClientConfig::default() -/// }; -/// let targets = vec![ManagedTargetConfig::new("edge", "127.0.0.1:2112")]; -/// let (task, handle) = ManagedClient::task(config, targets)?; -/// let mut events = handle.subscribe()?; -/// let driver = tokio::spawn(task); -/// -/// while let Ok(event) = events.recv().await { -/// if let ManagedEvent::Client { event: ClientEvent::EchoReply { .. }, .. } = event { -/// println!("received a probe reply"); -/// } -/// } -/// let outcome = driver.await?; -/// println!("run ended: {:?}", outcome.end_reason); -/// # Ok(()) -/// # } -/// ``` -#[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, - update_admission: AtomicU8, - wake: watch::Sender<()>, -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -#[repr(u8)] -enum UpdateAdmission { - Open, - Stopping, - Closed, -} - -impl UpdateAdmission { - fn from_raw(value: u8) -> Self { - match value { - value if value == Self::Open as u8 => Self::Open, - value if value == Self::Stopping as u8 => Self::Stopping, - value if value == Self::Closed as u8 => Self::Closed, - _ => unreachable!("managed update admission stores a known state"), - } - } -} - -impl StopSignal { - fn request(&self) { - self.begin_stopping(); - if !self.requested.swap(true, Ordering::AcqRel) { - self.wake.send_replace(()); - } - } - - fn update_admission(&self) -> UpdateAdmission { - UpdateAdmission::from_raw(self.update_admission.load(Ordering::Acquire)) - } - - fn begin_stopping(&self) { - let mut current = self.update_admission.load(Ordering::Acquire); - loop { - match UpdateAdmission::from_raw(current) { - UpdateAdmission::Open => match self.update_admission.compare_exchange_weak( - UpdateAdmission::Open as u8, - UpdateAdmission::Stopping as u8, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => return, - Err(observed) => current = observed, - }, - UpdateAdmission::Stopping | UpdateAdmission::Closed => return, - } - } - } - - fn close_updates(&self) { - self.update_admission - .store(UpdateAdmission::Closed as u8, Ordering::Release); - } - - fn is_requested(&self) -> bool { - self.requested.load(Ordering::Acquire) - } -} - -/// Cloneable control and observation capability for a managed task. -#[derive(Clone)] -pub struct ManagedClientHandle { - stop: Arc, - commands: mpsc::Sender, - max_live_target_generations: usize, - 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 authoritative durable latest-state observation. - /// - /// The current snapshot is marked as seen when subscribing. - /// The current snapshot is immediately available through `borrow()`. Updates - /// may coalesce; this is not a lossless event log. After the task terminates, - /// the sender closes but the final snapshot remains readable. Use - /// [`Self::subscribe`] for the lossy presentation-event stream. - pub fn subscribe_status(&self) -> ManagedStatusSubscription { - let mut status = self.status.clone(); - drop(status.borrow_and_update()); - status - } - - /// 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(self.subscribe_status()) - } - - /// Submit the complete desired target set, replacing any previous one. - /// - /// `targets` is not a delta: any live target whose id is missing from - /// `targets` is stopped and removed. A target that is still present with an - /// unchanged configuration continues undisturbed — it is not restarted and - /// does not get a new generation. A target present with a *changed* - /// configuration, or whose current generation already reached a terminal - /// state (e.g. failed, completed, or was closed by the peer), is treated as - /// a fresh start: it is retired and a new generation is created for it, even - /// if the configuration is byte-for-byte identical to the one that just - /// finished. - /// - /// This call enqueues the update without waiting for queue capacity; await - /// the returned [`ManagedCommandReceipt`] to observe whether — and how — it - /// was applied. - pub fn update_targets( - &self, - targets: Vec, - ) -> Result { - match self.stop.update_admission() { - UpdateAdmission::Open => {} - UpdateAdmission::Stopping => return Err(ManagedCommandError::Stopping), - UpdateAdmission::Closed => return Err(ManagedCommandError::DriverClosed), - } - if targets.len() > self.max_live_target_generations { - return Err(ManagedCommandError::TooManyTargets { - configured: targets.len(), - limit: self.max_live_target_generations, - }); - } - - // Observing `Open` only permits attempting submission. A successful - // `try_send` transfers command ownership to the driver. A concurrent - // stop may move admission to `Stopping` before this send, allowing the - // command to enqueue; the driver-side check immediately before - // `apply_targets` determines whether it applies or resolves as - // `Stopping`. Terminal sealing closes the receiver before draining - // accepted commands. - let (acknowledgement, receiver) = oneshot::channel(); - match self.commands.try_send(ManagedCommand::UpdateTargets { - targets, - acknowledgement, - }) { - Ok(()) => Ok(ManagedCommandReceipt { receiver }), - Err(mpsc::error::TrySendError::Full(_)) => match self.stop.update_admission() { - UpdateAdmission::Open => Err(ManagedCommandError::QueueFull), - UpdateAdmission::Stopping => Err(ManagedCommandError::Stopping), - UpdateAdmission::Closed => Err(ManagedCommandError::DriverClosed), - }, - Err(mpsc::error::TrySendError::Closed(_)) => Err(ManagedCommandError::DriverClosed), - } - } -} - -/// Receipt for a target-set transaction accepted by the driver queue. -#[must_use = "await the receipt to observe target-update application"] -pub struct ManagedCommandReceipt { - receiver: oneshot::Receiver>, -} - -impl ManagedCommandReceipt { - /// Block until the target-set transaction is applied or rejected. - /// - /// This is for synchronous callers such as [`crate::managed::BlockingManagedClient`]. It - /// must not be called from an asynchronous runtime; use the [`Future`] - /// implementation there instead. - pub fn blocking_wait(self) -> Result { - self.receiver - .blocking_recv() - .unwrap_or(Err(ManagedCommandApplyError::AcknowledgementDisconnected)) - } -} - -impl Future for ManagedCommandReceipt { - type Output = Result; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match Pin::new(&mut self.receiver).poll(cx) { - Poll::Ready(Ok(result)) => Poll::Ready(result), - Poll::Ready(Err(_)) => { - Poll::Ready(Err(ManagedCommandApplyError::AcknowledgementDisconnected)) - } - Poll::Pending => Poll::Pending, - } - } -} - -enum ManagedCommand { - UpdateTargets { - targets: Vec, - acknowledgement: - oneshot::Sender>, - }, -} - -/// Receipt resolving once stop is durably observed or the task is terminal. -/// -/// Resolves from authoritative status as soon as `stop_requested` is true or -/// the lifecycle is Completed, Failed, or Abandoned. It need not wait for -/// graceful cleanup after stop observation. -#[must_use = "await the receipt to observe durable stop acknowledgement"] -pub struct ManagedStopReceipt { - future: Pin + Send>>, -} - -impl ManagedStopReceipt { - fn new(mut status: ManagedStatusSubscription) -> Self { - let future = Box::pin(async move { - while !{ - let snapshot = status.borrow(); - snapshot.stop_requested - || matches!( - snapshot.lifecycle, - ManagedLifecycle::Completed - | ManagedLifecycle::Failed - | ManagedLifecycle::Abandoned - ) - } { - // A closed sender retains the terminal snapshot. Never wait - // indefinitely if the channel closes. - if status.changed().await.is_err() { - break; - } - } - }); - Self { future } - } -} - -impl Future for ManagedStopReceipt { - type Output = (); - - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - self.future.as_mut().poll(cx) - } -} - -type ConnectFuture = - Pin> + Send + 'static>>; type WakeFuture = Pin>> + Send + 'static>>; fn arm_wake(mut receiver: watch::Receiver<()>) -> WakeFuture { @@ -360,275 +46,6 @@ fn arm_wake(mut receiver: watch::Receiver<()>) -> WakeFuture { }) } -enum TargetPhase { - Pending { - client_config: ClientConfig, - }, - Connecting { - future: ConnectFuture, - }, - Opening { - client: AsyncClient, - open: Box, - }, - Active { - client: AsyncClient, - schedule: ProbeSchedule, - send_waiting: bool, - }, - 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 state always owns its finalized authoritative outcome. - Terminal { - outcome: Arc, - }, -} - -impl TargetPhase { - 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, - } - } - - fn client(&self) -> &AsyncClient { - match self { - Self::Opening { client, .. } - | Self::Active { client, .. } - | Self::Draining { client, .. } - | Self::Closing { client, .. } => client, - _ => panic!("managed target has no client in this phase"), - } - } - - fn into_client(self) -> AsyncClient { - match self { - Self::Opening { client, .. } - | Self::Active { client, .. } - | Self::Draining { client, .. } - | Self::Closing { client, .. } => client, - _ => panic!("managed target has no client in this phase"), - } - } -} - -#[derive(Default)] -struct TargetCounters { - packets_sent: u64, - /// Unique echo replies that arrived in order and before their probe timed - /// out. - /// - /// A reply that arrives behind an already-received sequence, or after its - /// own probe timed out, is counted by `late` and deliberately not here. - /// These counters are durable, unlike the lossy managed presentation - /// events. - replies_received: u64, - duplicates: u64, - late: u64, - warning_events: u64, -} - -impl TargetCounters { - /// Account for one client event. - /// - /// Events that carry no durable target-local count are ignored. Probe - /// sends are not counted here: `packets_sent` is authoritative in the - /// underlying client and is synchronized through - /// [`TargetRuntime::sync_packets_sent`]. - fn observe(&mut self, event: &ClientEvent) { - match event { - ClientEvent::EchoReply { .. } => { - self.replies_received = self.replies_received.saturating_add(1); - } - ClientEvent::DuplicateReply { .. } => { - self.duplicates = self.duplicates.saturating_add(1); - } - ClientEvent::LateReply { .. } => { - self.late = self.late.saturating_add(1); - } - ClientEvent::Warning { .. } => { - self.warning_events = self.warning_events.saturating_add(1); - } - ClientEvent::EchoSent { .. } - | ClientEvent::EchoLoss { .. } - | ClientEvent::SessionStarted(_) - | ClientEvent::NoTestCompleted(_) - | ClientEvent::SessionClosed { .. } => {} - } - } -} - -#[derive(Clone, Copy)] -enum RetirementReason { - Removed, - Replaced, -} - -impl From for ManagedTargetEndReason { - fn from(reason: RetirementReason) -> Self { - match reason { - RetirementReason::Removed => Self::Removed, - RetirementReason::Replaced => Self::Replaced, - } - } -} - -enum TargetMembership { - Desired, - Withdrawn(RetirementReason), -} - -impl TargetMembership { - fn is_desired(&self) -> bool { - matches!(self, Self::Desired) - } - - fn retirement_reason(&self) -> Option { - match self { - Self::Desired => None, - Self::Withdrawn(reason) => Some(*reason), - } - } -} - -struct TargetRuntime { - instance: TargetInstance, - config: ManagedTargetConfig, - membership: TargetMembership, - server_addr: Arc, - remote: Option, - counters: TargetCounters, - // None only while a synchronous operation owns the phase locally. - phase: Option, -} - -impl TargetRuntime { - fn phase(&self) -> &TargetPhase { - self.phase - .as_ref() - .expect("managed target phase is extracted") - } - - fn phase_mut(&mut self) -> &mut TargetPhase { - self.phase - .as_mut() - .expect("managed target phase is extracted") - } - - fn take_phase(&mut self) -> TargetPhase { - self.phase - .take() - .expect("managed target phase is extracted") - } - - fn restore_phase(&mut self, phase: TargetPhase) { - assert!( - self.phase.is_none(), - "managed target phase is already installed" - ); - self.phase = Some(phase); - } - - fn is_paced_active(&self) -> bool { - let phase = self.phase(); - self.membership.is_desired() && matches!(phase, TargetPhase::Active { .. }) - } - - /// Adopt the underlying client's authoritative sent count. - /// - /// The client owns `packets_sent`; this runtime only mirrors it so a - /// terminal outcome can report it after the client is gone. - fn sync_packets_sent(&mut self, client: &AsyncClient) { - self.counters.packets_sent = client.packets_sent(); - } - - fn sync_live_packets_sent(&mut self) { - self.counters.packets_sent = self.phase().client().packets_sent(); - } - - /// Build this target's durable outcome. - fn outcome( - &self, - end_reason: ManagedTargetEndReason, - cleanup_failure: Option, - ) -> ManagedTargetOutcome { - ManagedTargetOutcome { - target: self.instance.clone(), - server_addr: Arc::clone(&self.server_addr), - remote: self.remote, - end_reason, - packets_sent: self.counters.packets_sent, - replies_received: self.counters.replies_received, - duplicates: self.counters.duplicates, - late: self.counters.late, - warning_events: self.counters.warning_events, - cleanup_failure, - } - } -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -struct TimeoutBacklogEntry { - index: usize, - generation: u64, -} - -struct PlannedTarget { - target: ManagedTargetConfig, - client_config: ClientConfig, - generation: u64, -} - -struct PlannedRetirement { - index: usize, - reason: RetirementReason, - synchronous: bool, -} - -struct UpdatePlan { - retirements: Vec, - created: Vec, - next_generation: u64, - next_command_sequence: u64, - prospective_live_count: usize, -} - -fn synchronously_retireable(phase: &TargetPhase) -> bool { - match phase { - TargetPhase::Pending { .. } - | TargetPhase::Connecting { .. } - | TargetPhase::Terminal { .. } => true, - TargetPhase::Opening { open, .. } => !open.has_in_flight_work(), - TargetPhase::Active { .. } | TargetPhase::Draining { .. } | TargetPhase::Closing { .. } => { - false - } - } -} - -fn runtime_satisfies_target(runtime: &TargetRuntime, target: &ManagedTargetConfig) -> bool { - runtime.membership.is_desired() - && !matches!(runtime.phase(), TargetPhase::Terminal { .. }) - && runtime.config == *target -} - #[derive(Default)] struct OutcomeHistory { limit: usize, @@ -753,12 +170,6 @@ impl TaskPhase { } } -#[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 { @@ -864,86 +275,6 @@ impl ManagedClientTask { self.resources().status.send_replace(self.snapshot()); } - fn install_target_phase(&mut self, index: usize, phase: TargetPhase) { - let previous = self.targets[index].phase().lifecycle(); - self.targets[index].phase = Some(phase); - self.publish_target_phase(index, previous); - } - - fn publish_target_phase(&mut self, index: usize, previous: ManagedTargetLifecycle) { - let lifecycle = self.targets[index].phase().lifecycle(); - let desired = self.targets[index].membership.is_desired(); - let was_paced_active = desired && previous == ManagedTargetLifecycle::Active; - let paced_active = self.targets[index].is_paced_active(); - let now = Instant::now(); - match (was_paced_active, paced_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 { - self.targets[index].counters.observe(&event); - 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, - ) { - let outcome = Arc::new(self.targets[index].outcome(end_reason, cleanup_failure)); - // Finalize history before the phase publisher installs Terminal, updates - // pacing and publishes one complete status followed by the lifecycle event. - self.history.record(outcome.as_ref().clone()); - self.install_target_phase( - index, - TargetPhase::Terminal { - outcome: Arc::clone(&outcome), - }, - ); - self.publish_event(ManagedEvent::TargetFinished { 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, - phase: ManagedTargetFailurePhase, - error: ClientError, - now: Instant, - cleanup: OpenSessionFailureCleanup, - ) -> bool { - let primary_end = ManagedTargetEndReason::Failed(classify_client_error(phase, &error)); - let (TargetPhase::Opening { client, .. } | TargetPhase::Active { client, .. }) = - self.targets[index].phase_mut() - else { - unreachable!("open-session failure requires an opening or active target") - }; - client.discard_prepared_probe(); - self.targets[index].sync_live_packets_sent(); - match cleanup { - OpenSessionFailureCleanup::Drain => self.begin_drain(index, primary_end, now), - OpenSessionFailureCleanup::Close => self.begin_close(index, primary_end, None, now), - } - } - fn begin_running(&mut self) -> Result<(), ManagedDriverFailure> { tokio::runtime::Handle::try_current().map_err(|_| ManagedDriverFailure::NoTokioRuntime)?; self.phase = TaskPhase::Running; @@ -968,963 +299,6 @@ impl ManagedClientTask { self.burst_remaining = 0; } - fn effective_retirement(&self, index: usize) -> Option { - self.targets[index] - .membership - .retirement_reason() - .map(Into::into) - } - - fn process_commands(&mut self, cx: &mut Context<'_>) -> bool { - let mut immediate = false; - for _ in 0..COMMAND_WORK_BUDGET { - match self.commands.poll_recv(cx) { - Poll::Ready(Some(ManagedCommand::UpdateTargets { - targets, - acknowledgement, - })) => { - // This check is the driver-side application linearization point. - // `apply_targets` cannot suspend, so a stop observed here wins over - // this command; once it passes, this command may complete atomically. - let result = if !self.phase.is_running() - || self.resources().stop.update_admission() != UpdateAdmission::Open - { - Err(ManagedCommandApplyError::Stopping) - } else { - self.apply_targets(targets) - }; - let _ = acknowledgement.send(result); - immediate = true; - } - Poll::Ready(None) | Poll::Pending => break, - } - } - immediate - } - - fn apply_targets( - &mut self, - incoming: Vec, - ) -> Result { - let plan = self.plan_targets(incoming)?; - debug_assert!(plan.prospective_live_count <= self.config.max_live_target_generations); - - let now = Instant::now(); - let mut synchronous = HashSet::with_capacity(plan.retirements.len()); - let mut finished = Vec::new(); - for retirement in &plan.retirements { - let target = &mut self.targets[retirement.index]; - let was_paced_active = target.is_paced_active(); - target.membership = TargetMembership::Withdrawn(retirement.reason); - if retirement.synchronous { - synchronous.insert(retirement.index); - if !matches!(target.phase(), TargetPhase::Terminal { .. }) { - let outcome = Arc::new(target.outcome(retirement.reason.into(), None)); - self.history.record(outcome.as_ref().clone()); - finished.push(ManagedEvent::TargetFinished { outcome }); - } - } else if was_paced_active { - // Removal is directional: it may discard an elapsed gate but never - // lengthens an existing future stagger gate. - self.stagger_target_removed(now); - } - } - if !synchronous.is_empty() { - let mut index = 0; - self.targets.retain(|_| { - let keep = !synchronous.contains(&index); - index += 1; - keep - }); - self.rebase_target_cursors(); - } - - let mut created_instances = Vec::with_capacity(plan.created.len()); - for planned in plan.created { - let instance = TargetInstance { - id: planned.target.id.clone(), - generation: planned.generation, - }; - created_instances.push(instance.clone()); - self.targets.push(TargetRuntime { - instance, - server_addr: Arc::from(planned.target.server_addr.clone()), - config: planned.target, - membership: TargetMembership::Desired, - remote: None, - counters: TargetCounters::default(), - phase: Some(TargetPhase::Pending { - client_config: planned.client_config, - }), - }); - } - self.next_generation = plan.next_generation; - self.applied_command_sequence = plan.next_command_sequence; - let stopping = self - .targets - .iter() - .all(|runtime| !runtime.membership.is_desired()) - && self.config.completion == ManagedCompletionPolicy::FinishWhenQuiescent; - if stopping { - self.resources().stop.begin_stopping(); - self.phase = TaskPhase::Stopping; - } - - // A transaction has one externally visible durable point: all runtime and - // history changes above, then this exact status snapshot, then its events and - // acknowledgement. Do not use per-transition publishing helpers here. - let status = self.snapshot(); - self.resources().status.send_replace(Arc::clone(&status)); - if stopping { - self.publish_event(ManagedEvent::Stopping); - } - for event in finished { - self.publish_event(event); - } - for target in created_instances { - self.publish_event(ManagedEvent::TargetStateChanged { - target, - lifecycle: ManagedTargetLifecycle::Pending, - }); - } - Ok(ManagedCommandAcknowledgement { - sequence: self.applied_command_sequence, - status, - }) - } - - fn plan_targets( - &self, - incoming: Vec, - ) -> Result { - let mut ids = HashSet::with_capacity(incoming.len()); - let mut prepared = Vec::with_capacity(incoming.len()); - for target in incoming { - if !ids.insert(target.id.clone()) { - return Err(ManagedCommandApplyError::DuplicateTargetId { id: target.id }); - } - let mut client_config = self.config.client.clone(); - client_config.auth = target.auth.resolve(&self.config.client.auth); - validate_target_config(&client_config).map_err(|source| { - ManagedCommandApplyError::InvalidTarget { - id: target.id.clone(), - source, - } - })?; - prepared.push((target, client_config)); - } - let mut retirements = Vec::new(); - for (index, runtime) in self.targets.iter().enumerate() { - if !runtime.membership.is_desired() - || prepared - .iter() - .any(|(target, _)| runtime_satisfies_target(runtime, target)) - { - continue; - } - let reason = if prepared - .iter() - .any(|(target, _)| target.id == runtime.instance.id) - { - RetirementReason::Replaced - } else { - RetirementReason::Removed - }; - retirements.push(PlannedRetirement { - index, - reason, - synchronous: synchronously_retireable(runtime.phase()), - }); - } - - let mut next_generation = self.next_generation; - let mut created = Vec::new(); - for (target, client_config) in prepared { - if self - .targets - .iter() - .any(|runtime| runtime_satisfies_target(runtime, &target)) - { - continue; - } - let generation = next_generation; - next_generation = next_generation - .checked_add(1) - .ok_or(ManagedCommandApplyError::GenerationExhausted)?; - created.push(PlannedTarget { - target, - client_config, - generation, - }); - } - let next_command_sequence = self - .applied_command_sequence - .checked_add(1) - .ok_or(ManagedCommandApplyError::CommandSequenceExhausted)?; - let synchronous = retirements - .iter() - .filter(|retirement| retirement.synchronous) - .map(|retirement| retirement.index) - .collect::>(); - let prospective_live_count = self - .targets - .iter() - .enumerate() - .filter(|(index, _)| !synchronous.contains(index)) - .count() - .saturating_add(created.len()); - if prospective_live_count > self.config.max_live_target_generations { - return Err(ManagedCommandApplyError::LiveGenerationLimitExceeded { - required: prospective_live_count, - limit: self.config.max_live_target_generations, - }); - } - Ok(UpdatePlan { - retirements, - created, - next_generation, - next_command_sequence, - prospective_live_count, - }) - } - - fn prune_undesired_terminal(&mut self) { - if !self.targets.iter().any(|target| { - !target.membership.is_desired() - && matches!(target.phase(), TargetPhase::Terminal { .. }) - }) { - return; - } - self.targets.retain(|target| { - target.membership.is_desired() - || !matches!(target.phase(), TargetPhase::Terminal { .. }) - }); - self.rebase_target_cursors(); - self.replace_status(); - } - - fn rebase_target_cursors(&mut self) { - let len = self.targets.len(); - if len == 0 { - self.cursor = 0; - self.timeout_cursor = 0; - self.send_cursor = 0; - } else { - self.cursor %= len; - self.timeout_cursor %= len; - self.send_cursor %= len; - } - self.scan_remaining = self.scan_remaining.min(len); - self.burst_remaining = self.burst_remaining.min(len); - self.stagger_remaining = self.stagger_remaining.min(len); - } - - fn start_connecting(&mut self, index: usize) { - let TargetPhase::Pending { client_config } = self.targets[index].take_phase() else { - unreachable!("connecting starts from a pending target") - }; - let endpoint = Arc::clone(&self.targets[index].server_addr); - let future = Box::pin(async move { AsyncClient::connect(endpoint, client_config).await }); - self.targets[index].restore_phase(TargetPhase::Connecting { future }); - self.publish_target_phase(index, ManagedTargetLifecycle::Pending); - } - - fn poll_target(&mut self, index: usize, cx: &mut Context<'_>, now: Instant) -> bool { - let retirement = self.effective_retirement(index); - let stopping = self.phase.is_stopping() || retirement.is_some(); - match self.targets[index].phase_mut() { - TargetPhase::Pending { .. } => { - if stopping { - self.finish_target( - index, - retirement - .clone() - .unwrap_or(ManagedTargetEndReason::Stopped), - None, - ); - false - } else { - self.start_connecting(index); - true - } - } - TargetPhase::Connecting { future } => { - if stopping { - self.finish_target( - index, - retirement - .clone() - .unwrap_or(ManagedTargetEndReason::Stopped), - None, - ); - return false; - } - match future.as_mut().poll(cx) { - Poll::Pending => false, - Poll::Ready(Ok(client)) => { - self.targets[index].remote = Some(client.remote_addr()); - self.install_target_phase( - index, - TargetPhase::Opening { - client, - open: Box::new(AsyncOpenState::new()), - }, - ); - true - } - Poll::Ready(Err(error)) => { - self.fail_target(index, ManagedTargetFailurePhase::Connecting, error); - false - } - } - } - TargetPhase::Opening { client, open } => { - if stopping { - let retirement = retirement - .clone() - .unwrap_or(ManagedTargetEndReason::Stopped); - if !open.has_in_flight_work() { - self.targets[index].sync_live_packets_sent(); - self.finish_target(index, retirement, None); - return false; - } - open.request_stop_after_current_attempt(); - } - match client.poll_open(open, cx) { - Poll::Pending => false, - Poll::Ready(Ok(OpenOutcome::Started(started))) => { - let schedule = match ProbeSchedule::new( - started.at.mono, - &started.negotiation.accepted, - ) { - Ok(schedule) => schedule, - Err(error) => { - return self.begin_open_session_failure( - index, - ManagedTargetFailurePhase::Timing, - error, - now, - OpenSessionFailureCleanup::Close, - ) - } - }; - self.publish_client_events( - index, - vec![ClientEvent::SessionStarted(started)], - ); - if stopping { - self.targets[index].sync_live_packets_sent(); - self.begin_drain( - index, - retirement - .clone() - .unwrap_or(ManagedTargetEndReason::Stopped), - now, - ) - } else { - let previous = self.targets[index].take_phase(); - self.targets[index].restore_phase(TargetPhase::Active { - client: previous.into_client(), - schedule, - send_waiting: false, - }); - self.publish_target_phase(index, ManagedTargetLifecycle::Opening); - true - } - } - Poll::Ready(Ok(OpenOutcome::NoTestCompleted(completed))) => { - self.publish_client_events( - index, - vec![ClientEvent::NoTestCompleted(completed)], - ); - let end_reason = if stopping { - retirement - .clone() - .unwrap_or(ManagedTargetEndReason::Stopped) - } else { - ManagedTargetEndReason::NoTestComplete - }; - self.finish_target(index, end_reason, None); - false - } - Poll::Ready(Err(error)) => { - if stopping { - let TargetPhase::Opening { open, .. } = self.targets[index].phase() - else { - unreachable!() - }; - let cleanup_failure = - (!open.stopped_after_current_attempt()).then(|| { - classify_client_error( - ManagedTargetFailurePhase::Opening, - &error, - ) - }); - self.finish_target( - index, - retirement - .clone() - .unwrap_or(ManagedTargetEndReason::Stopped), - cleanup_failure, - ); - } else { - self.fail_target(index, ManagedTargetFailurePhase::Opening, error); - } - false - } - } - } - TargetPhase::Active { client, .. } => { - if stopping { - client.discard_prepared_probe(); - self.targets[index].sync_live_packets_sent(); - return self.begin_drain( - index, - retirement - .clone() - .unwrap_or(ManagedTargetEndReason::Stopped), - now, - ); - } - if client - .next_probe_timeout_deadline() - .is_some_and(|deadline| deadline <= now) - { - return true; - } - 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, - ManagedTargetFailurePhase::Receiving, - error, - now, - OpenSessionFailureCleanup::Close, - ); - } - }; - if self.targets[index].phase().client().is_peer_closed() { - self.targets[index].sync_live_packets_sent(); - self.finish_target(index, ManagedTargetEndReason::PeerClosed, None); - return false; - } - let TargetPhase::Active { - client, schedule, .. - } = self.targets[index].phase_mut() - else { - unreachable!() - }; - schedule.permit_probe_at(now); - if schedule.is_finished() && !client.has_pending_probes() { - self.targets[index].sync_live_packets_sent(); - return self.begin_drain(index, ManagedTargetEndReason::TestComplete, now); - } - received - } - TargetPhase::Draining { client, .. } => { - if client - .next_probe_timeout_deadline() - .is_some_and(|timeout| timeout <= now) - { - return true; - } - let mut retained_state_changed = false; - 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].sync_live_packets_sent(); - let TargetPhase::Draining { - primary_end, - cleanup_failure, - .. - } = self.targets[index].phase_mut() - else { - unreachable!() - }; - cleanup_failure.get_or_insert_with(|| { - classify_client_error(ManagedTargetFailurePhase::Receiving, &error) - }); - let primary_end = primary_end.clone(); - let cleanup_failure = cleanup_failure.clone(); - return self.begin_close(index, primary_end, cleanup_failure, now); - } - }; - if self.targets[index].phase().client().is_peer_closed() { - self.targets[index].sync_live_packets_sent(); - let TargetPhase::Draining { - cleanup_failure, .. - } = self.targets[index].phase() - else { - unreachable!() - }; - self.finish_target( - index, - ManagedTargetEndReason::PeerClosed, - cleanup_failure.clone(), - ); - return false; - } - if retained_state_changed { - let TargetPhase::Draining { - client, - drain_started_at, - .. - } = self.targets[index].phase() - else { - unreachable!() - }; - let candidate = self.drain_deadline(client, *drain_started_at); - let TargetPhase::Draining { - deadline, - primary_end, - cleanup_failure, - .. - } = self.targets[index].phase_mut() - else { - unreachable!() - }; - match candidate { - Some(candidate) => *deadline = (*deadline).min(candidate), - None => { - cleanup_failure.get_or_insert_with(duration_overflow_failure); - let primary_end = primary_end.clone(); - let cleanup_failure = cleanup_failure.clone(); - return self.begin_close(index, primary_end, cleanup_failure, now); - } - } - } - let TargetPhase::Draining { - deadline, - primary_end, - cleanup_failure, - post_deadline_receives_remaining, - .. - } = self.targets[index].phase_mut() - else { - unreachable!() - }; - if now >= *deadline { - if received && *post_deadline_receives_remaining > 1 { - *post_deadline_receives_remaining -= 1; - true - } else { - let primary_end = primary_end.clone(); - let cleanup_failure = cleanup_failure.clone(); - self.begin_close(index, primary_end, cleanup_failure, now) - } - } else { - received - } - } - TargetPhase::Closing { client, .. } => { - let result = client.poll_close(cx); - let TargetPhase::Closing { - deadline, - primary_end, - cleanup_failure, - .. - } = self.targets[index].phase() - else { - unreachable!() - }; - let deadline = *deadline; - let primary_end = primary_end.clone(); - let cleanup_failure = cleanup_failure.clone(); - match result { - Poll::Pending => { - if now >= deadline { - self.targets[index].sync_live_packets_sent(); - self.finish_target( - index, - primary_end, - cleanup_failure.or_else(|| Some(close_timeout_failure())), - ); - } - false - } - Poll::Ready(Ok(events)) => { - self.targets[index].sync_live_packets_sent(); - self.publish_client_events(index, events); - self.finish_target(index, primary_end, cleanup_failure); - false - } - Poll::Ready(Err(error)) => { - self.targets[index].sync_live_packets_sent(); - let cleanup = - classify_client_error(ManagedTargetFailurePhase::Closing, &error); - self.finish_target(index, primary_end, cleanup_failure.or(Some(cleanup))); - false - } - } - } - TargetPhase::Terminal { .. } => false, - } - } - - fn begin_drain( - &mut self, - index: usize, - primary_end: ManagedTargetEndReason, - _now: Instant, - ) -> bool { - let drain_started_at = Instant::now(); - let Some(deadline) = - self.drain_deadline(self.targets[index].phase().client(), drain_started_at) - else { - return self.begin_close( - index, - primary_end, - Some(duration_overflow_failure()), - drain_started_at, - ); - }; - let previous = self.targets[index].take_phase(); - let lifecycle = previous.lifecycle(); - self.targets[index].restore_phase(TargetPhase::Draining { - client: previous.into_client(), - drain_started_at, - deadline, - primary_end, - cleanup_failure: None, - post_deadline_receives_remaining: POST_DEADLINE_RECEIVE_BUDGET, - }); - self.publish_target_phase(index, lifecycle); - 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, - 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 - }); - let previous = self.targets[index].take_phase(); - let lifecycle = previous.lifecycle(); - self.targets[index].restore_phase(TargetPhase::Closing { - client: previous.into_client(), - deadline, - primary_end, - cleanup_failure, - }); - self.publish_target_phase(index, lifecycle); - true - } - - fn active_count(&self) -> usize { - self.targets - .iter() - .filter(|target| target.is_paced_active()) - .count() - } - - fn active_stagger_spacing(&self) -> Option { - let mut active = 0; - let mut minimum: Option = None; - for target in &self.targets { - let TargetPhase::Active { schedule, .. } = target.phase() else { - continue; - }; - if !target.membership.is_desired() { - continue; - } - active += 1; - minimum = minimum.into_iter().chain(Some(schedule.interval())).min(); - } - minimum.map(|interval| 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); - } - - /// The next stagger gate after a send was accepted at `accepted_at`. - /// - /// The gate advances on its **own** cadence rather than being re-anchored - /// on every observed accept time. Anchoring on `accepted_at` folds each - /// send's wakeup latency into the next slot and then into every slot after - /// it, so the gate — and with it the whole group's send cadence — drifts - /// later without bound while the per-target schedules stay on their - /// absolute grid. The drift is not hypothetical: it shows up as a - /// monotonically growing `EchoSent::timer_error` and, once it exceeds one - /// interval, as probe slots the schedule then skips. - /// - /// `accepted_at` is still the anchor in the two cases where continuing the - /// old cadence would be wrong: there is no previous gate, or the pacer has - /// fallen a full slot behind it. Re-anchoring there is what keeps a gate - /// stranded in the past from releasing a burst of catch-up sends once the - /// group starts moving again. - fn next_stagger_gate(&self, spacing: Duration, accepted_at: Instant) -> Option { - let base = match self.send_gate { - Some(gate) if gate <= accepted_at && accepted_at.duration_since(gate) < spacing => gate, - _ => accepted_at, - }; - base.checked_add(spacing) - } - - fn record_stagger_acceptance( - &mut self, - result: SendResult, - stagger_spacing: Option, - accepted_at: Instant, - ) { - if !result.accepted() { - return; - } - let Some(spacing) = stagger_spacing else { - return; - }; - self.last_stagger_send = Some(accepted_at); - self.send_gate = self.next_stagger_gate(spacing, accepted_at); - } - - fn poll_one_send( - &mut self, - index: usize, - cx: &mut Context<'_>, - now: Instant, - stagger_spacing: Option, - ) -> SendResult { - if !self.targets[index].membership.is_desired() { - if let TargetPhase::Active { send_waiting, .. } = self.targets[index].phase_mut() { - *send_waiting = false; - } - return SendResult::NotAttempted; - } - let phase = self.targets[index].take_phase(); - let TargetPhase::Active { - mut client, - mut schedule, - mut send_waiting, - } = phase - else { - self.targets[index].restore_phase(phase); - return SendResult::NotAttempted; - }; - // A burst pass may span several slots; validate this target against - // fresh time rather than the pass's earlier scheduling snapshot. - let now = now.max(Instant::now()); - if client - .next_probe_timeout_deadline() - .is_some_and(|deadline| deadline <= now) - { - send_waiting = false; - self.targets[index].restore_phase(TargetPhase::Active { - client, - schedule, - send_waiting, - }); - return SendResult::NotAttempted; - } - if !schedule.permit_probe_at(now) { - send_waiting = false; - let pending = client.has_pending_probes(); - if !pending { - self.targets[index].sync_packets_sent(&client); - } - self.targets[index].restore_phase(TargetPhase::Active { - client, - schedule, - send_waiting, - }); - if pending { - return SendResult::NotAttempted; - } - self.begin_drain(index, ManagedTargetEndReason::TestComplete, now); - return SendResult::Ready { accepted: false }; - } - if schedule - .next_send_deadline() - .is_none_or(|deadline| deadline > now) - { - send_waiting = false; - self.targets[index].restore_phase(TargetPhase::Active { - client, - schedule, - send_waiting, - }); - return SendResult::NotAttempted; - } - let commit = - match schedule.preflight_managed_commit(schedule.next_send_deadline().unwrap(), now) { - Ok(commit) => commit, - Err(error) => { - self.targets[index].restore_phase(TargetPhase::Active { - client, - schedule, - send_waiting, - }); - self.begin_open_session_failure( - index, - ManagedTargetFailurePhase::Timing, - error, - now, - OpenSessionFailureCleanup::Drain, - ); - return SendResult::Failed { accepted: false }; - } - }; - let scheduled_at = commit.scheduled_at; - let result = client.poll_send_probe(cx); - let receipt = match &result { - Poll::Ready(Ok(receipt)) - | Poll::Ready(Err(SendProbeError::AfterCommit { receipt, .. })) => Some(*receipt), - Poll::Pending | Poll::Ready(Err(SendProbeError::NotCommitted(_))) => None, - }; - let accepted = receipt.is_some(); - 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()); - if accepted { - schedule.commit(commit); - } - let schedule_error = if accepted { - schedule.skip_missed_slots_at(Instant::now()).err() - } else { - None - }; - self.targets[index].sync_packets_sent(&client); - send_waiting = result.is_pending(); - self.targets[index].restore_phase(TargetPhase::Active { - client, - schedule, - send_waiting, - }); - // A committed send must reach the lifecycle/accounting stream even - // when post-send processing fails. Publish before failure cleanup. - let sent_event = receipt.map(|receipt| { - let mut event = ClientEvent::from(receipt); - if let ClientEvent::EchoSent { - scheduled_at: intended, - timer_error, - .. - } = &mut event - { - *intended = Some(scheduled_at); - *timer_error = Some(instant_abs_diff(receipt.sent_at.mono, scheduled_at)); - } - event - }); - match result { - Poll::Pending => SendResult::Pending, - Poll::Ready(Ok(_)) => { - self.publish_client_events(index, vec![sent_event.unwrap()]); - if let Some(error) = schedule_error { - self.begin_open_session_failure( - index, - ManagedTargetFailurePhase::Timing, - error, - now, - OpenSessionFailureCleanup::Drain, - ); - SendResult::Failed { accepted } - } else { - SendResult::Ready { accepted } - } - } - Poll::Ready(Err(error)) => { - if let Some(event) = sent_event { - self.publish_client_events(index, vec![event]); - } - let error = match error { - SendProbeError::NotCommitted(source) => source, - SendProbeError::AfterCommit { source, .. } => *source, - }; - self.begin_open_session_failure( - index, - 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 !self.targets[index].is_paced_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; @@ -1943,161 +317,6 @@ impl ManagedClientTask { immediate || self.scan_remaining > 0 } - fn poll_timeout_pass(&mut self, now: Instant) -> bool { - let target_len = self.targets.len(); - if target_len == 0 { - return false; - } - - let mut transitions_remaining = TIMEOUT_WORK_BUDGET; - for _ in 0..target_len.min(TIMEOUT_WORK_BUDGET) { - if transitions_remaining == 0 { - break; - } - let index = self.timeout_cursor; - self.timeout_cursor = (self.timeout_cursor + 1) % target_len; - if self.target_has_due_timeout(index, now) { - let generation = self.targets[index].instance.generation; - let step = self.poll_target_timeout(index, now); - transitions_remaining -= 1; - if step.more_due { - self.enqueue_timeout_backlog(index, generation); - } - } - } - - while transitions_remaining > 0 { - let Some(entry) = self.timeout_backlog.pop_front() else { - break; - }; - let valid = self - .targets - .get(entry.index) - .is_some_and(|target| target.instance.generation == entry.generation) - && self.target_has_due_timeout(entry.index, now); - if !valid { - continue; - } - - let step = self.poll_target_timeout(entry.index, now); - transitions_remaining -= 1; - if step.more_due { - self.enqueue_timeout_backlog(entry.index, entry.generation); - } - } - - !self.timeout_backlog.is_empty() - } - - fn enqueue_timeout_backlog(&mut self, index: usize, generation: u64) { - if self.timeout_backlog.len() == TIMEOUT_WORK_BUDGET - || self - .timeout_backlog - .iter() - .any(|entry| entry.index == index && entry.generation == generation) - { - return; - } - self.timeout_backlog - .push_back(TimeoutBacklogEntry { index, generation }); - } - - fn target_has_due_timeout(&self, index: usize, now: Instant) -> bool { - #[cfg(test)] - self.timeout_inspections.borrow_mut().push(index); - match self.targets[index].phase() { - TargetPhase::Active { client, .. } => { - !self.phase.is_stopping() - && self.effective_retirement(index).is_none() - && client - .next_probe_timeout_deadline() - .is_some_and(|deadline| deadline <= now) - } - TargetPhase::Draining { client, .. } => client - .next_probe_timeout_deadline() - .is_some_and(|timeout| timeout <= now), - _ => false, - } - } - - fn poll_target_timeout(&mut self, index: usize, now: Instant) -> TimeoutStep { - match self.targets[index].phase_mut() { - TargetPhase::Active { client, .. } => match client.poll_timeouts_bounded_at(now, 1) { - Ok(batch) => { - self.publish_client_events(index, batch.events); - TimeoutStep { - more_due: batch.more_due, - } - } - Err(error) => { - self.begin_open_session_failure( - index, - ManagedTargetFailurePhase::Timing, - error, - now, - OpenSessionFailureCleanup::Close, - ); - TimeoutStep::NONE - } - }, - TargetPhase::Draining { client, .. } => match client.poll_timeouts_bounded_at(now, 1) { - Ok(batch) => { - self.publish_client_events(index, batch.events); - let TargetPhase::Draining { - client, - drain_started_at, - .. - } = self.targets[index].phase() - else { - unreachable!() - }; - let candidate = self.drain_deadline(client, *drain_started_at); - let TargetPhase::Draining { - deadline, - primary_end, - cleanup_failure, - .. - } = self.targets[index].phase_mut() - else { - unreachable!() - }; - match candidate { - Some(candidate) => *deadline = (*deadline).min(candidate), - None => { - cleanup_failure.get_or_insert_with(duration_overflow_failure); - let primary_end = primary_end.clone(); - let cleanup_failure = cleanup_failure.clone(); - self.begin_close(index, primary_end, cleanup_failure, now); - return TimeoutStep::NONE; - } - } - TimeoutStep { - more_due: batch.more_due, - } - } - Err(error) => { - self.targets[index].sync_live_packets_sent(); - let TargetPhase::Draining { - primary_end, - cleanup_failure, - .. - } = self.targets[index].phase_mut() - else { - unreachable!() - }; - cleanup_failure.get_or_insert_with(|| { - classify_client_error(ManagedTargetFailurePhase::Timing, &error) - }); - let primary_end = primary_end.clone(); - let cleanup_failure = cleanup_failure.clone(); - self.begin_close(index, primary_end, cleanup_failure, now); - TimeoutStep::NONE - } - }, - _ => TimeoutStep::NONE, - } - } - fn all_targets_terminal(&self) -> bool { self.targets .iter() @@ -2244,37 +463,6 @@ impl ManagedClientTask { } } -// A release gate delays a send opportunity; it cannot create one. -fn gated_send_deadline(cadence: Option, gate: Option) -> Option { - cadence.map(|deadline| gate.map_or(deadline, |gate| deadline.max(gate))) -} - -#[derive(Clone, Copy, Eq, PartialEq)] -enum SendResult { - NotAttempted, - Pending, - Ready { accepted: bool }, - Failed { accepted: bool }, -} - -#[derive(Clone, Copy)] -struct TimeoutStep { - more_due: bool, -} - -impl TimeoutStep { - const NONE: Self = Self { more_due: false }; -} - -impl SendResult { - fn accepted(self) -> bool { - matches!( - self, - Self::Ready { accepted: true } | Self::Failed { accepted: true } - ) - } -} - impl Future for ManagedClientTask { type Output = ManagedOutcome; @@ -2367,197 +555,15 @@ impl Drop for ManagedClientTask { } } -fn build_task( - config: ManagedClientConfig, - targets: Vec, -) -> Result<(ManagedClientTask, ManagedClientHandle), ManagedConfigError> { - if config.event_capacity == 0 { - return Err(ManagedConfigError::ZeroEventCapacity); - } - if config.event_capacity > MAX_BROADCAST_CHANNEL_CAPACITY { - return Err(ManagedConfigError::EventCapacityTooLarge { - configured: config.event_capacity, - maximum: MAX_BROADCAST_CHANNEL_CAPACITY, - }); - } - if config.command_capacity == 0 { - return Err(ManagedConfigError::ZeroCommandCapacity); - } - if config.command_capacity > tokio::sync::Semaphore::MAX_PERMITS { - return Err(ManagedConfigError::CommandCapacityTooLarge { - configured: config.command_capacity, - maximum: tokio::sync::Semaphore::MAX_PERMITS, - }); - } - 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.auth = target.auth.resolve(&config.client.auth); - validate_target_config(&client_config).map_err(|source| { - ManagedConfigError::InvalidTarget { - id: target.id.clone(), - source, - } - })?; - runtimes.push(TargetRuntime { - instance: TargetInstance { - id: target.id.clone(), - generation, - }, - config: target.clone(), - membership: TargetMembership::Desired, - server_addr: Arc::from(target.server_addr), - remote: None, - counters: TargetCounters::default(), - phase: Some(TargetPhase::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 stop = Arc::new(StopSignal { - requested: AtomicBool::new(false), - update_admission: AtomicU8::new(UpdateAdmission::Open as u8), - wake: wake_sender, - }); - let history = OutcomeHistory::new(config.outcome_history_limit); - let initial_targets = runtimes - .iter() - .map(|target| ManagedTargetStatus { - target: target.instance.clone(), - desired: target.membership.is_desired(), - lifecycle: ManagedTargetLifecycle::Pending, - outcome: None, - server_addr: Arc::clone(&target.server_addr), - remote: None, - }) - .collect::>(); - let phase = TaskPhase::NotStarted; - let initial = Arc::new(ManagedStatus { - lifecycle: phase.lifecycle(), - stop_requested: false, - applied_command_sequence: 0, - 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: phase.final_outcome(), - }); - let (status_sender, status_receiver) = watch::channel(initial); - let (command_sender, command_receiver) = mpsc::channel(config.command_capacity); - let max_live_target_generations = config.max_live_target_generations; - let resources = TaskResources { - status: status_sender, - events: Some(event_sender), - stop: Arc::clone(&stop), - }; - let task = ManagedClientTask { - #[cfg(test)] - timeout_inspections: std::cell::RefCell::new(Vec::new()), - phase, - config, - targets: runtimes, - commands: command_receiver, - history, - resources: Some(resources), - wake: Some(arm_wake(wake_receiver)), - timer: None, - cursor: 0, - timeout_cursor: 0, - timeout_backlog: VecDeque::with_capacity(TIMEOUT_WORK_BUDGET), - scan_remaining: 0, - send_cursor: 0, - burst_remaining: 0, - stagger_remaining: 0, - send_gate: None, - last_stagger_send: None, - stop_observed: false, - next_generation, - applied_command_sequence: 0, - }; - let handle = ManagedClientHandle { - stop, - commands: command_sender, - max_live_target_generations, - status: status_receiver, - events: weak_events, - }; - Ok((task, handle)) -} - -fn validate_target_config(config: &ClientConfig) -> Result<(), ClientError> { - config.open.validate()?; - SessionMachine::validate_config(config)?; - if let Some(ttl) = config.socket.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)] mod tests { use super::*; + use crate::{ + async_client::AsyncClient, + managed::{schedule::ProbeSchedule, ManagedTargetConfig}, + ClientConfig, + }; + use std::time::Duration; // Real EAGAIN needs privileged queue shaping. Inspect only the owning // deadline selection here, using a production socket without sending. @@ -2638,54 +644,4 @@ mod tests { } }); } - - #[test] - fn release_gate_only_delays_an_existing_send_deadline() { - let earlier = Instant::now(); - let later = earlier + Duration::from_secs(1); - for (cadence, gate, expected) in [ - (None, None, None), - (None, Some(earlier), None), - (None, Some(later), None), - (Some(earlier), None, Some(earlier)), - (Some(later), Some(earlier), Some(later)), - (Some(earlier), Some(earlier), Some(earlier)), - (Some(earlier), Some(later), Some(later)), - ] { - assert_eq!(gated_send_deadline(cadence, gate), expected); - } - } - - // Poll latency through public events cannot distinguish an O(n) discovery - // scan from a bounded scan reliably. This tiny inspection trace measures - // work directly, including empty targets, and checks eventual coverage - // without asserting a cursor value or a particular visitation order. - #[test] - fn timeout_discovery_has_bounded_work_and_eventually_visits_every_target() { - let count = TIMEOUT_WORK_BUDGET * 2 + 1; - let targets = (0..count) - .map(|i| ManagedTargetConfig::new(i.to_string(), "127.0.0.1:2112")) - .collect(); - let (mut task, _handle) = ManagedClient::task( - ManagedClientConfig { - max_live_target_generations: count, - ..ManagedClientConfig::default() - }, - targets, - ) - .unwrap(); - let mut visited = HashSet::new(); - for _ in 0..count.div_ceil(TIMEOUT_WORK_BUDGET) { - task.timeout_inspections.borrow_mut().clear(); - task.poll_timeout_pass(Instant::now()); - let inspected = task.timeout_inspections.borrow(); - assert!(inspected.len() <= TIMEOUT_WORK_BUDGET); - visited.extend(inspected.iter().copied()); - } - assert_eq!( - visited.len(), - count, - "timeout discovery must reach all targets" - ); - } } diff --git a/crates/irtt-client/src/managed/task/control.rs b/crates/irtt-client/src/managed/task/control.rs new file mode 100644 index 0000000..ad24508 --- /dev/null +++ b/crates/irtt-client/src/managed/task/control.rs @@ -0,0 +1,511 @@ +//! Task construction, public control/observation handles, and stop/update admission. + +use std::{ + collections::{HashSet, VecDeque}, + fmt, + future::Future, + pin::Pin, + sync::{ + atomic::{AtomicBool, AtomicU8, Ordering}, + Arc, + }, + task::{Context, Poll}, + time::Instant, +}; + +use tokio::sync::{broadcast, mpsc, oneshot, watch}; + +use crate::{ + managed::{ + ManagedClientConfig, ManagedCommandAcknowledgement, ManagedCommandApplyError, + ManagedCommandError, ManagedCompletionPolicy, ManagedConfigError, ManagedEvent, + ManagedEventSubscription, ManagedLifecycle, ManagedStatus, ManagedStatusSubscription, + ManagedSubscribeError, ManagedTargetConfig, ManagedTargetLifecycle, ManagedTargetStatus, + TargetInstance, + }, + session::machine::SessionMachine, + socket_options::validate_ttl, + ClientConfig, ClientError, +}; + +use super::{ + arm_wake, + target::{TargetCounters, TargetMembership, TargetPhase, TargetRuntime}, + timeouts::TIMEOUT_WORK_BUDGET, + ManagedClientTask, OutcomeHistory, TaskPhase, TaskResources, +}; + +// Tokio 1.53 broadcast asserts this exact bound before rounding capacity to a +// power of two; Tokio does not expose it as a public constant. +const MAX_BROADCAST_CHANNEL_CAPACITY: usize = usize::MAX >> 1; + +/// Entry point for constructing a unified Tokio managed task. +/// +/// # Example +/// +/// Construct the task and its separate control handle, subscribe before +/// starting the task, then run it under the caller's Tokio runtime. A reachable +/// server is required at runtime, so this example only type-checks. +/// +/// ```no_run +/// use irtt_client::{ClientConfig, ClientEvent}; +/// use irtt_client::managed::{ +/// ManagedClient, ManagedClientConfig, ManagedEvent, ManagedTargetConfig, +/// }; +/// +/// # async fn run() -> Result<(), Box> { +/// let config = ManagedClientConfig { +/// client: ClientConfig { +/// request: irtt_client::SessionRequest { +/// duration: Some(std::time::Duration::from_secs(10)), +/// ..irtt_client::SessionRequest::default() +/// }, +/// ..ClientConfig::default() +/// }, +/// ..ManagedClientConfig::default() +/// }; +/// let targets = vec![ManagedTargetConfig::new("edge", "127.0.0.1:2112")]; +/// let (task, handle) = ManagedClient::task(config, targets)?; +/// let mut events = handle.subscribe()?; +/// let driver = tokio::spawn(task); +/// +/// while let Ok(event) = events.recv().await { +/// if let ManagedEvent::Client { event: ClientEvent::EchoReply { .. }, .. } = event { +/// println!("received a probe reply"); +/// } +/// } +/// let outcome = driver.await?; +/// println!("run ended: {:?}", outcome.end_reason); +/// # Ok(()) +/// # } +/// ``` +#[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)] +pub(super) struct StopSignal { + requested: AtomicBool, + update_admission: AtomicU8, + wake: watch::Sender<()>, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[repr(u8)] +pub(super) enum UpdateAdmission { + Open, + Stopping, + Closed, +} + +impl UpdateAdmission { + fn from_raw(value: u8) -> Self { + match value { + value if value == Self::Open as u8 => Self::Open, + value if value == Self::Stopping as u8 => Self::Stopping, + value if value == Self::Closed as u8 => Self::Closed, + _ => unreachable!("managed update admission stores a known state"), + } + } +} + +impl StopSignal { + fn request(&self) { + self.begin_stopping(); + if !self.requested.swap(true, Ordering::AcqRel) { + self.wake.send_replace(()); + } + } + + pub(super) fn update_admission(&self) -> UpdateAdmission { + UpdateAdmission::from_raw(self.update_admission.load(Ordering::Acquire)) + } + + pub(super) fn begin_stopping(&self) { + let mut current = self.update_admission.load(Ordering::Acquire); + loop { + match UpdateAdmission::from_raw(current) { + UpdateAdmission::Open => match self.update_admission.compare_exchange_weak( + UpdateAdmission::Open as u8, + UpdateAdmission::Stopping as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return, + Err(observed) => current = observed, + }, + UpdateAdmission::Stopping | UpdateAdmission::Closed => return, + } + } + } + + pub(super) fn close_updates(&self) { + self.update_admission + .store(UpdateAdmission::Closed as u8, Ordering::Release); + } + + pub(super) fn is_requested(&self) -> bool { + self.requested.load(Ordering::Acquire) + } +} + +/// Cloneable control and observation capability for a managed task. +#[derive(Clone)] +pub struct ManagedClientHandle { + stop: Arc, + commands: mpsc::Sender, + max_live_target_generations: usize, + 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 authoritative durable latest-state observation. + /// + /// The current snapshot is marked as seen when subscribing. + /// The current snapshot is immediately available through `borrow()`. Updates + /// may coalesce; this is not a lossless event log. After the task terminates, + /// the sender closes but the final snapshot remains readable. Use + /// [`Self::subscribe`] for the lossy presentation-event stream. + pub fn subscribe_status(&self) -> ManagedStatusSubscription { + let mut status = self.status.clone(); + drop(status.borrow_and_update()); + status + } + + /// 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(self.subscribe_status()) + } + + /// Submit the complete desired target set, replacing any previous one. + /// + /// `targets` is not a delta: any live target whose id is missing from + /// `targets` is stopped and removed. A target that is still present with an + /// unchanged configuration continues undisturbed — it is not restarted and + /// does not get a new generation. A target present with a *changed* + /// configuration, or whose current generation already reached a terminal + /// state (e.g. failed, completed, or was closed by the peer), is treated as + /// a fresh start: it is retired and a new generation is created for it, even + /// if the configuration is byte-for-byte identical to the one that just + /// finished. + /// + /// This call enqueues the update without waiting for queue capacity; await + /// the returned [`ManagedCommandReceipt`] to observe whether — and how — it + /// was applied. + pub fn update_targets( + &self, + targets: Vec, + ) -> Result { + match self.stop.update_admission() { + UpdateAdmission::Open => {} + UpdateAdmission::Stopping => return Err(ManagedCommandError::Stopping), + UpdateAdmission::Closed => return Err(ManagedCommandError::DriverClosed), + } + if targets.len() > self.max_live_target_generations { + return Err(ManagedCommandError::TooManyTargets { + configured: targets.len(), + limit: self.max_live_target_generations, + }); + } + + // Observing `Open` only permits attempting submission. A successful + // `try_send` transfers command ownership to the driver. A concurrent + // stop may move admission to `Stopping` before this send, allowing the + // command to enqueue; the driver-side check immediately before + // `apply_targets` determines whether it applies or resolves as + // `Stopping`. Terminal sealing closes the receiver before draining + // accepted commands. + let (acknowledgement, receiver) = oneshot::channel(); + match self.commands.try_send(ManagedCommand::UpdateTargets { + targets, + acknowledgement, + }) { + Ok(()) => Ok(ManagedCommandReceipt { receiver }), + Err(mpsc::error::TrySendError::Full(_)) => match self.stop.update_admission() { + UpdateAdmission::Open => Err(ManagedCommandError::QueueFull), + UpdateAdmission::Stopping => Err(ManagedCommandError::Stopping), + UpdateAdmission::Closed => Err(ManagedCommandError::DriverClosed), + }, + Err(mpsc::error::TrySendError::Closed(_)) => Err(ManagedCommandError::DriverClosed), + } + } +} + +/// Receipt for a target-set transaction accepted by the driver queue. +#[must_use = "await the receipt to observe target-update application"] +pub struct ManagedCommandReceipt { + receiver: oneshot::Receiver>, +} + +impl ManagedCommandReceipt { + /// Block until the target-set transaction is applied or rejected. + /// + /// This is for synchronous callers such as [`crate::managed::BlockingManagedClient`]. It + /// must not be called from an asynchronous runtime; use the [`Future`] + /// implementation there instead. + pub fn blocking_wait(self) -> Result { + self.receiver + .blocking_recv() + .unwrap_or(Err(ManagedCommandApplyError::AcknowledgementDisconnected)) + } +} + +impl Future for ManagedCommandReceipt { + type Output = Result; + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match Pin::new(&mut self.receiver).poll(cx) { + Poll::Ready(Ok(result)) => Poll::Ready(result), + Poll::Ready(Err(_)) => { + Poll::Ready(Err(ManagedCommandApplyError::AcknowledgementDisconnected)) + } + Poll::Pending => Poll::Pending, + } + } +} + +pub(super) enum ManagedCommand { + UpdateTargets { + targets: Vec, + acknowledgement: + oneshot::Sender>, + }, +} + +/// Receipt resolving once stop is durably observed or the task is terminal. +/// +/// Resolves from authoritative status as soon as `stop_requested` is true or +/// the lifecycle is Completed, Failed, or Abandoned. It need not wait for +/// graceful cleanup after stop observation. +#[must_use = "await the receipt to observe durable stop acknowledgement"] +pub struct ManagedStopReceipt { + future: Pin + Send>>, +} + +impl ManagedStopReceipt { + fn new(mut status: ManagedStatusSubscription) -> Self { + let future = Box::pin(async move { + while !{ + let snapshot = status.borrow(); + snapshot.stop_requested + || matches!( + snapshot.lifecycle, + ManagedLifecycle::Completed + | ManagedLifecycle::Failed + | ManagedLifecycle::Abandoned + ) + } { + // A closed sender retains the terminal snapshot. Never wait + // indefinitely if the channel closes. + if status.changed().await.is_err() { + break; + } + } + }); + 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) + } +} + +fn build_task( + config: ManagedClientConfig, + targets: Vec, +) -> Result<(ManagedClientTask, ManagedClientHandle), ManagedConfigError> { + if config.event_capacity == 0 { + return Err(ManagedConfigError::ZeroEventCapacity); + } + if config.event_capacity > MAX_BROADCAST_CHANNEL_CAPACITY { + return Err(ManagedConfigError::EventCapacityTooLarge { + configured: config.event_capacity, + maximum: MAX_BROADCAST_CHANNEL_CAPACITY, + }); + } + if config.command_capacity == 0 { + return Err(ManagedConfigError::ZeroCommandCapacity); + } + if config.command_capacity > tokio::sync::Semaphore::MAX_PERMITS { + return Err(ManagedConfigError::CommandCapacityTooLarge { + configured: config.command_capacity, + maximum: tokio::sync::Semaphore::MAX_PERMITS, + }); + } + 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.auth = target.auth.resolve(&config.client.auth); + validate_target_config(&client_config).map_err(|source| { + ManagedConfigError::InvalidTarget { + id: target.id.clone(), + source, + } + })?; + runtimes.push(TargetRuntime { + instance: TargetInstance { + id: target.id.clone(), + generation, + }, + config: target.clone(), + membership: TargetMembership::Desired, + server_addr: Arc::from(target.server_addr), + remote: None, + counters: TargetCounters::default(), + phase: Some(TargetPhase::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 stop = Arc::new(StopSignal { + requested: AtomicBool::new(false), + update_admission: AtomicU8::new(UpdateAdmission::Open as u8), + wake: wake_sender, + }); + let history = OutcomeHistory::new(config.outcome_history_limit); + let initial_targets = runtimes + .iter() + .map(|target| ManagedTargetStatus { + target: target.instance.clone(), + desired: target.membership.is_desired(), + lifecycle: ManagedTargetLifecycle::Pending, + outcome: None, + server_addr: Arc::clone(&target.server_addr), + remote: None, + }) + .collect::>(); + let phase = TaskPhase::NotStarted; + let initial = Arc::new(ManagedStatus { + lifecycle: phase.lifecycle(), + stop_requested: false, + applied_command_sequence: 0, + 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: phase.final_outcome(), + }); + let (status_sender, status_receiver) = watch::channel(initial); + let (command_sender, command_receiver) = mpsc::channel(config.command_capacity); + let max_live_target_generations = config.max_live_target_generations; + let resources = TaskResources { + status: status_sender, + events: Some(event_sender), + stop: Arc::clone(&stop), + }; + let task = ManagedClientTask { + #[cfg(test)] + timeout_inspections: std::cell::RefCell::new(Vec::new()), + phase, + config, + targets: runtimes, + commands: command_receiver, + history, + resources: Some(resources), + wake: Some(arm_wake(wake_receiver)), + timer: None, + cursor: 0, + timeout_cursor: 0, + timeout_backlog: VecDeque::with_capacity(TIMEOUT_WORK_BUDGET), + scan_remaining: 0, + send_cursor: 0, + burst_remaining: 0, + stagger_remaining: 0, + send_gate: None, + last_stagger_send: None, + stop_observed: false, + next_generation, + applied_command_sequence: 0, + }; + let handle = ManagedClientHandle { + stop, + commands: command_sender, + max_live_target_generations, + status: status_receiver, + events: weak_events, + }; + Ok((task, handle)) +} + +pub(super) fn validate_target_config(config: &ClientConfig) -> Result<(), ClientError> { + config.open.validate()?; + SessionMachine::validate_config(config)?; + if let Some(ttl) = config.socket.ttl { + validate_ttl(ttl)?; + } + Ok(()) +} diff --git a/crates/irtt-client/src/managed/task/pacing.rs b/crates/irtt-client/src/managed/task/pacing.rs new file mode 100644 index 0000000..6ae5d6e --- /dev/null +++ b/crates/irtt-client/src/managed/task/pacing.rs @@ -0,0 +1,365 @@ +//! Burst/stagger send selection and committed-send cadence. + +use std::{ + task::{Context, Poll}, + time::{Duration, Instant}, +}; + +use crate::{ + managed::{schedule::instant_abs_diff, ManagedTargetEndReason, ManagedTargetFailurePhase}, + ClientEvent, SendProbeError, +}; + +use super::{ + target::{OpenSessionFailureCleanup, TargetPhase}, + ManagedClientTask, TARGET_WORK_BUDGET, +}; + +impl ManagedClientTask { + pub(super) fn active_count(&self) -> usize { + self.targets + .iter() + .filter(|target| target.is_paced_active()) + .count() + } + + fn active_stagger_spacing(&self) -> Option { + let mut active = 0; + let mut minimum: Option = None; + for target in &self.targets { + let TargetPhase::Active { schedule, .. } = target.phase() else { + continue; + }; + if !target.membership.is_desired() { + continue; + } + active += 1; + minimum = minimum.into_iter().chain(Some(schedule.interval())).min(); + } + minimum.map(|interval| stagger_spacing(interval, active)) + } + + pub(super) 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)); + } + + pub(super) fn stagger_target_removed(&mut self, now: Instant) { + self.send_gate = self.send_gate.filter(|gate| *gate > now); + } + + /// The next stagger gate after a send was accepted at `accepted_at`. + /// + /// The gate advances on its **own** cadence rather than being re-anchored + /// on every observed accept time. Anchoring on `accepted_at` folds each + /// send's wakeup latency into the next slot and then into every slot after + /// it, so the gate — and with it the whole group's send cadence — drifts + /// later without bound while the per-target schedules stay on their + /// absolute grid. The drift is not hypothetical: it shows up as a + /// monotonically growing `EchoSent::timer_error` and, once it exceeds one + /// interval, as probe slots the schedule then skips. + /// + /// `accepted_at` is still the anchor in the two cases where continuing the + /// old cadence would be wrong: there is no previous gate, or the pacer has + /// fallen a full slot behind it. Re-anchoring there is what keeps a gate + /// stranded in the past from releasing a burst of catch-up sends once the + /// group starts moving again. + fn next_stagger_gate(&self, spacing: Duration, accepted_at: Instant) -> Option { + let base = match self.send_gate { + Some(gate) if gate <= accepted_at && accepted_at.duration_since(gate) < spacing => gate, + _ => accepted_at, + }; + base.checked_add(spacing) + } + + fn record_stagger_acceptance( + &mut self, + result: SendResult, + stagger_spacing: Option, + accepted_at: Instant, + ) { + if !result.accepted() { + return; + } + let Some(spacing) = stagger_spacing else { + return; + }; + self.last_stagger_send = Some(accepted_at); + self.send_gate = self.next_stagger_gate(spacing, accepted_at); + } + + fn poll_one_send( + &mut self, + index: usize, + cx: &mut Context<'_>, + now: Instant, + stagger_spacing: Option, + ) -> SendResult { + if !self.targets[index].membership.is_desired() { + if let TargetPhase::Active { send_waiting, .. } = self.targets[index].phase_mut() { + *send_waiting = false; + } + return SendResult::NotAttempted; + } + let phase = self.targets[index].take_phase(); + let TargetPhase::Active { + mut client, + mut schedule, + mut send_waiting, + } = phase + else { + self.targets[index].restore_phase(phase); + return SendResult::NotAttempted; + }; + // A burst pass may span several slots; validate this target against + // fresh time rather than the pass's earlier scheduling snapshot. + let now = now.max(Instant::now()); + if client + .next_probe_timeout_deadline() + .is_some_and(|deadline| deadline <= now) + { + send_waiting = false; + self.targets[index].restore_phase(TargetPhase::Active { + client, + schedule, + send_waiting, + }); + return SendResult::NotAttempted; + } + if !schedule.permit_probe_at(now) { + send_waiting = false; + let pending = client.has_pending_probes(); + if !pending { + self.targets[index].sync_packets_sent(&client); + } + self.targets[index].restore_phase(TargetPhase::Active { + client, + schedule, + send_waiting, + }); + if pending { + return SendResult::NotAttempted; + } + self.begin_drain(index, ManagedTargetEndReason::TestComplete, now); + return SendResult::Ready { accepted: false }; + } + if schedule + .next_send_deadline() + .is_none_or(|deadline| deadline > now) + { + send_waiting = false; + self.targets[index].restore_phase(TargetPhase::Active { + client, + schedule, + send_waiting, + }); + return SendResult::NotAttempted; + } + let commit = + match schedule.preflight_managed_commit(schedule.next_send_deadline().unwrap(), now) { + Ok(commit) => commit, + Err(error) => { + self.targets[index].restore_phase(TargetPhase::Active { + client, + schedule, + send_waiting, + }); + self.begin_open_session_failure( + index, + ManagedTargetFailurePhase::Timing, + error, + now, + OpenSessionFailureCleanup::Drain, + ); + return SendResult::Failed { accepted: false }; + } + }; + let scheduled_at = commit.scheduled_at; + let result = client.poll_send_probe(cx); + let receipt = match &result { + Poll::Ready(Ok(receipt)) + | Poll::Ready(Err(SendProbeError::AfterCommit { receipt, .. })) => Some(*receipt), + Poll::Pending | Poll::Ready(Err(SendProbeError::NotCommitted(_))) => None, + }; + let accepted = receipt.is_some(); + 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()); + if accepted { + schedule.commit(commit); + } + let schedule_error = if accepted { + schedule.skip_missed_slots_at(Instant::now()).err() + } else { + None + }; + self.targets[index].sync_packets_sent(&client); + send_waiting = result.is_pending(); + self.targets[index].restore_phase(TargetPhase::Active { + client, + schedule, + send_waiting, + }); + // A committed send must reach the lifecycle/accounting stream even + // when post-send processing fails. Publish before failure cleanup. + let sent_event = receipt.map(|receipt| { + let mut event = ClientEvent::from(receipt); + if let ClientEvent::EchoSent { + scheduled_at: intended, + timer_error, + .. + } = &mut event + { + *intended = Some(scheduled_at); + *timer_error = Some(instant_abs_diff(receipt.sent_at.mono, scheduled_at)); + } + event + }); + match result { + Poll::Pending => SendResult::Pending, + Poll::Ready(Ok(_)) => { + self.publish_client_events(index, vec![sent_event.unwrap()]); + if let Some(error) = schedule_error { + self.begin_open_session_failure( + index, + ManagedTargetFailurePhase::Timing, + error, + now, + OpenSessionFailureCleanup::Drain, + ); + SendResult::Failed { accepted } + } else { + SendResult::Ready { accepted } + } + } + Poll::Ready(Err(error)) => { + if let Some(event) = sent_event { + self.publish_client_events(index, vec![event]); + } + let error = match error { + SendProbeError::NotCommitted(source) => source, + SendProbeError::AfterCommit { source, .. } => *source, + }; + self.begin_open_session_failure( + index, + ManagedTargetFailurePhase::Sending, + error, + now, + OpenSessionFailureCleanup::Drain, + ); + SendResult::Failed { accepted } + } + } + } + + pub(super) 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 !self.targets[index].is_paced_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 + } + + pub(super) 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 + } +} + +// A release gate delays a send opportunity; it cannot create one. +pub(super) fn gated_send_deadline( + cadence: Option, + gate: Option, +) -> Option { + cadence.map(|deadline| gate.map_or(deadline, |gate| deadline.max(gate))) +} + +#[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 } + ) + } +} + +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) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn release_gate_only_delays_an_existing_send_deadline() { + let earlier = Instant::now(); + let later = earlier + Duration::from_secs(1); + for (cadence, gate, expected) in [ + (None, None, None), + (None, Some(earlier), None), + (None, Some(later), None), + (Some(earlier), None, Some(earlier)), + (Some(later), Some(earlier), Some(later)), + (Some(earlier), Some(earlier), Some(earlier)), + (Some(earlier), Some(later), Some(later)), + ] { + assert_eq!(gated_send_deadline(cadence, gate), expected); + } + } +} diff --git a/crates/irtt-client/src/managed/task/target.rs b/crates/irtt-client/src/managed/task/target.rs new file mode 100644 index 0000000..0a71ef4 --- /dev/null +++ b/crates/irtt-client/src/managed/task/target.rs @@ -0,0 +1,780 @@ +//! Target-generation state, phase transitions, accounting, and completion. + +use std::{ + future::Future, + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::Instant, +}; + +use crate::{ + async_client::{AsyncClient, AsyncOpenState}, + managed::schedule::ProbeSchedule, + managed::{ + classify_client_error, ManagedEvent, ManagedTargetConfig, ManagedTargetEndReason, + ManagedTargetFailure, ManagedTargetFailureKind, ManagedTargetFailurePhase, + ManagedTargetLifecycle, ManagedTargetOutcome, TargetInstance, + }, + ClientConfig, ClientError, ClientEvent, OpenOutcome, +}; + +use super::{ManagedClientTask, TARGET_WORK_BUDGET}; + +const POST_DEADLINE_RECEIVE_BUDGET: usize = TARGET_WORK_BUDGET; + +type ConnectFuture = + Pin> + Send + 'static>>; + +pub(super) enum TargetPhase { + Pending { + client_config: ClientConfig, + }, + Connecting { + future: ConnectFuture, + }, + Opening { + client: AsyncClient, + open: Box, + }, + Active { + client: AsyncClient, + schedule: ProbeSchedule, + send_waiting: bool, + }, + 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 state always owns its finalized authoritative outcome. + Terminal { + outcome: Arc, + }, +} + +impl TargetPhase { + pub(super) 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, + } + } + + pub(super) fn client(&self) -> &AsyncClient { + match self { + Self::Opening { client, .. } + | Self::Active { client, .. } + | Self::Draining { client, .. } + | Self::Closing { client, .. } => client, + _ => panic!("managed target has no client in this phase"), + } + } + + fn into_client(self) -> AsyncClient { + match self { + Self::Opening { client, .. } + | Self::Active { client, .. } + | Self::Draining { client, .. } + | Self::Closing { client, .. } => client, + _ => panic!("managed target has no client in this phase"), + } + } +} + +#[derive(Default)] +pub(super) struct TargetCounters { + packets_sent: u64, + /// Unique echo replies that arrived in order and before their probe timed + /// out. + /// + /// A reply that arrives behind an already-received sequence, or after its + /// own probe timed out, is counted by `late` and deliberately not here. + /// These counters are durable, unlike the lossy managed presentation + /// events. + replies_received: u64, + duplicates: u64, + late: u64, + warning_events: u64, +} + +impl TargetCounters { + /// Account for one client event. + /// + /// Events that carry no durable target-local count are ignored. Probe + /// sends are not counted here: `packets_sent` is authoritative in the + /// underlying client and is synchronized through + /// [`TargetRuntime::sync_packets_sent`]. + fn observe(&mut self, event: &ClientEvent) { + match event { + ClientEvent::EchoReply { .. } => { + self.replies_received = self.replies_received.saturating_add(1); + } + ClientEvent::DuplicateReply { .. } => { + self.duplicates = self.duplicates.saturating_add(1); + } + ClientEvent::LateReply { .. } => { + self.late = self.late.saturating_add(1); + } + ClientEvent::Warning { .. } => { + self.warning_events = self.warning_events.saturating_add(1); + } + ClientEvent::EchoSent { .. } + | ClientEvent::EchoLoss { .. } + | ClientEvent::SessionStarted(_) + | ClientEvent::NoTestCompleted(_) + | ClientEvent::SessionClosed { .. } => {} + } + } +} + +#[derive(Clone, Copy)] +pub(super) enum RetirementReason { + Removed, + Replaced, +} + +impl From for ManagedTargetEndReason { + fn from(reason: RetirementReason) -> Self { + match reason { + RetirementReason::Removed => Self::Removed, + RetirementReason::Replaced => Self::Replaced, + } + } +} + +pub(super) enum TargetMembership { + Desired, + Withdrawn(RetirementReason), +} + +impl TargetMembership { + pub(super) fn is_desired(&self) -> bool { + matches!(self, Self::Desired) + } + + fn retirement_reason(&self) -> Option { + match self { + Self::Desired => None, + Self::Withdrawn(reason) => Some(*reason), + } + } +} + +pub(super) struct TargetRuntime { + pub(super) instance: TargetInstance, + pub(super) config: ManagedTargetConfig, + pub(super) membership: TargetMembership, + pub(super) server_addr: Arc, + pub(super) remote: Option, + pub(super) counters: TargetCounters, + // None only while a synchronous operation owns the phase locally. + pub(super) phase: Option, +} + +impl TargetRuntime { + pub(super) fn phase(&self) -> &TargetPhase { + self.phase + .as_ref() + .expect("managed target phase is extracted") + } + + pub(super) fn phase_mut(&mut self) -> &mut TargetPhase { + self.phase + .as_mut() + .expect("managed target phase is extracted") + } + + pub(super) fn take_phase(&mut self) -> TargetPhase { + self.phase + .take() + .expect("managed target phase is extracted") + } + + pub(super) fn restore_phase(&mut self, phase: TargetPhase) { + assert!( + self.phase.is_none(), + "managed target phase is already installed" + ); + self.phase = Some(phase); + } + + pub(super) fn is_paced_active(&self) -> bool { + let phase = self.phase(); + self.membership.is_desired() && matches!(phase, TargetPhase::Active { .. }) + } + + /// Adopt the underlying client's authoritative sent count. + /// + /// The client owns `packets_sent`; this runtime only mirrors it so a + /// terminal outcome can report it after the client is gone. + pub(super) fn sync_packets_sent(&mut self, client: &AsyncClient) { + self.counters.packets_sent = client.packets_sent(); + } + + pub(super) fn sync_live_packets_sent(&mut self) { + self.counters.packets_sent = self.phase().client().packets_sent(); + } + + /// Build this target's durable outcome. + pub(super) fn outcome( + &self, + end_reason: ManagedTargetEndReason, + cleanup_failure: Option, + ) -> ManagedTargetOutcome { + ManagedTargetOutcome { + target: self.instance.clone(), + server_addr: Arc::clone(&self.server_addr), + remote: self.remote, + end_reason, + packets_sent: self.counters.packets_sent, + replies_received: self.counters.replies_received, + duplicates: self.counters.duplicates, + late: self.counters.late, + warning_events: self.counters.warning_events, + cleanup_failure, + } + } +} + +#[derive(Clone, Copy)] +pub(super) enum OpenSessionFailureCleanup { + Drain, + Close, +} + +impl ManagedClientTask { + pub(super) fn install_target_phase(&mut self, index: usize, phase: TargetPhase) { + let previous = self.targets[index].phase().lifecycle(); + self.targets[index].phase = Some(phase); + self.publish_target_phase(index, previous); + } + + fn publish_target_phase(&mut self, index: usize, previous: ManagedTargetLifecycle) { + let lifecycle = self.targets[index].phase().lifecycle(); + let desired = self.targets[index].membership.is_desired(); + let was_paced_active = desired && previous == ManagedTargetLifecycle::Active; + let paced_active = self.targets[index].is_paced_active(); + let now = Instant::now(); + match (was_paced_active, paced_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, + }); + } + + pub(super) fn publish_client_events(&mut self, index: usize, events: Vec) { + for event in events { + self.targets[index].counters.observe(&event); + 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, + ) { + let outcome = Arc::new(self.targets[index].outcome(end_reason, cleanup_failure)); + // Finalize history before the phase publisher installs Terminal, updates + // pacing and publishes one complete status followed by the lifecycle event. + self.history.record(outcome.as_ref().clone()); + self.install_target_phase( + index, + TargetPhase::Terminal { + outcome: Arc::clone(&outcome), + }, + ); + self.publish_event(ManagedEvent::TargetFinished { 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); + } + + pub(super) fn begin_open_session_failure( + &mut self, + index: usize, + phase: ManagedTargetFailurePhase, + error: ClientError, + now: Instant, + cleanup: OpenSessionFailureCleanup, + ) -> bool { + let primary_end = ManagedTargetEndReason::Failed(classify_client_error(phase, &error)); + let (TargetPhase::Opening { client, .. } | TargetPhase::Active { client, .. }) = + self.targets[index].phase_mut() + else { + unreachable!("open-session failure requires an opening or active target") + }; + client.discard_prepared_probe(); + self.targets[index].sync_live_packets_sent(); + match cleanup { + OpenSessionFailureCleanup::Drain => self.begin_drain(index, primary_end, now), + OpenSessionFailureCleanup::Close => self.begin_close(index, primary_end, None, now), + } + } + + pub(super) fn effective_retirement(&self, index: usize) -> Option { + self.targets[index] + .membership + .retirement_reason() + .map(Into::into) + } + + fn start_connecting(&mut self, index: usize) { + let TargetPhase::Pending { client_config } = self.targets[index].take_phase() else { + unreachable!("connecting starts from a pending target") + }; + let endpoint = Arc::clone(&self.targets[index].server_addr); + let future = Box::pin(async move { AsyncClient::connect(endpoint, client_config).await }); + self.targets[index].restore_phase(TargetPhase::Connecting { future }); + self.publish_target_phase(index, ManagedTargetLifecycle::Pending); + } + + pub(super) fn poll_target(&mut self, index: usize, cx: &mut Context<'_>, now: Instant) -> bool { + let retirement = self.effective_retirement(index); + let stopping = self.phase.is_stopping() || retirement.is_some(); + match self.targets[index].phase_mut() { + TargetPhase::Pending { .. } => { + if stopping { + self.finish_target( + index, + retirement + .clone() + .unwrap_or(ManagedTargetEndReason::Stopped), + None, + ); + false + } else { + self.start_connecting(index); + true + } + } + TargetPhase::Connecting { future } => { + if stopping { + self.finish_target( + index, + retirement + .clone() + .unwrap_or(ManagedTargetEndReason::Stopped), + None, + ); + return false; + } + match future.as_mut().poll(cx) { + Poll::Pending => false, + Poll::Ready(Ok(client)) => { + self.targets[index].remote = Some(client.remote_addr()); + self.install_target_phase( + index, + TargetPhase::Opening { + client, + open: Box::new(AsyncOpenState::new()), + }, + ); + true + } + Poll::Ready(Err(error)) => { + self.fail_target(index, ManagedTargetFailurePhase::Connecting, error); + false + } + } + } + TargetPhase::Opening { client, open } => { + if stopping { + let retirement = retirement + .clone() + .unwrap_or(ManagedTargetEndReason::Stopped); + if !open.has_in_flight_work() { + self.targets[index].sync_live_packets_sent(); + self.finish_target(index, retirement, None); + return false; + } + open.request_stop_after_current_attempt(); + } + match client.poll_open(open, cx) { + Poll::Pending => false, + Poll::Ready(Ok(OpenOutcome::Started(started))) => { + let schedule = match ProbeSchedule::new( + started.at.mono, + &started.negotiation.accepted, + ) { + Ok(schedule) => schedule, + Err(error) => { + return self.begin_open_session_failure( + index, + ManagedTargetFailurePhase::Timing, + error, + now, + OpenSessionFailureCleanup::Close, + ) + } + }; + self.publish_client_events( + index, + vec![ClientEvent::SessionStarted(started)], + ); + if stopping { + self.targets[index].sync_live_packets_sent(); + self.begin_drain( + index, + retirement + .clone() + .unwrap_or(ManagedTargetEndReason::Stopped), + now, + ) + } else { + let previous = self.targets[index].take_phase(); + self.targets[index].restore_phase(TargetPhase::Active { + client: previous.into_client(), + schedule, + send_waiting: false, + }); + self.publish_target_phase(index, ManagedTargetLifecycle::Opening); + true + } + } + Poll::Ready(Ok(OpenOutcome::NoTestCompleted(completed))) => { + self.publish_client_events( + index, + vec![ClientEvent::NoTestCompleted(completed)], + ); + let end_reason = if stopping { + retirement + .clone() + .unwrap_or(ManagedTargetEndReason::Stopped) + } else { + ManagedTargetEndReason::NoTestComplete + }; + self.finish_target(index, end_reason, None); + false + } + Poll::Ready(Err(error)) => { + if stopping { + let TargetPhase::Opening { open, .. } = self.targets[index].phase() + else { + unreachable!() + }; + let cleanup_failure = + (!open.stopped_after_current_attempt()).then(|| { + classify_client_error( + ManagedTargetFailurePhase::Opening, + &error, + ) + }); + self.finish_target( + index, + retirement + .clone() + .unwrap_or(ManagedTargetEndReason::Stopped), + cleanup_failure, + ); + } else { + self.fail_target(index, ManagedTargetFailurePhase::Opening, error); + } + false + } + } + } + TargetPhase::Active { client, .. } => { + if stopping { + client.discard_prepared_probe(); + self.targets[index].sync_live_packets_sent(); + return self.begin_drain( + index, + retirement + .clone() + .unwrap_or(ManagedTargetEndReason::Stopped), + now, + ); + } + if client + .next_probe_timeout_deadline() + .is_some_and(|deadline| deadline <= now) + { + return true; + } + 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, + ManagedTargetFailurePhase::Receiving, + error, + now, + OpenSessionFailureCleanup::Close, + ); + } + }; + if self.targets[index].phase().client().is_peer_closed() { + self.targets[index].sync_live_packets_sent(); + self.finish_target(index, ManagedTargetEndReason::PeerClosed, None); + return false; + } + let TargetPhase::Active { + client, schedule, .. + } = self.targets[index].phase_mut() + else { + unreachable!() + }; + schedule.permit_probe_at(now); + if schedule.is_finished() && !client.has_pending_probes() { + self.targets[index].sync_live_packets_sent(); + return self.begin_drain(index, ManagedTargetEndReason::TestComplete, now); + } + received + } + TargetPhase::Draining { client, .. } => { + if client + .next_probe_timeout_deadline() + .is_some_and(|timeout| timeout <= now) + { + return true; + } + let mut retained_state_changed = false; + 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].sync_live_packets_sent(); + let TargetPhase::Draining { + primary_end, + cleanup_failure, + .. + } = self.targets[index].phase_mut() + else { + unreachable!() + }; + cleanup_failure.get_or_insert_with(|| { + classify_client_error(ManagedTargetFailurePhase::Receiving, &error) + }); + let primary_end = primary_end.clone(); + let cleanup_failure = cleanup_failure.clone(); + return self.begin_close(index, primary_end, cleanup_failure, now); + } + }; + if self.targets[index].phase().client().is_peer_closed() { + self.targets[index].sync_live_packets_sent(); + let TargetPhase::Draining { + cleanup_failure, .. + } = self.targets[index].phase() + else { + unreachable!() + }; + self.finish_target( + index, + ManagedTargetEndReason::PeerClosed, + cleanup_failure.clone(), + ); + return false; + } + if retained_state_changed { + let TargetPhase::Draining { + client, + drain_started_at, + .. + } = self.targets[index].phase() + else { + unreachable!() + }; + let candidate = self.drain_deadline(client, *drain_started_at); + let TargetPhase::Draining { + deadline, + primary_end, + cleanup_failure, + .. + } = self.targets[index].phase_mut() + else { + unreachable!() + }; + match candidate { + Some(candidate) => *deadline = (*deadline).min(candidate), + None => { + cleanup_failure.get_or_insert_with(duration_overflow_failure); + let primary_end = primary_end.clone(); + let cleanup_failure = cleanup_failure.clone(); + return self.begin_close(index, primary_end, cleanup_failure, now); + } + } + } + let TargetPhase::Draining { + deadline, + primary_end, + cleanup_failure, + post_deadline_receives_remaining, + .. + } = self.targets[index].phase_mut() + else { + unreachable!() + }; + if now >= *deadline { + if received && *post_deadline_receives_remaining > 1 { + *post_deadline_receives_remaining -= 1; + true + } else { + let primary_end = primary_end.clone(); + let cleanup_failure = cleanup_failure.clone(); + self.begin_close(index, primary_end, cleanup_failure, now) + } + } else { + received + } + } + TargetPhase::Closing { client, .. } => { + let result = client.poll_close(cx); + let TargetPhase::Closing { + deadline, + primary_end, + cleanup_failure, + .. + } = self.targets[index].phase() + else { + unreachable!() + }; + let deadline = *deadline; + let primary_end = primary_end.clone(); + let cleanup_failure = cleanup_failure.clone(); + match result { + Poll::Pending => { + if now >= deadline { + self.targets[index].sync_live_packets_sent(); + self.finish_target( + index, + primary_end, + cleanup_failure.or_else(|| Some(close_timeout_failure())), + ); + } + false + } + Poll::Ready(Ok(events)) => { + self.targets[index].sync_live_packets_sent(); + self.publish_client_events(index, events); + self.finish_target(index, primary_end, cleanup_failure); + false + } + Poll::Ready(Err(error)) => { + self.targets[index].sync_live_packets_sent(); + let cleanup = + classify_client_error(ManagedTargetFailurePhase::Closing, &error); + self.finish_target(index, primary_end, cleanup_failure.or(Some(cleanup))); + false + } + } + } + TargetPhase::Terminal { .. } => false, + } + } + + pub(super) fn begin_drain( + &mut self, + index: usize, + primary_end: ManagedTargetEndReason, + _now: Instant, + ) -> bool { + let drain_started_at = Instant::now(); + let Some(deadline) = + self.drain_deadline(self.targets[index].phase().client(), drain_started_at) + else { + return self.begin_close( + index, + primary_end, + Some(duration_overflow_failure()), + drain_started_at, + ); + }; + let previous = self.targets[index].take_phase(); + let lifecycle = previous.lifecycle(); + self.targets[index].restore_phase(TargetPhase::Draining { + client: previous.into_client(), + drain_started_at, + deadline, + primary_end, + cleanup_failure: None, + post_deadline_receives_remaining: POST_DEADLINE_RECEIVE_BUDGET, + }); + self.publish_target_phase(index, lifecycle); + true + } + + pub(super) 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) + } + + pub(super) fn begin_close( + &mut self, + index: usize, + 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 + }); + let previous = self.targets[index].take_phase(); + let lifecycle = previous.lifecycle(); + self.targets[index].restore_phase(TargetPhase::Closing { + client: previous.into_client(), + deadline, + primary_end, + cleanup_failure, + }); + self.publish_target_phase(index, lifecycle); + true + } +} + +pub(super) 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"), + } +} diff --git a/crates/irtt-client/src/managed/task/timeouts.rs b/crates/irtt-client/src/managed/task/timeouts.rs new file mode 100644 index 0000000..dd15bb4 --- /dev/null +++ b/crates/irtt-client/src/managed/task/timeouts.rs @@ -0,0 +1,226 @@ +//! Bounded timeout discovery and processing across target generations. + +use std::time::Instant; + +use crate::managed::{classify_client_error, ManagedTargetFailurePhase}; + +use super::{ + target::{duration_overflow_failure, OpenSessionFailureCleanup, TargetPhase}, + ManagedClientTask, +}; + +pub(super) const TIMEOUT_WORK_BUDGET: usize = 128; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct TimeoutBacklogEntry { + index: usize, + generation: u64, +} + +#[derive(Clone, Copy)] +struct TimeoutStep { + more_due: bool, +} + +impl TimeoutStep { + const NONE: Self = Self { more_due: false }; +} + +impl ManagedClientTask { + pub(super) fn poll_timeout_pass(&mut self, now: Instant) -> bool { + let target_len = self.targets.len(); + if target_len == 0 { + return false; + } + + let mut transitions_remaining = TIMEOUT_WORK_BUDGET; + for _ in 0..target_len.min(TIMEOUT_WORK_BUDGET) { + if transitions_remaining == 0 { + break; + } + let index = self.timeout_cursor; + self.timeout_cursor = (self.timeout_cursor + 1) % target_len; + if self.target_has_due_timeout(index, now) { + let generation = self.targets[index].instance.generation; + let step = self.poll_target_timeout(index, now); + transitions_remaining -= 1; + if step.more_due { + self.enqueue_timeout_backlog(index, generation); + } + } + } + + while transitions_remaining > 0 { + let Some(entry) = self.timeout_backlog.pop_front() else { + break; + }; + let valid = self + .targets + .get(entry.index) + .is_some_and(|target| target.instance.generation == entry.generation) + && self.target_has_due_timeout(entry.index, now); + if !valid { + continue; + } + + let step = self.poll_target_timeout(entry.index, now); + transitions_remaining -= 1; + if step.more_due { + self.enqueue_timeout_backlog(entry.index, entry.generation); + } + } + + !self.timeout_backlog.is_empty() + } + + fn enqueue_timeout_backlog(&mut self, index: usize, generation: u64) { + if self.timeout_backlog.len() == TIMEOUT_WORK_BUDGET + || self + .timeout_backlog + .iter() + .any(|entry| entry.index == index && entry.generation == generation) + { + return; + } + self.timeout_backlog + .push_back(TimeoutBacklogEntry { index, generation }); + } + + fn target_has_due_timeout(&self, index: usize, now: Instant) -> bool { + #[cfg(test)] + self.timeout_inspections.borrow_mut().push(index); + match self.targets[index].phase() { + TargetPhase::Active { client, .. } => { + !self.phase.is_stopping() + && self.effective_retirement(index).is_none() + && client + .next_probe_timeout_deadline() + .is_some_and(|deadline| deadline <= now) + } + TargetPhase::Draining { client, .. } => client + .next_probe_timeout_deadline() + .is_some_and(|timeout| timeout <= now), + _ => false, + } + } + + fn poll_target_timeout(&mut self, index: usize, now: Instant) -> TimeoutStep { + match self.targets[index].phase_mut() { + TargetPhase::Active { client, .. } => match client.poll_timeouts_bounded_at(now, 1) { + Ok(batch) => { + self.publish_client_events(index, batch.events); + TimeoutStep { + more_due: batch.more_due, + } + } + Err(error) => { + self.begin_open_session_failure( + index, + ManagedTargetFailurePhase::Timing, + error, + now, + OpenSessionFailureCleanup::Close, + ); + TimeoutStep::NONE + } + }, + TargetPhase::Draining { client, .. } => match client.poll_timeouts_bounded_at(now, 1) { + Ok(batch) => { + self.publish_client_events(index, batch.events); + let TargetPhase::Draining { + client, + drain_started_at, + .. + } = self.targets[index].phase() + else { + unreachable!() + }; + let candidate = self.drain_deadline(client, *drain_started_at); + let TargetPhase::Draining { + deadline, + primary_end, + cleanup_failure, + .. + } = self.targets[index].phase_mut() + else { + unreachable!() + }; + match candidate { + Some(candidate) => *deadline = (*deadline).min(candidate), + None => { + cleanup_failure.get_or_insert_with(duration_overflow_failure); + let primary_end = primary_end.clone(); + let cleanup_failure = cleanup_failure.clone(); + self.begin_close(index, primary_end, cleanup_failure, now); + return TimeoutStep::NONE; + } + } + TimeoutStep { + more_due: batch.more_due, + } + } + Err(error) => { + self.targets[index].sync_live_packets_sent(); + let TargetPhase::Draining { + primary_end, + cleanup_failure, + .. + } = self.targets[index].phase_mut() + else { + unreachable!() + }; + cleanup_failure.get_or_insert_with(|| { + classify_client_error(ManagedTargetFailurePhase::Timing, &error) + }); + let primary_end = primary_end.clone(); + let cleanup_failure = cleanup_failure.clone(); + self.begin_close(index, primary_end, cleanup_failure, now); + TimeoutStep::NONE + } + }, + _ => TimeoutStep::NONE, + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use crate::managed::{ManagedClientConfig, ManagedTargetConfig}; + + use super::{super::ManagedClient, Instant, TIMEOUT_WORK_BUDGET}; + + // Poll latency through public events cannot distinguish an O(n) discovery + // scan from a bounded scan reliably. This tiny inspection trace measures + // work directly, including empty targets, and checks eventual coverage + // without asserting a cursor value or a particular visitation order. + #[test] + fn timeout_discovery_has_bounded_work_and_eventually_visits_every_target() { + let count = TIMEOUT_WORK_BUDGET * 2 + 1; + let targets = (0..count) + .map(|i| ManagedTargetConfig::new(i.to_string(), "127.0.0.1:2112")) + .collect(); + let (mut task, _handle) = ManagedClient::task( + ManagedClientConfig { + max_live_target_generations: count, + ..ManagedClientConfig::default() + }, + targets, + ) + .unwrap(); + let mut visited = HashSet::new(); + for _ in 0..count.div_ceil(TIMEOUT_WORK_BUDGET) { + task.timeout_inspections.borrow_mut().clear(); + task.poll_timeout_pass(Instant::now()); + let inspected = task.timeout_inspections.borrow(); + assert!(inspected.len() <= TIMEOUT_WORK_BUDGET); + visited.extend(inspected.iter().copied()); + } + assert_eq!( + visited.len(), + count, + "timeout discovery must reach all targets" + ); + } +} diff --git a/crates/irtt-client/src/managed/task/update.rs b/crates/irtt-client/src/managed/task/update.rs new file mode 100644 index 0000000..f1bd036 --- /dev/null +++ b/crates/irtt-client/src/managed/task/update.rs @@ -0,0 +1,308 @@ +//! Atomic desired-target transactions, generation retirement, and pruning. + +use std::{ + collections::HashSet, + sync::Arc, + task::{Context, Poll}, + time::Instant, +}; + +use crate::{ + managed::{ + ManagedCommandAcknowledgement, ManagedCommandApplyError, ManagedCompletionPolicy, + ManagedEvent, ManagedTargetConfig, ManagedTargetLifecycle, TargetInstance, + }, + ClientConfig, +}; + +use super::{ + control::{validate_target_config, ManagedCommand, UpdateAdmission}, + target::{RetirementReason, TargetCounters, TargetMembership, TargetPhase, TargetRuntime}, + ManagedClientTask, TaskPhase, +}; + +const COMMAND_WORK_BUDGET: usize = 32; + +struct PlannedTarget { + target: ManagedTargetConfig, + client_config: ClientConfig, + generation: u64, +} + +struct PlannedRetirement { + index: usize, + reason: RetirementReason, + synchronous: bool, +} + +struct UpdatePlan { + retirements: Vec, + created: Vec, + next_generation: u64, + next_command_sequence: u64, + prospective_live_count: usize, +} + +fn synchronously_retireable(phase: &TargetPhase) -> bool { + match phase { + TargetPhase::Pending { .. } + | TargetPhase::Connecting { .. } + | TargetPhase::Terminal { .. } => true, + TargetPhase::Opening { open, .. } => !open.has_in_flight_work(), + TargetPhase::Active { .. } | TargetPhase::Draining { .. } | TargetPhase::Closing { .. } => { + false + } + } +} + +fn runtime_satisfies_target(runtime: &TargetRuntime, target: &ManagedTargetConfig) -> bool { + runtime.membership.is_desired() + && !matches!(runtime.phase(), TargetPhase::Terminal { .. }) + && runtime.config == *target +} + +impl ManagedClientTask { + pub(super) fn process_commands(&mut self, cx: &mut Context<'_>) -> bool { + let mut immediate = false; + for _ in 0..COMMAND_WORK_BUDGET { + match self.commands.poll_recv(cx) { + Poll::Ready(Some(ManagedCommand::UpdateTargets { + targets, + acknowledgement, + })) => { + // This check is the driver-side application linearization point. + // `apply_targets` cannot suspend, so a stop observed here wins over + // this command; once it passes, this command may complete atomically. + let result = if !self.phase.is_running() + || self.resources().stop.update_admission() != UpdateAdmission::Open + { + Err(ManagedCommandApplyError::Stopping) + } else { + self.apply_targets(targets) + }; + let _ = acknowledgement.send(result); + immediate = true; + } + Poll::Ready(None) | Poll::Pending => break, + } + } + immediate + } + + fn apply_targets( + &mut self, + incoming: Vec, + ) -> Result { + let plan = self.plan_targets(incoming)?; + debug_assert!(plan.prospective_live_count <= self.config.max_live_target_generations); + + let now = Instant::now(); + let mut synchronous = HashSet::with_capacity(plan.retirements.len()); + let mut finished = Vec::new(); + for retirement in &plan.retirements { + let target = &mut self.targets[retirement.index]; + let was_paced_active = target.is_paced_active(); + target.membership = TargetMembership::Withdrawn(retirement.reason); + if retirement.synchronous { + synchronous.insert(retirement.index); + if !matches!(target.phase(), TargetPhase::Terminal { .. }) { + let outcome = Arc::new(target.outcome(retirement.reason.into(), None)); + self.history.record(outcome.as_ref().clone()); + finished.push(ManagedEvent::TargetFinished { outcome }); + } + } else if was_paced_active { + // Removal is directional: it may discard an elapsed gate but never + // lengthens an existing future stagger gate. + self.stagger_target_removed(now); + } + } + if !synchronous.is_empty() { + let mut index = 0; + self.targets.retain(|_| { + let keep = !synchronous.contains(&index); + index += 1; + keep + }); + self.rebase_target_cursors(); + } + + let mut created_instances = Vec::with_capacity(plan.created.len()); + for planned in plan.created { + let instance = TargetInstance { + id: planned.target.id.clone(), + generation: planned.generation, + }; + created_instances.push(instance.clone()); + self.targets.push(TargetRuntime { + instance, + server_addr: Arc::from(planned.target.server_addr.clone()), + config: planned.target, + membership: TargetMembership::Desired, + remote: None, + counters: TargetCounters::default(), + phase: Some(TargetPhase::Pending { + client_config: planned.client_config, + }), + }); + } + self.next_generation = plan.next_generation; + self.applied_command_sequence = plan.next_command_sequence; + let stopping = self + .targets + .iter() + .all(|runtime| !runtime.membership.is_desired()) + && self.config.completion == ManagedCompletionPolicy::FinishWhenQuiescent; + if stopping { + self.resources().stop.begin_stopping(); + self.phase = TaskPhase::Stopping; + } + + // A transaction has one externally visible durable point: all runtime and + // history changes above, then this exact status snapshot, then its events and + // acknowledgement. Do not use per-transition publishing helpers here. + let status = self.snapshot(); + self.resources().status.send_replace(Arc::clone(&status)); + if stopping { + self.publish_event(ManagedEvent::Stopping); + } + for event in finished { + self.publish_event(event); + } + for target in created_instances { + self.publish_event(ManagedEvent::TargetStateChanged { + target, + lifecycle: ManagedTargetLifecycle::Pending, + }); + } + Ok(ManagedCommandAcknowledgement { + sequence: self.applied_command_sequence, + status, + }) + } + + fn plan_targets( + &self, + incoming: Vec, + ) -> Result { + let mut ids = HashSet::with_capacity(incoming.len()); + let mut prepared = Vec::with_capacity(incoming.len()); + for target in incoming { + if !ids.insert(target.id.clone()) { + return Err(ManagedCommandApplyError::DuplicateTargetId { id: target.id }); + } + let mut client_config = self.config.client.clone(); + client_config.auth = target.auth.resolve(&self.config.client.auth); + validate_target_config(&client_config).map_err(|source| { + ManagedCommandApplyError::InvalidTarget { + id: target.id.clone(), + source, + } + })?; + prepared.push((target, client_config)); + } + let mut retirements = Vec::new(); + for (index, runtime) in self.targets.iter().enumerate() { + if !runtime.membership.is_desired() + || prepared + .iter() + .any(|(target, _)| runtime_satisfies_target(runtime, target)) + { + continue; + } + let reason = if prepared + .iter() + .any(|(target, _)| target.id == runtime.instance.id) + { + RetirementReason::Replaced + } else { + RetirementReason::Removed + }; + retirements.push(PlannedRetirement { + index, + reason, + synchronous: synchronously_retireable(runtime.phase()), + }); + } + + let mut next_generation = self.next_generation; + let mut created = Vec::new(); + for (target, client_config) in prepared { + if self + .targets + .iter() + .any(|runtime| runtime_satisfies_target(runtime, &target)) + { + continue; + } + let generation = next_generation; + next_generation = next_generation + .checked_add(1) + .ok_or(ManagedCommandApplyError::GenerationExhausted)?; + created.push(PlannedTarget { + target, + client_config, + generation, + }); + } + let next_command_sequence = self + .applied_command_sequence + .checked_add(1) + .ok_or(ManagedCommandApplyError::CommandSequenceExhausted)?; + let synchronous = retirements + .iter() + .filter(|retirement| retirement.synchronous) + .map(|retirement| retirement.index) + .collect::>(); + let prospective_live_count = self + .targets + .iter() + .enumerate() + .filter(|(index, _)| !synchronous.contains(index)) + .count() + .saturating_add(created.len()); + if prospective_live_count > self.config.max_live_target_generations { + return Err(ManagedCommandApplyError::LiveGenerationLimitExceeded { + required: prospective_live_count, + limit: self.config.max_live_target_generations, + }); + } + Ok(UpdatePlan { + retirements, + created, + next_generation, + next_command_sequence, + prospective_live_count, + }) + } + + pub(super) fn prune_undesired_terminal(&mut self) { + if !self.targets.iter().any(|target| { + !target.membership.is_desired() + && matches!(target.phase(), TargetPhase::Terminal { .. }) + }) { + return; + } + self.targets.retain(|target| { + target.membership.is_desired() + || !matches!(target.phase(), TargetPhase::Terminal { .. }) + }); + self.rebase_target_cursors(); + self.replace_status(); + } + + fn rebase_target_cursors(&mut self) { + let len = self.targets.len(); + if len == 0 { + self.cursor = 0; + self.timeout_cursor = 0; + self.send_cursor = 0; + } else { + self.cursor %= len; + self.timeout_cursor %= len; + self.send_cursor %= len; + } + self.scan_remaining = self.scan_remaining.min(len); + self.burst_remaining = self.burst_remaining.min(len); + self.stagger_remaining = self.stagger_remaining.min(len); + } +}