515 lines
19 KiB
Rust
515 lines
19 KiB
Rust
//! 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.address_family = target
|
|
.address_family
|
|
.unwrap_or(config.client.address_family);
|
|
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(())
|
|
}
|