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.
This commit is contained in:
6 files changed
+2220
-2074
No files matched your search
@@ -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<dyn std::error::Error>> {
|
||||
/// 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<ManagedTargetConfig>,
|
||||
) -> 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<StopSignal>,
|
||||
commands: mpsc::Sender<ManagedCommand>,
|
||||
max_live_target_generations: usize,
|
||||
status: watch::Receiver<Arc<ManagedStatus>>,
|
||||
events: broadcast::WeakSender<ManagedEvent>,
|
||||
}
|
||||
|
||||
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<ManagedStatus> {
|
||||
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<ManagedEventSubscription, ManagedSubscribeError> {
|
||||
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<ManagedTargetConfig>,
|
||||
) -> Result<ManagedCommandReceipt, ManagedCommandError> {
|
||||
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<Result<ManagedCommandAcknowledgement, ManagedCommandApplyError>>,
|
||||
}
|
||||
|
||||
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<ManagedCommandAcknowledgement, ManagedCommandApplyError> {
|
||||
self.receiver
|
||||
.blocking_recv()
|
||||
.unwrap_or(Err(ManagedCommandApplyError::AcknowledgementDisconnected))
|
||||
}
|
||||
}
|
||||
|
||||
impl Future for ManagedCommandReceipt {
|
||||
type Output = Result<ManagedCommandAcknowledgement, ManagedCommandApplyError>;
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
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<ManagedTargetConfig>,
|
||||
acknowledgement:
|
||||
oneshot::Sender<Result<ManagedCommandAcknowledgement, ManagedCommandApplyError>>,
|
||||
},
|
||||
}
|
||||
|
||||
/// 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<Box<dyn Future<Output = ()> + 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::Output> {
|
||||
self.future.as_mut().poll(cx)
|
||||
}
|
||||
}
|
||||
|
||||
fn build_task(
|
||||
config: ManagedClientConfig,
|
||||
targets: Vec<ManagedTargetConfig>,
|
||||
) -> 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::<Vec<_>>();
|
||||
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(())
|
||||
}
|
||||
Reference in new issue
Block a user