Nils Andreas Svee
dd68608d56
The counter tests instantiated the private TargetCounters and mirrored its match arms, duplicating outcome-level coverage that already asserts replies_received and late through ManagedTargetOutcome, and coupling a private refactor to persistent test churn. The new counter documentation also defined late replies in timeout-only terms. A reply that arrives behind an already-received sequence is classified late by process_echo_reply even when its own probe has not timed out, so both counters are now described in terms of order and timeout.
2409 lines
86 KiB
Rust
2409 lines
86 KiB
Rust
use std::{
|
|
collections::{HashSet, VecDeque},
|
|
fmt,
|
|
future::Future,
|
|
mem,
|
|
pin::Pin,
|
|
sync::{
|
|
atomic::{AtomicBool, AtomicU8, Ordering},
|
|
Arc,
|
|
},
|
|
task::{Context, Poll},
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use tokio::{
|
|
sync::{broadcast, mpsc, oneshot, watch},
|
|
time::{self, Sleep},
|
|
};
|
|
|
|
use crate::{
|
|
async_client::{AsyncClient, AsyncOpenState},
|
|
session::machine::SessionMachine,
|
|
socket::validate_open_timeouts,
|
|
socket_options::validate_ttl,
|
|
ClientConfig, ClientError, ClientEvent, OpenOutcome,
|
|
};
|
|
|
|
use super::{
|
|
classify_client_error, ManagedClientConfig, ManagedCommandAcknowledgement,
|
|
ManagedCommandApplyError, ManagedCommandError, ManagedCompletionPolicy, ManagedConfigError,
|
|
ManagedDriverFailure, ManagedEndReason, ManagedEvent, ManagedEventSubscription,
|
|
ManagedLifecycle, ManagedOutcome, ManagedPacing, ManagedStatus, ManagedSubscribeError,
|
|
ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetFailure, ManagedTargetFailureKind,
|
|
ManagedTargetFailurePhase, ManagedTargetLifecycle, ManagedTargetOutcome, ManagedTargetStatus,
|
|
TargetInstance,
|
|
};
|
|
|
|
const TARGET_WORK_BUDGET: usize = 128;
|
|
const 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.
|
|
#[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)]
|
|
struct StopSignal {
|
|
requested: 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.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)
|
|
}
|
|
|
|
fn acknowledge(&self) {
|
|
if !self.acknowledged.swap(true, Ordering::AcqRel) {
|
|
self.acknowledgement.send_replace(());
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Cloneable control and observation capability for a managed task.
|
|
#[derive(Clone)]
|
|
pub struct ManagedClientHandle {
|
|
stop: Arc<StopSignal>,
|
|
commands: mpsc::Sender<ManagedCommand>,
|
|
max_live_target_generations: usize,
|
|
status: watch::Receiver<Arc<ManagedStatus>>,
|
|
events: broadcast::WeakSender<ManagedEvent>,
|
|
#[cfg(test)]
|
|
update_submission_hook: Option<Arc<UpdateSubmissionTestHook>>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[derive(Default)]
|
|
struct UpdateSubmissionTestHook {
|
|
state: std::sync::Mutex<UpdateSubmissionTestState>,
|
|
arrived: std::sync::Condvar,
|
|
released: std::sync::Condvar,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[derive(Default)]
|
|
struct UpdateSubmissionTestState {
|
|
armed: bool,
|
|
arrived: bool,
|
|
released: bool,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl UpdateSubmissionTestHook {
|
|
fn arm(&self) {
|
|
self.state
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
|
.armed = true;
|
|
}
|
|
|
|
fn pause_after_admission(&self) {
|
|
let mut state = self
|
|
.state
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
if !state.armed {
|
|
return;
|
|
}
|
|
state.arrived = true;
|
|
self.arrived.notify_all();
|
|
while !state.released {
|
|
state = self
|
|
.released
|
|
.wait(state)
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
}
|
|
}
|
|
|
|
fn wait_until_arrived(&self) {
|
|
let state = self
|
|
.state
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
let (state, timeout) = self
|
|
.arrived
|
|
.wait_timeout_while(state, Duration::from_secs(2), |state| !state.arrived)
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
assert!(state.arrived && !timeout.timed_out());
|
|
}
|
|
|
|
fn release(&self) {
|
|
let mut state = self
|
|
.state
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
state.released = true;
|
|
self.released.notify_all();
|
|
}
|
|
}
|
|
|
|
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 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(Arc::clone(&self.stop))
|
|
}
|
|
|
|
/// Submit one complete desired target set without awaiting queue capacity.
|
|
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,
|
|
});
|
|
}
|
|
|
|
#[cfg(test)]
|
|
if let Some(hook) = &self.update_submission_hook {
|
|
hook.pause_after_admission();
|
|
}
|
|
|
|
// 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 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,
|
|
}
|
|
}
|
|
}
|
|
|
|
enum ManagedCommand {
|
|
UpdateTargets {
|
|
targets: Vec<ManagedTargetConfig>,
|
|
acknowledgement:
|
|
oneshot::Sender<Result<ManagedCommandAcknowledgement, ManagedCommandApplyError>>,
|
|
},
|
|
}
|
|
|
|
/// Receipt resolving once stop is durably observed or the task is terminal.
|
|
#[must_use = "await the receipt to observe durable stop acknowledgement"]
|
|
pub struct ManagedStopReceipt {
|
|
future: Pin<Box<dyn Future<Output = ()> + Send>>,
|
|
}
|
|
|
|
impl ManagedStopReceipt {
|
|
fn new(stop: Arc<StopSignal>) -> Self {
|
|
let mut acknowledgement = stop.acknowledgement.subscribe();
|
|
let future = Box::pin(async move {
|
|
while !stop.acknowledged.load(Ordering::Acquire) {
|
|
let _ = acknowledgement.changed().await;
|
|
}
|
|
});
|
|
Self { future }
|
|
}
|
|
}
|
|
|
|
impl Future for ManagedStopReceipt {
|
|
type Output = ();
|
|
|
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
self.future.as_mut().poll(cx)
|
|
}
|
|
}
|
|
|
|
type ConnectFuture =
|
|
Pin<Box<dyn Future<Output = Result<AsyncClient, ClientError>> + Send + 'static>>;
|
|
type WakeFuture = Pin<Box<dyn Future<Output = Option<watch::Receiver<()>>> + Send + 'static>>;
|
|
#[cfg(test)]
|
|
type EventObservations = Arc<std::sync::Mutex<Vec<(ManagedEvent, Arc<ManagedStatus>)>>>;
|
|
#[cfg(test)]
|
|
type StaggerObservations = Arc<std::sync::Mutex<Vec<(usize, Duration)>>>;
|
|
|
|
fn arm_wake(mut receiver: watch::Receiver<()>) -> WakeFuture {
|
|
Box::pin(async move {
|
|
receiver.changed().await.ok()?;
|
|
Some(receiver)
|
|
})
|
|
}
|
|
|
|
enum TargetState {
|
|
Pending {
|
|
client_config: ClientConfig,
|
|
},
|
|
Connecting {
|
|
future: ConnectFuture,
|
|
},
|
|
Opening {
|
|
client: AsyncClient,
|
|
open: Box<AsyncOpenState>,
|
|
},
|
|
Active {
|
|
client: AsyncClient,
|
|
},
|
|
Draining {
|
|
client: AsyncClient,
|
|
drain_started_at: Instant,
|
|
deadline: Instant,
|
|
primary_end: ManagedTargetEndReason,
|
|
cleanup_failure: Option<ManagedTargetFailure>,
|
|
post_deadline_receives_remaining: usize,
|
|
},
|
|
Closing {
|
|
client: AsyncClient,
|
|
deadline: Instant,
|
|
primary_end: ManagedTargetEndReason,
|
|
cleanup_failure: Option<ManagedTargetFailure>,
|
|
},
|
|
Terminal,
|
|
}
|
|
|
|
impl TargetState {
|
|
fn lifecycle(&self) -> ManagedTargetLifecycle {
|
|
match self {
|
|
Self::Pending { .. } => ManagedTargetLifecycle::Pending,
|
|
Self::Connecting { .. } => ManagedTargetLifecycle::Connecting,
|
|
Self::Opening { .. } => ManagedTargetLifecycle::Opening,
|
|
Self::Active { .. } => ManagedTargetLifecycle::Active,
|
|
Self::Draining { .. } => ManagedTargetLifecycle::Draining,
|
|
Self::Closing { .. } => ManagedTargetLifecycle::Closing,
|
|
Self::Terminal => ManagedTargetLifecycle::Terminal,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct TargetCounters {
|
|
packets_sent: u64,
|
|
/// 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 { .. } => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
struct TargetRuntime {
|
|
instance: TargetInstance,
|
|
config: ManagedTargetConfig,
|
|
desired: bool,
|
|
retirement: Option<ManagedTargetEndReason>,
|
|
server_addr: Arc<str>,
|
|
remote: Option<std::net::SocketAddr>,
|
|
counters: TargetCounters,
|
|
send_waiting: bool,
|
|
active: bool,
|
|
state: TargetState,
|
|
}
|
|
|
|
impl TargetRuntime {
|
|
/// 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.record_packets_sent(client.packets_sent());
|
|
}
|
|
|
|
/// Adopt an already-read authoritative sent count.
|
|
fn record_packets_sent(&mut self, packets_sent: u64) {
|
|
self.counters.packets_sent = packets_sent;
|
|
}
|
|
|
|
/// Account for one client event against this target's durable counters.
|
|
fn observe(&mut self, event: &ClientEvent) {
|
|
self.counters.observe(event);
|
|
}
|
|
|
|
/// Build this target's durable outcome.
|
|
fn outcome(
|
|
&self,
|
|
end_reason: ManagedTargetEndReason,
|
|
cleanup_failure: Option<ManagedTargetFailure>,
|
|
) -> 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: ManagedTargetEndReason,
|
|
synchronous: bool,
|
|
}
|
|
|
|
struct UpdatePlan {
|
|
retirements: Vec<PlannedRetirement>,
|
|
created: Vec<PlannedTarget>,
|
|
next_generation: u64,
|
|
next_command_sequence: u64,
|
|
prospective_live_count: usize,
|
|
}
|
|
|
|
fn synchronously_retireable(state: &TargetState) -> bool {
|
|
match state {
|
|
TargetState::Pending { .. } | TargetState::Connecting { .. } | TargetState::Terminal => {
|
|
true
|
|
}
|
|
TargetState::Opening { open, .. } => !open.has_in_flight_work(),
|
|
TargetState::Active { .. } | TargetState::Draining { .. } | TargetState::Closing { .. } => {
|
|
false
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct OutcomeHistory {
|
|
limit: usize,
|
|
recent: VecDeque<ManagedTargetOutcome>,
|
|
total: u64,
|
|
successful: u64,
|
|
failed: u64,
|
|
peer_closed: u64,
|
|
discarded: u64,
|
|
}
|
|
|
|
impl OutcomeHistory {
|
|
fn new(limit: usize) -> Self {
|
|
Self {
|
|
limit,
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
fn record(&mut self, outcome: ManagedTargetOutcome) {
|
|
self.total = self.total.saturating_add(1);
|
|
match &outcome.end_reason {
|
|
ManagedTargetEndReason::Failed(_) => self.failed = self.failed.saturating_add(1),
|
|
ManagedTargetEndReason::PeerClosed => {
|
|
self.successful = self.successful.saturating_add(1);
|
|
self.peer_closed = self.peer_closed.saturating_add(1);
|
|
}
|
|
ManagedTargetEndReason::TestComplete
|
|
| ManagedTargetEndReason::NoTestComplete
|
|
| ManagedTargetEndReason::Removed
|
|
| ManagedTargetEndReason::Replaced
|
|
| ManagedTargetEndReason::Stopped => {
|
|
self.successful = self.successful.saturating_add(1);
|
|
}
|
|
}
|
|
if self.limit == 0 {
|
|
self.discarded = self.discarded.saturating_add(1);
|
|
return;
|
|
}
|
|
if self.recent.len() == self.limit {
|
|
self.recent.pop_front();
|
|
self.discarded = self.discarded.saturating_add(1);
|
|
}
|
|
self.recent.push_back(outcome);
|
|
}
|
|
|
|
fn recent(&self) -> Arc<[ManagedTargetOutcome]> {
|
|
Arc::from(
|
|
self.recent
|
|
.iter()
|
|
.cloned()
|
|
.collect::<Vec<_>>()
|
|
.into_boxed_slice(),
|
|
)
|
|
}
|
|
|
|
fn outcome(
|
|
&self,
|
|
end_reason: ManagedEndReason,
|
|
applied_command_sequence: u64,
|
|
) -> ManagedOutcome {
|
|
ManagedOutcome {
|
|
end_reason,
|
|
applied_command_sequence,
|
|
total_target_outcomes: self.total,
|
|
successful_target_outcomes: self.successful,
|
|
failed_target_outcomes: self.failed,
|
|
peer_closed_target_outcomes: self.peer_closed,
|
|
discarded_target_outcomes: self.discarded,
|
|
recent_target_outcomes: self.recent(),
|
|
}
|
|
}
|
|
}
|
|
|
|
struct TaskResources {
|
|
status: watch::Sender<Arc<ManagedStatus>>,
|
|
events: Option<broadcast::Sender<ManagedEvent>>,
|
|
stop: Arc<StopSignal>,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Eq, PartialEq)]
|
|
enum DriverState {
|
|
NotStarted,
|
|
Running,
|
|
Stopping,
|
|
Finished,
|
|
}
|
|
|
|
#[derive(Clone, Copy)]
|
|
enum OpenSessionFailureCleanup {
|
|
Drain,
|
|
Close,
|
|
}
|
|
|
|
/// The sole authoritative zero-or-many target driver.
|
|
#[must_use = "ManagedClientTask must be awaited or deliberately dropped"]
|
|
pub struct ManagedClientTask {
|
|
state: DriverState,
|
|
lifecycle: ManagedLifecycle,
|
|
config: ManagedClientConfig,
|
|
targets: Vec<TargetRuntime>,
|
|
commands: mpsc::Receiver<ManagedCommand>,
|
|
history: OutcomeHistory,
|
|
resources: Option<TaskResources>,
|
|
wake: Option<WakeFuture>,
|
|
timer: Option<Pin<Box<Sleep>>>,
|
|
cursor: usize,
|
|
timeout_cursor: usize,
|
|
timeout_backlog: VecDeque<TimeoutBacklogEntry>,
|
|
scan_remaining: usize,
|
|
send_cursor: usize,
|
|
burst_remaining: usize,
|
|
stagger_remaining: usize,
|
|
send_gate: Option<Instant>,
|
|
last_stagger_send: Option<Instant>,
|
|
stop_observed: bool,
|
|
final_outcome: Option<Arc<ManagedOutcome>>,
|
|
next_generation: u64,
|
|
applied_command_sequence: u64,
|
|
#[cfg(test)]
|
|
event_observations: Option<EventObservations>,
|
|
#[cfg(test)]
|
|
stagger_observations: Option<StaggerObservations>,
|
|
#[cfg(test)]
|
|
drain_test_hook: DrainTestHook,
|
|
#[cfg(test)]
|
|
timeout_inspections: usize,
|
|
#[cfg(test)]
|
|
deadline_inspections: usize,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[derive(Default)]
|
|
struct DrainTestHook {
|
|
defer_work_until_deadline: bool,
|
|
fail_receive: bool,
|
|
fail_close: bool,
|
|
}
|
|
|
|
impl ManagedClientTask {
|
|
fn resources(&self) -> &TaskResources {
|
|
self.resources
|
|
.as_ref()
|
|
.expect("managed task resources exist before terminal sealing")
|
|
}
|
|
|
|
fn publish_event(&self, event: ManagedEvent) {
|
|
#[cfg(test)]
|
|
if let Some(observations) = &self.event_observations {
|
|
observations
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
|
.push((event.clone(), Arc::clone(&self.resources().status.borrow())));
|
|
}
|
|
if let Some(events) = &self.resources().events {
|
|
let _ = events.send(event);
|
|
}
|
|
}
|
|
|
|
fn snapshot(&self) -> Arc<ManagedStatus> {
|
|
let mut connecting = 0;
|
|
let mut opening = 0;
|
|
let mut active = 0;
|
|
let mut draining = 0;
|
|
let mut closing = 0;
|
|
let mut terminal = 0;
|
|
let targets = self
|
|
.targets
|
|
.iter()
|
|
.map(|target| {
|
|
let lifecycle = target.state.lifecycle();
|
|
match lifecycle {
|
|
ManagedTargetLifecycle::Pending => {}
|
|
ManagedTargetLifecycle::Connecting => connecting += 1,
|
|
ManagedTargetLifecycle::Opening => opening += 1,
|
|
ManagedTargetLifecycle::Active => active += 1,
|
|
ManagedTargetLifecycle::Draining => draining += 1,
|
|
ManagedTargetLifecycle::Closing => closing += 1,
|
|
ManagedTargetLifecycle::Terminal => terminal += 1,
|
|
}
|
|
ManagedTargetStatus {
|
|
target: target.instance.clone(),
|
|
desired: target.desired,
|
|
lifecycle,
|
|
server_addr: Arc::clone(&target.server_addr),
|
|
remote: target.remote,
|
|
}
|
|
})
|
|
.collect::<Vec<_>>();
|
|
Arc::new(ManagedStatus {
|
|
lifecycle: self.lifecycle,
|
|
stop_requested: self.stop_observed,
|
|
applied_command_sequence: self.applied_command_sequence,
|
|
desired_target_count: self.targets.iter().filter(|target| target.desired).count(),
|
|
connecting_target_count: connecting,
|
|
opening_target_count: opening,
|
|
active_target_count: active,
|
|
draining_target_count: draining,
|
|
closing_target_count: closing,
|
|
terminal_target_count: terminal,
|
|
total_target_outcomes: self.history.total,
|
|
successful_target_outcomes: self.history.successful,
|
|
failed_target_outcomes: self.history.failed,
|
|
peer_closed_target_outcomes: self.history.peer_closed,
|
|
discarded_target_outcomes: self.history.discarded,
|
|
targets: Arc::from(targets.into_boxed_slice()),
|
|
recent_target_outcomes: self.history.recent(),
|
|
final_outcome: self.final_outcome.clone(),
|
|
})
|
|
}
|
|
|
|
fn replace_status(&self) {
|
|
self.resources().status.send_replace(self.snapshot());
|
|
}
|
|
|
|
fn install_target_state(&mut self, index: usize, state: TargetState) {
|
|
let lifecycle = state.lifecycle();
|
|
let active = matches!(state, TargetState::Active { .. });
|
|
let was_active = self.targets[index].active;
|
|
self.targets[index].active = active;
|
|
self.targets[index].state = state;
|
|
let now = Instant::now();
|
|
match (was_active, active) {
|
|
(false, true) => self.stagger_target_added(now),
|
|
(true, false) => self.stagger_target_removed(now),
|
|
_ => {}
|
|
}
|
|
self.replace_status();
|
|
self.publish_event(ManagedEvent::TargetStateChanged {
|
|
target: self.targets[index].instance.clone(),
|
|
lifecycle,
|
|
});
|
|
}
|
|
|
|
fn publish_client_events(&mut self, index: usize, events: Vec<ClientEvent>) {
|
|
for event in events {
|
|
self.targets[index].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<ManagedTargetFailure>,
|
|
) {
|
|
self.install_target_state(index, TargetState::Terminal);
|
|
let outcome = self.targets[index].outcome(end_reason, cleanup_failure);
|
|
self.history.record(outcome.clone());
|
|
self.replace_status();
|
|
self.publish_event(ManagedEvent::TargetFinished {
|
|
outcome: Arc::new(outcome),
|
|
});
|
|
}
|
|
|
|
fn fail_target(&mut self, index: usize, phase: ManagedTargetFailurePhase, error: ClientError) {
|
|
let failure = classify_client_error(phase, &error);
|
|
self.finish_target(index, ManagedTargetEndReason::Failed(failure), None);
|
|
}
|
|
|
|
fn begin_open_session_failure(
|
|
&mut self,
|
|
index: usize,
|
|
mut client: AsyncClient,
|
|
phase: ManagedTargetFailurePhase,
|
|
error: ClientError,
|
|
now: Instant,
|
|
cleanup: OpenSessionFailureCleanup,
|
|
) -> bool {
|
|
let primary_end = ManagedTargetEndReason::Failed(classify_client_error(phase, &error));
|
|
client.discard_prepared_probe();
|
|
self.targets[index].sync_packets_sent(&client);
|
|
match cleanup {
|
|
OpenSessionFailureCleanup::Drain => self.begin_drain(index, client, primary_end, now),
|
|
OpenSessionFailureCleanup::Close => {
|
|
self.begin_close(index, client, primary_end, None, now)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn begin_running(&mut self) -> Result<(), ManagedDriverFailure> {
|
|
tokio::runtime::Handle::try_current().map_err(|_| ManagedDriverFailure::NoTokioRuntime)?;
|
|
self.state = DriverState::Running;
|
|
self.lifecycle = ManagedLifecycle::Running;
|
|
self.replace_status();
|
|
self.publish_event(ManagedEvent::Started);
|
|
Ok(())
|
|
}
|
|
|
|
fn observe_stop(&mut self) {
|
|
if self.stop_observed {
|
|
return;
|
|
}
|
|
self.stop_observed = true;
|
|
self.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();
|
|
if lifecycle_transition {
|
|
self.publish_event(ManagedEvent::Stopping);
|
|
}
|
|
self.resources().stop.acknowledge();
|
|
self.scan_remaining = self.targets.len();
|
|
self.burst_remaining = 0;
|
|
}
|
|
|
|
fn effective_retirement(&self, index: usize) -> Option<ManagedTargetEndReason> {
|
|
self.targets[index].retirement.clone()
|
|
}
|
|
|
|
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.state != DriverState::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<ManagedTargetConfig>,
|
|
) -> Result<ManagedCommandAcknowledgement, ManagedCommandApplyError> {
|
|
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];
|
|
target.desired = false;
|
|
if retirement.synchronous {
|
|
synchronous.insert(retirement.index);
|
|
if !matches!(target.state, TargetState::Terminal) {
|
|
let outcome = target.outcome(retirement.reason.clone(), None);
|
|
self.history.record(outcome.clone());
|
|
finished.push(ManagedEvent::TargetFinished {
|
|
outcome: Arc::new(outcome),
|
|
});
|
|
}
|
|
} else {
|
|
target.retirement = Some(retirement.reason.clone());
|
|
if target.active {
|
|
target.active = false;
|
|
// 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,
|
|
desired: true,
|
|
retirement: None,
|
|
remote: None,
|
|
counters: TargetCounters::default(),
|
|
send_waiting: false,
|
|
active: false,
|
|
state: TargetState::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.desired)
|
|
&& self.config.completion == ManagedCompletionPolicy::FinishWhenQuiescent;
|
|
if stopping {
|
|
self.resources().stop.begin_stopping();
|
|
self.state = DriverState::Stopping;
|
|
self.lifecycle = ManagedLifecycle::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<ManagedTargetConfig>,
|
|
) -> Result<UpdatePlan, ManagedCommandApplyError> {
|
|
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.server_addr.clone_from(&target.server_addr);
|
|
if let Some(auth) = &target.auth {
|
|
client_config.hmac_key.clone_from(&auth.hmac_key);
|
|
}
|
|
validate_target_config(&client_config).map_err(|source| {
|
|
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.desired || prepared.iter().any(|(target, _)| target == &runtime.config) {
|
|
continue;
|
|
}
|
|
let reason = if prepared
|
|
.iter()
|
|
.any(|(target, _)| target.id == runtime.instance.id)
|
|
{
|
|
ManagedTargetEndReason::Replaced
|
|
} else {
|
|
ManagedTargetEndReason::Removed
|
|
};
|
|
retirements.push(PlannedRetirement {
|
|
index,
|
|
reason,
|
|
synchronous: synchronously_retireable(&runtime.state),
|
|
});
|
|
}
|
|
|
|
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.desired && runtime.config == 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::<HashSet<_>>();
|
|
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.desired && matches!(target.state, TargetState::Terminal))
|
|
{
|
|
return;
|
|
}
|
|
self.targets
|
|
.retain(|target| target.desired || !matches!(target.state, TargetState::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, config: ClientConfig) {
|
|
let future = Box::pin(AsyncClient::connect(config));
|
|
self.install_target_state(index, TargetState::Connecting { future });
|
|
}
|
|
|
|
fn poll_target(&mut self, index: usize, cx: &mut Context<'_>, now: Instant) -> bool {
|
|
let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal);
|
|
match state {
|
|
TargetState::Pending { client_config } => {
|
|
if self.state == DriverState::Stopping || self.effective_retirement(index).is_some()
|
|
{
|
|
self.finish_target(
|
|
index,
|
|
self.effective_retirement(index)
|
|
.unwrap_or(ManagedTargetEndReason::Stopped),
|
|
None,
|
|
);
|
|
false
|
|
} else {
|
|
self.start_connecting(index, client_config);
|
|
true
|
|
}
|
|
}
|
|
TargetState::Connecting { mut future } => {
|
|
if self.state == DriverState::Stopping || self.effective_retirement(index).is_some()
|
|
{
|
|
self.finish_target(
|
|
index,
|
|
self.effective_retirement(index)
|
|
.unwrap_or(ManagedTargetEndReason::Stopped),
|
|
None,
|
|
);
|
|
return false;
|
|
}
|
|
match future.as_mut().poll(cx) {
|
|
Poll::Pending => {
|
|
self.targets[index].state = TargetState::Connecting { future };
|
|
false
|
|
}
|
|
Poll::Ready(Ok(client)) => {
|
|
self.targets[index].remote = Some(client.remote_addr());
|
|
self.install_target_state(
|
|
index,
|
|
TargetState::Opening {
|
|
client,
|
|
open: Box::new(AsyncOpenState::new()),
|
|
},
|
|
);
|
|
true
|
|
}
|
|
Poll::Ready(Err(error)) => {
|
|
self.fail_target(index, ManagedTargetFailurePhase::Connecting, error);
|
|
false
|
|
}
|
|
}
|
|
}
|
|
TargetState::Opening {
|
|
mut client,
|
|
mut open,
|
|
} => {
|
|
if self.state == DriverState::Stopping || self.effective_retirement(index).is_some()
|
|
{
|
|
let retirement = self
|
|
.effective_retirement(index)
|
|
.unwrap_or(ManagedTargetEndReason::Stopped);
|
|
if !open.has_in_flight_work() {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
self.finish_target(index, retirement, None);
|
|
return false;
|
|
}
|
|
open.request_stop_after_current_attempt();
|
|
}
|
|
match client.poll_open(&mut open, cx) {
|
|
Poll::Pending => {
|
|
self.targets[index].state = TargetState::Opening { client, open };
|
|
false
|
|
}
|
|
Poll::Ready(Ok(OpenOutcome::Started { event, .. })) => {
|
|
self.publish_client_events(index, vec![event]);
|
|
if self.state == DriverState::Stopping
|
|
|| self.effective_retirement(index).is_some()
|
|
{
|
|
self.targets[index].sync_packets_sent(&client);
|
|
self.begin_drain(
|
|
index,
|
|
client,
|
|
self.effective_retirement(index)
|
|
.unwrap_or(ManagedTargetEndReason::Stopped),
|
|
now,
|
|
)
|
|
} else {
|
|
self.install_target_state(index, TargetState::Active { client });
|
|
true
|
|
}
|
|
}
|
|
Poll::Ready(Ok(OpenOutcome::NoTestCompleted { event, .. })) => {
|
|
self.publish_client_events(index, vec![event]);
|
|
let end_reason = if self.state == DriverState::Stopping
|
|
|| self.effective_retirement(index).is_some()
|
|
{
|
|
self.effective_retirement(index)
|
|
.unwrap_or(ManagedTargetEndReason::Stopped)
|
|
} else {
|
|
ManagedTargetEndReason::NoTestComplete
|
|
};
|
|
self.finish_target(index, end_reason, None);
|
|
false
|
|
}
|
|
Poll::Ready(Err(error)) => {
|
|
if self.state == DriverState::Stopping
|
|
|| self.effective_retirement(index).is_some()
|
|
{
|
|
let cleanup_failure =
|
|
(!open.stopped_after_current_attempt()).then(|| {
|
|
classify_client_error(
|
|
ManagedTargetFailurePhase::Opening,
|
|
&error,
|
|
)
|
|
});
|
|
self.finish_target(
|
|
index,
|
|
self.effective_retirement(index)
|
|
.unwrap_or(ManagedTargetEndReason::Stopped),
|
|
cleanup_failure,
|
|
);
|
|
} else {
|
|
self.fail_target(index, ManagedTargetFailurePhase::Opening, error);
|
|
}
|
|
false
|
|
}
|
|
}
|
|
}
|
|
TargetState::Active { mut client } => {
|
|
if self.state == DriverState::Stopping || self.effective_retirement(index).is_some()
|
|
{
|
|
client.discard_prepared_probe();
|
|
self.targets[index].sync_packets_sent(&client);
|
|
return self.begin_drain(
|
|
index,
|
|
client,
|
|
self.effective_retirement(index)
|
|
.unwrap_or(ManagedTargetEndReason::Stopped),
|
|
now,
|
|
);
|
|
}
|
|
if client
|
|
.next_probe_timeout_deadline()
|
|
.is_some_and(|deadline| deadline <= now)
|
|
{
|
|
self.targets[index].state = TargetState::Active { client };
|
|
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,
|
|
client,
|
|
ManagedTargetFailurePhase::Receiving,
|
|
error,
|
|
now,
|
|
OpenSessionFailureCleanup::Close,
|
|
);
|
|
}
|
|
};
|
|
if client.is_peer_closed() {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
self.finish_target(index, ManagedTargetEndReason::PeerClosed, None);
|
|
return false;
|
|
}
|
|
if client.is_run_complete() {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
return self.begin_drain(
|
|
index,
|
|
client,
|
|
ManagedTargetEndReason::TestComplete,
|
|
now,
|
|
);
|
|
}
|
|
self.targets[index].state = TargetState::Active { client };
|
|
received
|
|
}
|
|
TargetState::Draining {
|
|
mut client,
|
|
drain_started_at,
|
|
deadline,
|
|
primary_end,
|
|
mut cleanup_failure,
|
|
mut post_deadline_receives_remaining,
|
|
} => {
|
|
#[cfg(test)]
|
|
let defer_work = self.drain_test_hook.defer_work_until_deadline && now < deadline;
|
|
#[cfg(not(test))]
|
|
let defer_work = false;
|
|
if defer_work {
|
|
self.targets[index].state = TargetState::Draining {
|
|
client,
|
|
drain_started_at,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure,
|
|
post_deadline_receives_remaining,
|
|
};
|
|
return false;
|
|
}
|
|
|
|
if client
|
|
.next_probe_timeout_deadline()
|
|
.is_some_and(|timeout| timeout <= now)
|
|
{
|
|
self.targets[index].state = TargetState::Draining {
|
|
client,
|
|
drain_started_at,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure,
|
|
post_deadline_receives_remaining,
|
|
};
|
|
return true;
|
|
}
|
|
|
|
#[cfg(test)]
|
|
let injected_receive_failure = self.take_drain_receive_failure();
|
|
#[cfg(not(test))]
|
|
let injected_receive_failure = None;
|
|
let mut retained_state_changed = false;
|
|
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_packets_sent(&client);
|
|
cleanup_failure.get_or_insert_with(|| {
|
|
classify_client_error(ManagedTargetFailurePhase::Receiving, &error)
|
|
});
|
|
return self.begin_close(index, client, primary_end, cleanup_failure, now);
|
|
}
|
|
};
|
|
if let Some(error) = injected_receive_failure {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
cleanup_failure.get_or_insert_with(|| {
|
|
classify_client_error(ManagedTargetFailurePhase::Receiving, &error)
|
|
});
|
|
return self.begin_close(index, client, primary_end, cleanup_failure, now);
|
|
}
|
|
if client.is_peer_closed() {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
self.finish_target(index, ManagedTargetEndReason::PeerClosed, cleanup_failure);
|
|
return false;
|
|
}
|
|
let deadline = if retained_state_changed {
|
|
match self.drain_deadline(&client, drain_started_at) {
|
|
Some(candidate) => deadline.min(candidate),
|
|
None => {
|
|
cleanup_failure.get_or_insert_with(duration_overflow_failure);
|
|
return self.begin_close(
|
|
index,
|
|
client,
|
|
primary_end,
|
|
cleanup_failure,
|
|
now,
|
|
);
|
|
}
|
|
}
|
|
} else {
|
|
deadline
|
|
};
|
|
if now >= deadline {
|
|
if received && post_deadline_receives_remaining > 1 {
|
|
post_deadline_receives_remaining -= 1;
|
|
self.targets[index].state = TargetState::Draining {
|
|
client,
|
|
drain_started_at,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure,
|
|
post_deadline_receives_remaining,
|
|
};
|
|
true
|
|
} else {
|
|
self.begin_close(index, client, primary_end, cleanup_failure, now)
|
|
}
|
|
} else {
|
|
self.targets[index].state = TargetState::Draining {
|
|
client,
|
|
drain_started_at,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure,
|
|
post_deadline_receives_remaining,
|
|
};
|
|
received
|
|
}
|
|
}
|
|
TargetState::Closing {
|
|
mut client,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure,
|
|
} => match client.poll_close(cx) {
|
|
Poll::Pending => {
|
|
if now >= deadline {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
self.finish_target(
|
|
index,
|
|
primary_end,
|
|
cleanup_failure.or_else(|| Some(close_timeout_failure())),
|
|
);
|
|
false
|
|
} else {
|
|
self.targets[index].state = TargetState::Closing {
|
|
client,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure,
|
|
};
|
|
false
|
|
}
|
|
}
|
|
Poll::Ready(Ok(events)) => {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
self.publish_client_events(index, events);
|
|
#[cfg(test)]
|
|
let injected_close_failure = self.take_drain_close_failure().map(|error| {
|
|
classify_client_error(ManagedTargetFailurePhase::Closing, &error)
|
|
});
|
|
#[cfg(not(test))]
|
|
let injected_close_failure = None;
|
|
self.finish_target(
|
|
index,
|
|
primary_end,
|
|
cleanup_failure.or(injected_close_failure),
|
|
);
|
|
false
|
|
}
|
|
Poll::Ready(Err(error)) => {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
let cleanup = classify_client_error(ManagedTargetFailurePhase::Closing, &error);
|
|
self.finish_target(index, primary_end, cleanup_failure.or(Some(cleanup)));
|
|
false
|
|
}
|
|
},
|
|
TargetState::Terminal => {
|
|
self.targets[index].state = TargetState::Terminal;
|
|
false
|
|
}
|
|
}
|
|
}
|
|
|
|
fn begin_drain(
|
|
&mut self,
|
|
index: usize,
|
|
client: AsyncClient,
|
|
primary_end: ManagedTargetEndReason,
|
|
_now: Instant,
|
|
) -> bool {
|
|
let drain_started_at = Instant::now();
|
|
let Some(deadline) = self.drain_deadline(&client, drain_started_at) else {
|
|
let cleanup_failure = Some(duration_overflow_failure());
|
|
return self.begin_close(
|
|
index,
|
|
client,
|
|
primary_end,
|
|
cleanup_failure,
|
|
drain_started_at,
|
|
);
|
|
};
|
|
self.install_target_state(
|
|
index,
|
|
TargetState::Draining {
|
|
client,
|
|
drain_started_at,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure: None,
|
|
post_deadline_receives_remaining: POST_DEADLINE_RECEIVE_BUDGET,
|
|
},
|
|
);
|
|
true
|
|
}
|
|
|
|
fn drain_deadline(&self, client: &AsyncClient, drain_started_at: Instant) -> Option<Instant> {
|
|
client
|
|
.latest_probe_timeout_deadline()
|
|
.unwrap_or(drain_started_at)
|
|
.max(drain_started_at)
|
|
.checked_add(self.config.final_drain)
|
|
}
|
|
|
|
fn begin_close(
|
|
&mut self,
|
|
index: usize,
|
|
client: AsyncClient,
|
|
primary_end: ManagedTargetEndReason,
|
|
mut cleanup_failure: Option<ManagedTargetFailure>,
|
|
now: Instant,
|
|
) -> bool {
|
|
let deadline = now.checked_add(self.config.final_drain).unwrap_or_else(|| {
|
|
cleanup_failure.get_or_insert_with(duration_overflow_failure);
|
|
now
|
|
});
|
|
self.install_target_state(
|
|
index,
|
|
TargetState::Closing {
|
|
client,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure,
|
|
},
|
|
);
|
|
true
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn take_drain_receive_failure(&mut self) -> Option<ClientError> {
|
|
if !mem::take(&mut self.drain_test_hook.fail_receive) {
|
|
return None;
|
|
}
|
|
Some(ClientError::Socket(std::io::Error::other(
|
|
"injected drain receive failure",
|
|
)))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn take_drain_close_failure(&mut self) -> Option<ClientError> {
|
|
if !mem::take(&mut self.drain_test_hook.fail_close) {
|
|
return None;
|
|
}
|
|
Some(ClientError::Socket(std::io::Error::other(
|
|
"injected drain close failure",
|
|
)))
|
|
}
|
|
|
|
fn active_count(&self) -> usize {
|
|
self.targets
|
|
.iter()
|
|
.filter(|target| target.active && target.desired && target.retirement.is_none())
|
|
.count()
|
|
}
|
|
|
|
fn active_stagger_spacing(&self) -> Option<(usize, Duration)> {
|
|
let mut active = 0;
|
|
let mut minimum: Option<Duration> = None;
|
|
for target in &self.targets {
|
|
if !target.desired || target.retirement.is_some() {
|
|
continue;
|
|
}
|
|
let TargetState::Active { client } = &target.state else {
|
|
continue;
|
|
};
|
|
active += 1;
|
|
let interval = client
|
|
.probe_interval()
|
|
.expect("active managed targets have a committed probe schedule");
|
|
minimum = minimum.into_iter().chain(Some(interval)).min();
|
|
}
|
|
minimum.map(|interval| (active, stagger_spacing(interval, active)))
|
|
}
|
|
|
|
fn stagger_target_added(&mut self, now: Instant) {
|
|
let Some(existing) = self.send_gate.filter(|gate| *gate > now) else {
|
|
self.send_gate = None;
|
|
return;
|
|
};
|
|
let candidate = self
|
|
.last_stagger_send
|
|
.zip(self.active_stagger_spacing())
|
|
.and_then(|(last, (_, spacing))| last.checked_add(spacing))
|
|
.filter(|gate| *gate > now);
|
|
self.send_gate = candidate.map(|candidate| existing.min(candidate));
|
|
}
|
|
|
|
fn stagger_target_removed(&mut self, now: Instant) {
|
|
self.send_gate = self.send_gate.filter(|gate| *gate > now);
|
|
}
|
|
|
|
fn record_stagger_acceptance(
|
|
&mut self,
|
|
result: SendResult,
|
|
stagger_spacing: Option<(usize, Duration)>,
|
|
accepted_at: Instant,
|
|
) {
|
|
if !result.accepted() {
|
|
return;
|
|
}
|
|
let Some((_active, spacing)) = stagger_spacing else {
|
|
return;
|
|
};
|
|
self.last_stagger_send = Some(accepted_at);
|
|
self.send_gate = accepted_at.checked_add(spacing);
|
|
#[cfg(test)]
|
|
if let Some(observations) = &self.stagger_observations {
|
|
observations.lock().unwrap().push((_active, spacing));
|
|
}
|
|
}
|
|
|
|
fn poll_one_send(
|
|
&mut self,
|
|
index: usize,
|
|
cx: &mut Context<'_>,
|
|
now: Instant,
|
|
stagger_spacing: Option<(usize, Duration)>,
|
|
) -> SendResult {
|
|
if !self.targets[index].desired || self.targets[index].retirement.is_some() {
|
|
self.targets[index].send_waiting = false;
|
|
return SendResult::NotAttempted;
|
|
}
|
|
let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal);
|
|
let TargetState::Active { mut client } = state else {
|
|
self.targets[index].state = state;
|
|
return SendResult::NotAttempted;
|
|
};
|
|
if client
|
|
.next_probe_timeout_deadline()
|
|
.is_some_and(|deadline| deadline <= now)
|
|
{
|
|
self.targets[index].send_waiting = false;
|
|
self.targets[index].state = TargetState::Active { client };
|
|
return SendResult::NotAttempted;
|
|
}
|
|
if client
|
|
.next_send_deadline()
|
|
.is_none_or(|deadline| deadline > now)
|
|
{
|
|
self.targets[index].send_waiting = false;
|
|
self.targets[index].state = TargetState::Active { client };
|
|
return SendResult::NotAttempted;
|
|
}
|
|
let before = client.packets_sent();
|
|
let result = client.poll_send_probe(cx);
|
|
let after = client.packets_sent();
|
|
let accepted = after > before;
|
|
let send_result = match &result {
|
|
Poll::Pending => SendResult::Pending,
|
|
Poll::Ready(Ok(_)) => SendResult::Ready { accepted },
|
|
Poll::Ready(Err(_)) => SendResult::Failed { accepted },
|
|
};
|
|
self.record_stagger_acceptance(send_result, stagger_spacing, Instant::now());
|
|
let schedule_error = accepted
|
|
.then(|| client.skip_missed_probe_slots_at(Instant::now()))
|
|
.transpose()
|
|
.err();
|
|
self.targets[index].record_packets_sent(after);
|
|
match result {
|
|
Poll::Pending => {
|
|
self.targets[index].send_waiting = true;
|
|
self.targets[index].state = TargetState::Active { client };
|
|
SendResult::Pending
|
|
}
|
|
Poll::Ready(Ok(events)) => {
|
|
self.targets[index].send_waiting = false;
|
|
self.targets[index].state = TargetState::Active { client };
|
|
self.publish_client_events(index, events);
|
|
if let Some(error) = schedule_error {
|
|
let TargetState::Active { client } =
|
|
mem::replace(&mut self.targets[index].state, TargetState::Terminal)
|
|
else {
|
|
unreachable!("probe sender remained active while publishing send events");
|
|
};
|
|
self.begin_open_session_failure(
|
|
index,
|
|
client,
|
|
ManagedTargetFailurePhase::Timing,
|
|
error,
|
|
now,
|
|
OpenSessionFailureCleanup::Drain,
|
|
);
|
|
SendResult::Failed { accepted }
|
|
} else {
|
|
SendResult::Ready { accepted }
|
|
}
|
|
}
|
|
Poll::Ready(Err(error)) => {
|
|
self.targets[index].send_waiting = false;
|
|
self.begin_open_session_failure(
|
|
index,
|
|
client,
|
|
ManagedTargetFailurePhase::Sending,
|
|
error,
|
|
now,
|
|
OpenSessionFailureCleanup::Drain,
|
|
);
|
|
SendResult::Failed { accepted }
|
|
}
|
|
}
|
|
}
|
|
|
|
fn poll_staggered_send(&mut self, cx: &mut Context<'_>, now: Instant) -> bool {
|
|
if self.active_count() == 0 || self.send_gate.is_some_and(|gate| gate > now) {
|
|
return false;
|
|
}
|
|
if self.stagger_remaining == 0 {
|
|
self.stagger_remaining = self.targets.len();
|
|
}
|
|
let work = self.stagger_remaining.min(TARGET_WORK_BUDGET);
|
|
for _ in 0..work {
|
|
let index = self.send_cursor % self.targets.len();
|
|
self.send_cursor = (self.send_cursor + 1) % self.targets.len();
|
|
self.stagger_remaining -= 1;
|
|
if !self.targets[index].desired
|
|
|| self.targets[index].retirement.is_some()
|
|
|| !matches!(self.targets[index].state, TargetState::Active { .. })
|
|
{
|
|
continue;
|
|
}
|
|
let stagger_spacing = self.active_stagger_spacing();
|
|
let result = self.poll_one_send(index, cx, now, stagger_spacing);
|
|
if result == SendResult::NotAttempted {
|
|
continue;
|
|
}
|
|
self.stagger_remaining = 0;
|
|
return matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. });
|
|
}
|
|
self.stagger_remaining > 0
|
|
}
|
|
|
|
fn poll_burst_sends(&mut self, cx: &mut Context<'_>, now: Instant) -> bool {
|
|
if self.targets.is_empty() {
|
|
return false;
|
|
}
|
|
if self.burst_remaining == 0 {
|
|
self.burst_remaining = self.targets.len();
|
|
}
|
|
let work = self.burst_remaining.min(TARGET_WORK_BUDGET);
|
|
let mut immediate = false;
|
|
for _ in 0..work {
|
|
let index = self.send_cursor % self.targets.len();
|
|
self.send_cursor = (self.send_cursor + 1) % self.targets.len();
|
|
self.burst_remaining -= 1;
|
|
let result = self.poll_one_send(index, cx, now, None);
|
|
immediate |= matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. });
|
|
}
|
|
immediate || self.burst_remaining > 0
|
|
}
|
|
|
|
fn poll_target_pass(&mut self, cx: &mut Context<'_>, now: Instant) -> bool {
|
|
if self.targets.is_empty() {
|
|
return false;
|
|
}
|
|
if self.scan_remaining == 0 {
|
|
self.scan_remaining = self.targets.len();
|
|
}
|
|
let work = self.scan_remaining.min(TARGET_WORK_BUDGET);
|
|
let mut immediate = false;
|
|
for _ in 0..work {
|
|
let index = self.cursor % self.targets.len();
|
|
self.cursor = (self.cursor + 1) % self.targets.len();
|
|
self.scan_remaining -= 1;
|
|
immediate |= self.poll_target(index, cx, now);
|
|
}
|
|
immediate || self.scan_remaining > 0
|
|
}
|
|
|
|
fn 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;
|
|
#[cfg(test)]
|
|
{
|
|
self.timeout_inspections += 1;
|
|
}
|
|
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 {
|
|
match &self.targets[index].state {
|
|
TargetState::Active { client } => {
|
|
self.state != DriverState::Stopping
|
|
&& self.effective_retirement(index).is_none()
|
|
&& client
|
|
.next_probe_timeout_deadline()
|
|
.is_some_and(|deadline| deadline <= now)
|
|
}
|
|
TargetState::Draining {
|
|
client,
|
|
deadline: _deadline,
|
|
..
|
|
} => {
|
|
#[cfg(test)]
|
|
let defer_work = self.drain_test_hook.defer_work_until_deadline && now < *_deadline;
|
|
#[cfg(not(test))]
|
|
let defer_work = false;
|
|
!defer_work
|
|
&& client
|
|
.next_probe_timeout_deadline()
|
|
.is_some_and(|timeout| timeout <= now)
|
|
}
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
fn poll_target_timeout(&mut self, index: usize, now: Instant) -> TimeoutStep {
|
|
let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal);
|
|
match state {
|
|
TargetState::Active { mut client } => match client.poll_timeouts_bounded_at(now, 1) {
|
|
Ok(batch) => {
|
|
self.publish_client_events(index, batch.events);
|
|
self.targets[index].state = TargetState::Active { client };
|
|
TimeoutStep {
|
|
more_due: batch.more_due,
|
|
}
|
|
}
|
|
Err(error) => {
|
|
self.begin_open_session_failure(
|
|
index,
|
|
client,
|
|
ManagedTargetFailurePhase::Timing,
|
|
error,
|
|
now,
|
|
OpenSessionFailureCleanup::Close,
|
|
);
|
|
TimeoutStep::NONE
|
|
}
|
|
},
|
|
TargetState::Draining {
|
|
mut client,
|
|
drain_started_at,
|
|
deadline,
|
|
primary_end,
|
|
mut cleanup_failure,
|
|
post_deadline_receives_remaining,
|
|
} => match client.poll_timeouts_bounded_at(now, 1) {
|
|
Ok(batch) => {
|
|
self.publish_client_events(index, batch.events);
|
|
let deadline = match self.drain_deadline(&client, drain_started_at) {
|
|
Some(candidate) => deadline.min(candidate),
|
|
None => {
|
|
cleanup_failure.get_or_insert_with(duration_overflow_failure);
|
|
self.begin_close(index, client, primary_end, cleanup_failure, now);
|
|
return TimeoutStep::NONE;
|
|
}
|
|
};
|
|
self.targets[index].state = TargetState::Draining {
|
|
client,
|
|
drain_started_at,
|
|
deadline,
|
|
primary_end,
|
|
cleanup_failure,
|
|
post_deadline_receives_remaining,
|
|
};
|
|
TimeoutStep {
|
|
more_due: batch.more_due,
|
|
}
|
|
}
|
|
Err(error) => {
|
|
self.targets[index].sync_packets_sent(&client);
|
|
cleanup_failure.get_or_insert_with(|| {
|
|
classify_client_error(ManagedTargetFailurePhase::Timing, &error)
|
|
});
|
|
self.begin_close(index, client, primary_end, cleanup_failure, now);
|
|
TimeoutStep::NONE
|
|
}
|
|
},
|
|
state => {
|
|
self.targets[index].state = state;
|
|
TimeoutStep::NONE
|
|
}
|
|
}
|
|
}
|
|
|
|
fn all_targets_terminal(&self) -> bool {
|
|
self.targets
|
|
.iter()
|
|
.all(|target| matches!(target.state, TargetState::Terminal))
|
|
}
|
|
|
|
fn next_deadline(&mut self) -> Option<Instant> {
|
|
#[cfg(test)]
|
|
{
|
|
self.deadline_inspections += self.targets.len();
|
|
}
|
|
let mut non_send_deadline = None;
|
|
let mut send_deadline = None;
|
|
for target in &self.targets {
|
|
match &target.state {
|
|
TargetState::Active { client } => {
|
|
non_send_deadline = non_send_deadline
|
|
.into_iter()
|
|
.chain(client.next_probe_timeout_deadline())
|
|
.min();
|
|
if target.desired && target.retirement.is_none() && !target.send_waiting {
|
|
send_deadline = send_deadline
|
|
.into_iter()
|
|
.chain(client.next_send_deadline())
|
|
.min();
|
|
}
|
|
}
|
|
TargetState::Draining {
|
|
client, deadline, ..
|
|
} => {
|
|
#[cfg(test)]
|
|
let defer_work = self.drain_test_hook.defer_work_until_deadline;
|
|
#[cfg(not(test))]
|
|
let defer_work = false;
|
|
non_send_deadline = non_send_deadline
|
|
.into_iter()
|
|
.chain(
|
|
(!defer_work)
|
|
.then(|| client.next_probe_timeout_deadline())
|
|
.flatten(),
|
|
)
|
|
.chain(Some(*deadline))
|
|
.min();
|
|
}
|
|
TargetState::Closing { deadline, .. } => {
|
|
non_send_deadline = non_send_deadline.into_iter().chain(Some(*deadline)).min();
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
if self.state == DriverState::Running
|
|
&& self.config.pacing == ManagedPacing::Staggered
|
|
&& self.active_count() > 0
|
|
{
|
|
send_deadline = match (send_deadline, self.send_gate) {
|
|
(Some(target), Some(gate)) => Some(target.max(gate)),
|
|
(None, gate) => gate,
|
|
(target, None) => target,
|
|
};
|
|
}
|
|
non_send_deadline.into_iter().chain(send_deadline).min()
|
|
}
|
|
|
|
fn register_timer(&mut self, cx: &mut Context<'_>) -> bool {
|
|
let Some(deadline) = self.next_deadline() else {
|
|
self.timer = None;
|
|
return false;
|
|
};
|
|
let timer = self
|
|
.timer
|
|
.get_or_insert_with(|| Box::pin(time::sleep_until(deadline.into())));
|
|
timer.as_mut().reset(deadline.into());
|
|
timer.as_mut().poll(cx).is_ready()
|
|
}
|
|
|
|
fn register_stop_wake(&mut self, cx: &mut Context<'_>) -> bool {
|
|
let Some(mut wake) = self.wake.take() else {
|
|
return false;
|
|
};
|
|
match wake.as_mut().poll(cx) {
|
|
Poll::Pending => {
|
|
self.wake = Some(wake);
|
|
false
|
|
}
|
|
Poll::Ready(Some(receiver)) => {
|
|
self.wake = Some(arm_wake(receiver));
|
|
self.resources().stop.is_requested()
|
|
}
|
|
Poll::Ready(None) => false,
|
|
}
|
|
}
|
|
|
|
fn seal(&mut self, mut end_reason: ManagedEndReason, failed: bool) -> Poll<ManagedOutcome> {
|
|
// 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()
|
|
{
|
|
let error = match &end_reason {
|
|
ManagedEndReason::DriverFailed(failure) => {
|
|
ManagedCommandApplyError::DriverFailed(failure.clone())
|
|
}
|
|
_ => ManagedCommandApplyError::Stopping,
|
|
};
|
|
let _ = acknowledgement.send(Err(error));
|
|
}
|
|
let outcome = Arc::new(
|
|
self.history
|
|
.outcome(end_reason, self.applied_command_sequence),
|
|
);
|
|
self.final_outcome = Some(Arc::clone(&outcome));
|
|
self.lifecycle = if failed {
|
|
ManagedLifecycle::Failed
|
|
} else {
|
|
ManagedLifecycle::Completed
|
|
};
|
|
self.replace_status();
|
|
self.publish_event(if failed {
|
|
ManagedEvent::Failed {
|
|
outcome: Arc::clone(&outcome),
|
|
}
|
|
} else {
|
|
ManagedEvent::Completed {
|
|
outcome: Arc::clone(&outcome),
|
|
}
|
|
});
|
|
self.resources().stop.acknowledge();
|
|
if let Some(resources) = self.resources.as_mut() {
|
|
resources.events.take();
|
|
}
|
|
self.state = DriverState::Finished;
|
|
self.wake = None;
|
|
self.timer = None;
|
|
self.resources.take();
|
|
Poll::Ready((*outcome).clone())
|
|
}
|
|
|
|
fn fail_driver(&mut self, failure: ManagedDriverFailure) -> Poll<ManagedOutcome> {
|
|
self.seal(ManagedEndReason::DriverFailed(failure), true)
|
|
}
|
|
}
|
|
|
|
#[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;
|
|
|
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
let this = self.as_mut().get_mut();
|
|
if this.state == DriverState::Finished {
|
|
panic!("ManagedClientTask polled after completion");
|
|
}
|
|
|
|
if this.state == DriverState::NotStarted {
|
|
if this.resources().stop.is_requested() {
|
|
this.observe_stop();
|
|
} else if let Err(failure) = this.begin_running() {
|
|
return this.fail_driver(failure);
|
|
}
|
|
}
|
|
if this.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_timeout_pass(now);
|
|
immediate |= this.poll_target_pass(cx, now);
|
|
this.prune_undesired_terminal();
|
|
if this.state == DriverState::Running {
|
|
immediate |= match this.config.pacing {
|
|
ManagedPacing::Staggered => this.poll_staggered_send(cx, now),
|
|
ManagedPacing::Burst => this.poll_burst_sends(cx, now),
|
|
};
|
|
}
|
|
|
|
if this.resources().stop.is_requested() {
|
|
this.observe_stop();
|
|
}
|
|
|
|
if this.all_targets_terminal() {
|
|
match this.state {
|
|
DriverState::Stopping => {
|
|
return this.seal(
|
|
if this.stop_observed {
|
|
ManagedEndReason::StopRequested
|
|
} else {
|
|
ManagedEndReason::TargetsComplete
|
|
},
|
|
false,
|
|
);
|
|
}
|
|
DriverState::Running
|
|
if this.config.completion == ManagedCompletionPolicy::FinishWhenQuiescent =>
|
|
{
|
|
this.state = DriverState::Stopping;
|
|
this.lifecycle = ManagedLifecycle::Stopping;
|
|
this.resources().stop.begin_stopping();
|
|
this.replace_status();
|
|
this.publish_event(ManagedEvent::Stopping);
|
|
return this.seal(ManagedEndReason::TargetsComplete, false);
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
immediate |= this.register_stop_wake(cx);
|
|
if !immediate {
|
|
immediate |= this.register_timer(cx);
|
|
}
|
|
if immediate {
|
|
cx.waker().wake_by_ref();
|
|
}
|
|
Poll::Pending
|
|
}
|
|
}
|
|
|
|
impl Drop for ManagedClientTask {
|
|
fn drop(&mut self) {
|
|
if self.state == DriverState::Finished || self.resources.is_none() {
|
|
return;
|
|
}
|
|
self.lifecycle = ManagedLifecycle::Abandoned;
|
|
self.final_outcome = None;
|
|
self.resources().stop.close_updates();
|
|
self.commands.close();
|
|
self.replace_status();
|
|
self.publish_event(ManagedEvent::Abandoned);
|
|
self.resources().stop.acknowledge();
|
|
if let Some(resources) = self.resources.as_mut() {
|
|
resources.events.take();
|
|
}
|
|
self.resources.take();
|
|
}
|
|
}
|
|
|
|
fn build_task(
|
|
config: ManagedClientConfig,
|
|
targets: Vec<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.server_addr.clone_from(&target.server_addr);
|
|
if let Some(auth) = &target.auth {
|
|
client_config.hmac_key.clone_from(&auth.hmac_key);
|
|
}
|
|
validate_target_config(&client_config).map_err(|source| {
|
|
ManagedConfigError::InvalidTarget {
|
|
id: target.id.clone(),
|
|
source,
|
|
}
|
|
})?;
|
|
runtimes.push(TargetRuntime {
|
|
instance: TargetInstance {
|
|
id: target.id.clone(),
|
|
generation,
|
|
},
|
|
config: target.clone(),
|
|
desired: true,
|
|
retirement: None,
|
|
server_addr: Arc::from(target.server_addr),
|
|
remote: None,
|
|
counters: TargetCounters::default(),
|
|
send_waiting: false,
|
|
active: false,
|
|
state: TargetState::Pending { client_config },
|
|
});
|
|
}
|
|
|
|
let (event_sender, _) = broadcast::channel(config.event_capacity);
|
|
let weak_events = event_sender.downgrade();
|
|
let (wake_sender, wake_receiver) = watch::channel(());
|
|
let (acknowledgement_sender, _) = watch::channel(());
|
|
let stop = Arc::new(StopSignal {
|
|
requested: AtomicBool::new(false),
|
|
update_admission: AtomicU8::new(UpdateAdmission::Open as u8),
|
|
wake: wake_sender,
|
|
acknowledged: AtomicBool::new(false),
|
|
acknowledgement: acknowledgement_sender,
|
|
});
|
|
let history = OutcomeHistory::new(config.outcome_history_limit);
|
|
let initial_targets = runtimes
|
|
.iter()
|
|
.map(|target| ManagedTargetStatus {
|
|
target: target.instance.clone(),
|
|
desired: true,
|
|
lifecycle: ManagedTargetLifecycle::Pending,
|
|
server_addr: Arc::clone(&target.server_addr),
|
|
remote: None,
|
|
})
|
|
.collect::<Vec<_>>();
|
|
let initial = Arc::new(ManagedStatus {
|
|
lifecycle: ManagedLifecycle::NotStarted,
|
|
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: None,
|
|
});
|
|
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 {
|
|
state: DriverState::NotStarted,
|
|
lifecycle: ManagedLifecycle::NotStarted,
|
|
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,
|
|
final_outcome: None,
|
|
next_generation,
|
|
applied_command_sequence: 0,
|
|
#[cfg(test)]
|
|
event_observations: None,
|
|
#[cfg(test)]
|
|
stagger_observations: None,
|
|
#[cfg(test)]
|
|
drain_test_hook: DrainTestHook::default(),
|
|
#[cfg(test)]
|
|
timeout_inspections: 0,
|
|
#[cfg(test)]
|
|
deadline_inspections: 0,
|
|
};
|
|
let handle = ManagedClientHandle {
|
|
stop,
|
|
commands: command_sender,
|
|
max_live_target_generations,
|
|
status: status_receiver,
|
|
events: weak_events,
|
|
#[cfg(test)]
|
|
update_submission_hook: None,
|
|
};
|
|
Ok((task, handle))
|
|
}
|
|
|
|
fn validate_target_config(config: &ClientConfig) -> Result<(), ClientError> {
|
|
validate_open_timeouts(&config.open_timeouts)?;
|
|
SessionMachine::validate_config(config)?;
|
|
if config.socket_config.ipv4_only && config.socket_config.ipv6_only {
|
|
return Err(ClientError::InvalidConfig {
|
|
reason: "ipv4_only and ipv6_only cannot both be true".to_owned(),
|
|
});
|
|
}
|
|
if let Some(ttl) = config.socket_config.ttl {
|
|
validate_ttl(ttl)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn stagger_spacing(interval: Duration, active_targets: usize) -> Duration {
|
|
let divisor = u128::try_from(active_targets.max(1)).unwrap_or(u128::MAX);
|
|
let nanos = (interval.as_nanos() / divisor).max(1);
|
|
let seconds = u64::try_from(nanos / 1_000_000_000).unwrap_or(u64::MAX);
|
|
let subsec = u32::try_from(nanos % 1_000_000_000).unwrap_or(999_999_999);
|
|
Duration::new(seconds, subsec)
|
|
}
|
|
|
|
fn duration_overflow_failure() -> ManagedTargetFailure {
|
|
classify_client_error(
|
|
ManagedTargetFailurePhase::Timing,
|
|
&ClientError::DurationOverflow,
|
|
)
|
|
}
|
|
|
|
fn close_timeout_failure() -> ManagedTargetFailure {
|
|
ManagedTargetFailure {
|
|
phase: ManagedTargetFailurePhase::Closing,
|
|
kind: ManagedTargetFailureKind::Timeout,
|
|
message: Arc::from("best-effort close did not become writable before its deadline"),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "tests.rs"]
|
|
mod tests;
|