diff --git a/crates/irtt-client/src/managed/model.rs b/crates/irtt-client/src/managed/model.rs index f972c2d..fac057f 100644 --- a/crates/irtt-client/src/managed/model.rs +++ b/crates/irtt-client/src/managed/model.rs @@ -268,8 +268,12 @@ pub enum ManagedConfigError { TooManyTargets { configured: usize, limit: usize }, #[error("managed event capacity must be greater than zero")] ZeroEventCapacity, + #[error("managed event capacity {configured} exceeds Tokio's maximum {maximum}")] + EventCapacityTooLarge { configured: usize, maximum: usize }, #[error("managed command capacity must be greater than zero")] ZeroCommandCapacity, + #[error("managed command capacity {configured} exceeds Tokio's maximum {maximum}")] + CommandCapacityTooLarge { configured: usize, maximum: usize }, #[error("managed live-generation limit must be greater than zero")] ZeroLiveGenerationLimit, #[error("managed final drain {duration:?} cannot be scheduled from the current instant")] @@ -289,6 +293,8 @@ pub enum ManagedConfigError { pub enum ManagedCommandError { #[error("managed target updates are no longer accepted")] Stopping, + #[error("managed target update contains {configured} targets but the limit is {limit}")] + TooManyTargets { configured: usize, limit: usize }, #[error("managed command queue is full")] QueueFull, #[error("managed command receiver is closed")] diff --git a/crates/irtt-client/src/managed/task.rs b/crates/irtt-client/src/managed/task.rs index c68cb40..63b4845 100644 --- a/crates/irtt-client/src/managed/task.rs +++ b/crates/irtt-client/src/managed/task.rs @@ -5,7 +5,7 @@ use std::{ mem, pin::Pin, sync::{ - atomic::{AtomicBool, Ordering}, + atomic::{AtomicBool, AtomicU8, Ordering}, Arc, }, task::{Context, Poll}, @@ -38,6 +38,9 @@ use super::{ const TARGET_WORK_BUDGET: usize = 128; const COMMAND_WORK_BUDGET: usize = 32; 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. #[derive(Debug, Default)] @@ -56,25 +59,64 @@ impl ManagedClient { #[derive(Debug)] struct StopSignal { requested: AtomicBool, - updates_accepted: AtomicBool, + update_admission: AtomicU8, wake: watch::Sender<()>, acknowledged: AtomicBool, acknowledgement: 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.updates_accepted.store(false, Ordering::Release); + self.begin_stopping(); if !self.requested.swap(true, Ordering::AcqRel) { self.wake.send_replace(()); } } - fn updates_accepted(&self) -> bool { - self.updates_accepted.load(Ordering::Acquire) + 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.updates_accepted.store(false, Ordering::Release); + self.update_admission + .store(UpdateAdmission::Closed as u8, Ordering::Release); } fn is_requested(&self) -> bool { @@ -93,6 +135,7 @@ impl StopSignal { pub struct ManagedClientHandle { stop: Arc, commands: mpsc::Sender, + max_live_target_generations: usize, status: watch::Receiver>, events: broadcast::WeakSender, } @@ -130,16 +173,32 @@ impl ManagedClientHandle { &self, targets: Vec, ) -> Result { - if !self.stop.updates_accepted() { - return Err(ManagedCommandError::Stopping); + 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, + }); + } + + // Submission linearizes at the successful `try_send` below. The + // driver checks stop immediately before non-suspending application; + // terminal sealing closes the receiver before draining accepted work. let (acknowledgement, receiver) = oneshot::channel(); match self.commands.try_send(ManagedCommand::UpdateTargets { targets, acknowledgement, }) { Ok(()) => Ok(ManagedCommandReceipt { receiver }), - Err(mpsc::error::TrySendError::Full(_)) => Err(ManagedCommandError::QueueFull), + 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), } } @@ -657,11 +716,15 @@ impl ManagedClientTask { return; } self.stop_observed = true; - self.resources().stop.close_updates(); + self.resources().stop.begin_stopping(); + let lifecycle_transition = + self.state != DriverState::Stopping || self.lifecycle != ManagedLifecycle::Stopping; self.state = DriverState::Stopping; self.lifecycle = ManagedLifecycle::Stopping; self.replace_status(); - self.publish_event(ManagedEvent::Stopping); + if lifecycle_transition { + self.publish_event(ManagedEvent::Stopping); + } self.resources().stop.acknowledge(); self.scan_remaining = self.targets.len(); self.burst_remaining = 0; @@ -683,7 +746,7 @@ impl ManagedClientTask { // `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.state != DriverState::Running - || self.resources().stop.is_requested() + || self.resources().stop.update_admission() != UpdateAdmission::Open { Err(ManagedCommandApplyError::Stopping) } else { @@ -767,7 +830,7 @@ impl ManagedClientTask { let stopping = self.targets.iter().all(|runtime| !runtime.desired) && self.config.completion == ManagedCompletionPolicy::FinishWhenQuiescent; if stopping { - self.resources().stop.close_updates(); + self.resources().stop.begin_stopping(); self.state = DriverState::Stopping; self.lifecycle = ManagedLifecycle::Stopping; } @@ -1677,7 +1740,22 @@ impl ManagedClientTask { } } - fn seal(&mut self, end_reason: ManagedEndReason, failed: bool) -> Poll { + fn seal(&mut self, mut end_reason: ManagedEndReason, failed: bool) -> Poll { + // Sealing is the final stop-observation point. A stop that reaches the + // latch before this terminal linearization becomes durable even when + // quiescence had already entered `Stopping`. + if self.resources().stop.is_requested() { + self.observe_stop(); + } + if self.stop_observed && matches!(end_reason, ManagedEndReason::TargetsComplete) { + end_reason = ManagedEndReason::StopRequested; + } + + // Closing admission before closing and draining the receiver makes every + // accepted command task-owned. Tokio still permits buffered commands + // to be drained after `Receiver::close`, but rejects later sends. + self.resources().stop.close_updates(); + self.commands.close(); while let Ok(ManagedCommand::UpdateTargets { acknowledgement, .. }) = self.commands.try_recv() @@ -1759,11 +1837,14 @@ impl Future for ManagedClientTask { return this.fail_driver(failure); } } - if this.state == DriverState::Running && this.resources().stop.is_requested() { + if this.resources().stop.is_requested() { this.observe_stop(); } let mut immediate = this.process_commands(cx); + if this.resources().stop.is_requested() { + this.observe_stop(); + } let now = Instant::now(); immediate |= this.poll_target_pass(cx, now); this.prune_undesired_terminal(); @@ -1774,6 +1855,10 @@ impl Future for ManagedClientTask { }; } + if this.resources().stop.is_requested() { + this.observe_stop(); + } + if this.all_targets_terminal() { match this.state { DriverState::Stopping => { @@ -1791,7 +1876,7 @@ impl Future for ManagedClientTask { { this.state = DriverState::Stopping; this.lifecycle = ManagedLifecycle::Stopping; - this.resources().stop.close_updates(); + this.resources().stop.begin_stopping(); this.replace_status(); this.publish_event(ManagedEvent::Stopping); return this.seal(ManagedEndReason::TargetsComplete, false); @@ -1816,6 +1901,8 @@ impl Drop for ManagedClientTask { } self.lifecycle = ManagedLifecycle::Abandoned; self.final_outcome = None; + self.resources().stop.close_updates(); + self.commands.close(); self.replace_status(); self.publish_event(ManagedEvent::Abandoned); self.resources().stop.acknowledge(); @@ -1833,9 +1920,21 @@ fn build_task( 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); } @@ -1899,7 +1998,7 @@ fn build_task( let (acknowledgement_sender, _) = watch::channel(()); let stop = Arc::new(StopSignal { requested: AtomicBool::new(false), - updates_accepted: AtomicBool::new(true), + update_admission: AtomicU8::new(UpdateAdmission::Open as u8), wake: wake_sender, acknowledged: AtomicBool::new(false), acknowledgement: acknowledgement_sender, @@ -1937,6 +2036,7 @@ fn build_task( }); 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), @@ -1973,6 +2073,7 @@ fn build_task( let handle = ManagedClientHandle { stop, commands: command_sender, + max_live_target_generations, status: status_receiver, events: weak_events, };