Merge pull request #50 from Lochnair/feat/tokio-managed-target-updates

feat(managed): add dynamic target generations
This commit was merged in pull request #50.
This commit is contained in:
Nils Andreas Svee authored and GitHub committed 2026-08-09 19:03:54 +07:00
commit 53d29f4cf9
3 files changed
+954 -42

No files matched your search

+51
View File
@@ -8,6 +8,8 @@ use super::TargetId;
/// Default capacity of the lossy managed event channel.
pub const DEFAULT_MANAGED_EVENT_CAPACITY: usize = 256;
/// Default capacity of the bounded managed control queue.
pub const DEFAULT_MANAGED_COMMAND_CAPACITY: usize = 64;
/// Default number of recent target outcomes retained in status and output.
pub const DEFAULT_MANAGED_OUTCOME_HISTORY_LIMIT: usize = 256;
/// Default maximum number of simultaneously live target generations.
@@ -65,6 +67,7 @@ pub struct ManagedClientConfig {
pub pacing: ManagedPacing,
pub completion: ManagedCompletionPolicy,
pub event_capacity: usize,
pub command_capacity: usize,
pub outcome_history_limit: usize,
pub max_live_target_generations: usize,
pub final_drain: Duration,
@@ -77,6 +80,7 @@ impl Default for ManagedClientConfig {
pacing: ManagedPacing::default(),
completion: ManagedCompletionPolicy::default(),
event_capacity: DEFAULT_MANAGED_EVENT_CAPACITY,
command_capacity: DEFAULT_MANAGED_COMMAND_CAPACITY,
outcome_history_limit: DEFAULT_MANAGED_OUTCOME_HISTORY_LIMIT,
max_live_target_generations: DEFAULT_MANAGED_MAX_LIVE_TARGET_GENERATIONS,
final_drain: DEFAULT_MANAGED_FINAL_DRAIN,
@@ -111,6 +115,7 @@ pub enum ManagedTargetLifecycle {
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ManagedTargetStatus {
pub target: TargetInstance,
pub desired: bool,
pub lifecycle: ManagedTargetLifecycle,
pub server_addr: Arc<str>,
pub remote: Option<SocketAddr>,
@@ -121,6 +126,7 @@ pub struct ManagedTargetStatus {
pub struct ManagedStatus {
pub lifecycle: ManagedLifecycle,
pub stop_requested: bool,
pub applied_command_sequence: u64,
pub desired_target_count: usize,
pub connecting_target_count: usize,
pub opening_target_count: usize,
@@ -171,6 +177,7 @@ pub type ManagedEventSubscription = tokio::sync::broadcast::Receiver<ManagedEven
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ManagedOutcome {
pub end_reason: ManagedEndReason,
pub applied_command_sequence: u64,
pub total_target_outcomes: u64,
pub successful_target_outcomes: u64,
pub failed_target_outcomes: u64,
@@ -258,6 +265,8 @@ pub enum ManagedConfigError {
TooManyTargets { configured: usize, limit: usize },
#[error("managed event capacity must be greater than zero")]
ZeroEventCapacity,
#[error("managed command capacity must be greater than zero")]
ZeroCommandCapacity,
#[error("managed live-generation limit must be greater than zero")]
ZeroLiveGenerationLimit,
#[error("managed final drain {duration:?} cannot be scheduled from the current instant")]
@@ -272,6 +281,48 @@ pub enum ManagedConfigError {
GenerationExhausted,
}
/// Immediate failure to submit a managed target update.
#[derive(Clone, Debug, Eq, Error, PartialEq)]
pub enum ManagedCommandError {
#[error("managed target updates are no longer accepted")]
Stopping,
#[error("managed command queue is full")]
QueueFull,
#[error("managed command receiver is closed")]
DriverClosed,
}
/// Failure while applying an accepted managed target update.
#[derive(Debug, Error)]
pub enum ManagedCommandApplyError {
#[error("managed target updates are no longer accepted")]
Stopping,
#[error("duplicate target id {id}")]
DuplicateTargetId { id: TargetId },
#[error("invalid configuration for target {id}: {source}")]
InvalidTarget {
id: TargetId,
#[source]
source: ClientError,
},
#[error("target update requires {required} live generations but limit is {limit}")]
LiveGenerationLimitExceeded { required: usize, limit: usize },
#[error("managed target generation space is exhausted")]
GenerationExhausted,
#[error("managed command sequence is exhausted")]
CommandSequenceExhausted,
#[error("managed task ended before command acknowledgement")]
AcknowledgementDisconnected,
#[error("managed driver failed before command application: {0}")]
DriverFailed(ManagedDriverFailure),
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ManagedCommandAcknowledgement {
pub sequence: u64,
pub status: Arc<ManagedStatus>,
}
/// Failure to create a new event subscription.
#[derive(Clone, Debug, Eq, Error, PartialEq)]
pub enum ManagedSubscribeError {
+487 -27
View File
@@ -13,7 +13,7 @@ use std::{
};
use tokio::{
sync::{broadcast, watch},
sync::{broadcast, mpsc, oneshot, watch},
time::{self, Sleep},
};
@@ -26,7 +26,8 @@ use crate::{
};
use super::{
classify_client_error, ManagedClientConfig, ManagedCompletionPolicy, ManagedConfigError,
classify_client_error, ManagedClientConfig, ManagedCommandAcknowledgement,
ManagedCommandApplyError, ManagedCommandError, ManagedCompletionPolicy, ManagedConfigError,
ManagedDriverFailure, ManagedEndReason, ManagedEvent, ManagedEventSubscription,
ManagedLifecycle, ManagedOutcome, ManagedPacing, ManagedStatus, ManagedSubscribeError,
ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetFailure, ManagedTargetFailureKind,
@@ -35,6 +36,7 @@ use super::{
};
const TARGET_WORK_BUDGET: usize = 128;
const COMMAND_WORK_BUDGET: usize = 32;
const POST_DEADLINE_RECEIVE_BUDGET: usize = TARGET_WORK_BUDGET;
/// Entry point for constructing a unified Tokio managed task.
@@ -54,6 +56,7 @@ impl ManagedClient {
#[derive(Debug)]
struct StopSignal {
requested: AtomicBool,
updates_accepted: AtomicBool,
wake: watch::Sender<()>,
acknowledged: AtomicBool,
acknowledgement: watch::Sender<()>,
@@ -61,11 +64,19 @@ struct StopSignal {
impl StopSignal {
fn request(&self) {
self.updates_accepted.store(false, Ordering::Release);
if !self.requested.swap(true, Ordering::AcqRel) {
self.wake.send_replace(());
}
}
fn updates_accepted(&self) -> bool {
self.updates_accepted.load(Ordering::Acquire)
}
fn close_updates(&self) {
self.updates_accepted.store(false, Ordering::Release);
}
fn is_requested(&self) -> bool {
self.requested.load(Ordering::Acquire)
}
@@ -81,6 +92,7 @@ impl StopSignal {
#[derive(Clone)]
pub struct ManagedClientHandle {
stop: Arc<StopSignal>,
commands: mpsc::Sender<ManagedCommand>,
status: watch::Receiver<Arc<ManagedStatus>>,
events: broadcast::WeakSender<ManagedEvent>,
}
@@ -112,6 +124,52 @@ impl ManagedClientHandle {
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> {
if !self.stop.updates_accepted() {
return Err(ManagedCommandError::Stopping);
}
let (acknowledgement, receiver) = oneshot::channel();
match self.commands.try_send(ManagedCommand::UpdateTargets {
targets,
acknowledgement,
}) {
Ok(()) => Ok(ManagedCommandReceipt { receiver }),
Err(mpsc::error::TrySendError::Full(_)) => Err(ManagedCommandError::QueueFull),
Err(mpsc::error::TrySendError::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.
@@ -211,6 +269,9 @@ struct TargetCounters {
struct TargetRuntime {
instance: TargetInstance,
config: ManagedTargetConfig,
desired: bool,
retirement: Option<ManagedTargetEndReason>,
server_addr: Arc<str>,
remote: Option<std::net::SocketAddr>,
counters: TargetCounters,
@@ -219,6 +280,57 @@ struct TargetRuntime {
state: TargetState,
}
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
}
}
}
fn target_outcome(
target: &TargetRuntime,
end_reason: ManagedTargetEndReason,
cleanup_failure: Option<ManagedTargetFailure>,
) -> ManagedTargetOutcome {
ManagedTargetOutcome {
target: target.instance.clone(),
server_addr: Arc::clone(&target.server_addr),
remote: target.remote,
end_reason,
packets_sent: target.counters.packets_sent,
replies_received: target.counters.replies_received,
duplicates: target.counters.duplicates,
late: target.counters.late,
warning_events: target.counters.warning_events,
cleanup_failure,
}
}
#[derive(Default)]
struct OutcomeHistory {
limit: usize,
@@ -275,9 +387,14 @@ impl OutcomeHistory {
)
}
fn outcome(&self, end_reason: ManagedEndReason) -> ManagedOutcome {
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,
@@ -315,6 +432,7 @@ pub struct ManagedClientTask {
lifecycle: ManagedLifecycle,
config: ManagedClientConfig,
targets: Vec<TargetRuntime>,
commands: mpsc::Receiver<ManagedCommand>,
history: OutcomeHistory,
resources: Option<TaskResources>,
wake: Option<WakeFuture>,
@@ -328,7 +446,8 @@ pub struct ManagedClientTask {
last_stagger_send: Option<Instant>,
stop_observed: bool,
final_outcome: Option<Arc<ManagedOutcome>>,
_next_generation: u64,
next_generation: u64,
applied_command_sequence: u64,
#[cfg(test)]
event_observations: Option<EventObservations>,
#[cfg(test)]
@@ -388,6 +507,7 @@ impl ManagedClientTask {
}
ManagedTargetStatus {
target: target.instance.clone(),
desired: target.desired,
lifecycle,
server_addr: Arc::clone(&target.server_addr),
remote: target.remote,
@@ -397,7 +517,8 @@ impl ManagedClientTask {
Arc::new(ManagedStatus {
lifecycle: self.lifecycle,
stop_requested: self.stop_observed,
desired_target_count: self.targets.len(),
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,
@@ -536,6 +657,7 @@ impl ManagedClientTask {
return;
}
self.stop_observed = true;
self.resources().stop.close_updates();
self.state = DriverState::Stopping;
self.lifecycle = ManagedLifecycle::Stopping;
self.replace_status();
@@ -545,6 +667,256 @@ impl ManagedClientTask {
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.is_requested()
{
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(target, 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.close_updates();
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.send_cursor = 0;
} else {
self.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 });
@@ -554,8 +926,14 @@ impl ManagedClientTask {
let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal);
match state {
TargetState::Pending { client_config } => {
if self.state == DriverState::Stopping {
self.finish_target(index, ManagedTargetEndReason::Stopped, None);
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);
@@ -563,8 +941,14 @@ impl ManagedClientTask {
}
}
TargetState::Connecting { mut future } => {
if self.state == DriverState::Stopping {
self.finish_target(index, ManagedTargetEndReason::Stopped, None);
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) {
@@ -593,10 +977,14 @@ impl ManagedClientTask {
mut client,
mut open,
} => {
if self.state == DriverState::Stopping {
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].counters.packets_sent = client.packets_sent();
self.finish_target(index, ManagedTargetEndReason::Stopped, None);
self.finish_target(index, retirement, None);
return false;
}
open.request_stop_after_current_attempt();
@@ -608,9 +996,17 @@ impl ManagedClientTask {
}
Poll::Ready(Ok(OpenOutcome::Started { event, .. })) => {
self.publish_client_events(index, vec![event]);
if self.state == DriverState::Stopping {
if self.state == DriverState::Stopping
|| self.effective_retirement(index).is_some()
{
self.targets[index].counters.packets_sent = client.packets_sent();
self.begin_drain(index, client, ManagedTargetEndReason::Stopped, now)
self.begin_drain(
index,
client,
self.effective_retirement(index)
.unwrap_or(ManagedTargetEndReason::Stopped),
now,
)
} else {
self.install_target_state(index, TargetState::Active { client });
true
@@ -618,8 +1014,11 @@ impl ManagedClientTask {
}
Poll::Ready(Ok(OpenOutcome::NoTestCompleted { event, .. })) => {
self.publish_client_events(index, vec![event]);
let end_reason = if self.state == DriverState::Stopping {
ManagedTargetEndReason::Stopped
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
};
@@ -627,7 +1026,9 @@ impl ManagedClientTask {
false
}
Poll::Ready(Err(error)) => {
if self.state == DriverState::Stopping {
if self.state == DriverState::Stopping
|| self.effective_retirement(index).is_some()
{
let cleanup_failure =
(!open.stopped_after_current_attempt()).then(|| {
classify_client_error(
@@ -637,7 +1038,8 @@ impl ManagedClientTask {
});
self.finish_target(
index,
ManagedTargetEndReason::Stopped,
self.effective_retirement(index)
.unwrap_or(ManagedTargetEndReason::Stopped),
cleanup_failure,
);
} else {
@@ -648,10 +1050,17 @@ impl ManagedClientTask {
}
}
TargetState::Active { mut client } => {
if self.state == DriverState::Stopping {
if self.state == DriverState::Stopping || self.effective_retirement(index).is_some()
{
client.discard_prepared_probe();
self.targets[index].counters.packets_sent = client.packets_sent();
return self.begin_drain(index, client, ManagedTargetEndReason::Stopped, now);
return self.begin_drain(
index,
client,
self.effective_retirement(index)
.unwrap_or(ManagedTargetEndReason::Stopped),
now,
);
}
if client
.next_probe_timeout_deadline()
@@ -969,13 +1378,19 @@ impl ManagedClientTask {
}
fn active_count(&self) -> usize {
self.targets.iter().filter(|target| target.active).count()
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;
};
@@ -1032,6 +1447,10 @@ impl ManagedClientTask {
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;
@@ -1116,7 +1535,10 @@ impl ManagedClientTask {
let index = self.send_cursor % self.targets.len();
self.send_cursor = (self.send_cursor + 1) % self.targets.len();
self.stagger_remaining -= 1;
if !matches!(self.targets[index].state, TargetState::Active { .. }) {
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();
@@ -1183,7 +1605,7 @@ impl ManagedClientTask {
.into_iter()
.chain(client.next_probe_timeout_deadline())
.min();
if !target.send_waiting {
if target.desired && target.retirement.is_none() && !target.send_waiting {
send_deadline = send_deadline
.into_iter()
.chain(client.next_send_deadline())
@@ -1256,7 +1678,22 @@ impl ManagedClientTask {
}
fn seal(&mut self, end_reason: ManagedEndReason, failed: bool) -> Poll<ManagedOutcome> {
let outcome = Arc::new(self.history.outcome(end_reason));
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
@@ -1326,8 +1763,10 @@ impl Future for ManagedClientTask {
this.observe_stop();
}
let mut immediate = this.process_commands(cx);
let now = Instant::now();
let mut immediate = this.poll_target_pass(cx, 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),
@@ -1338,13 +1777,21 @@ impl Future for ManagedClientTask {
if this.all_targets_terminal() {
match this.state {
DriverState::Stopping => {
return this.seal(ManagedEndReason::StopRequested, false);
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.close_updates();
this.replace_status();
this.publish_event(ManagedEvent::Stopping);
return this.seal(ManagedEndReason::TargetsComplete, false);
@@ -1386,6 +1833,9 @@ fn build_task(
if config.event_capacity == 0 {
return Err(ManagedConfigError::ZeroEventCapacity);
}
if config.command_capacity == 0 {
return Err(ManagedConfigError::ZeroCommandCapacity);
}
if config.max_live_target_generations == 0 {
return Err(ManagedConfigError::ZeroLiveGenerationLimit);
}
@@ -1428,9 +1878,12 @@ fn build_task(
})?;
runtimes.push(TargetRuntime {
instance: TargetInstance {
id: target.id,
id: target.id.clone(),
generation,
},
config: target.clone(),
desired: true,
retirement: None,
server_addr: Arc::from(target.server_addr),
remote: None,
counters: TargetCounters::default(),
@@ -1446,6 +1899,7 @@ fn build_task(
let (acknowledgement_sender, _) = watch::channel(());
let stop = Arc::new(StopSignal {
requested: AtomicBool::new(false),
updates_accepted: AtomicBool::new(true),
wake: wake_sender,
acknowledged: AtomicBool::new(false),
acknowledgement: acknowledgement_sender,
@@ -1455,6 +1909,7 @@ fn build_task(
.iter()
.map(|target| ManagedTargetStatus {
target: target.instance.clone(),
desired: true,
lifecycle: ManagedTargetLifecycle::Pending,
server_addr: Arc::clone(&target.server_addr),
remote: None,
@@ -1463,6 +1918,7 @@ fn build_task(
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,
@@ -1480,6 +1936,7 @@ fn build_task(
final_outcome: None,
});
let (status_sender, status_receiver) = watch::channel(initial);
let (command_sender, command_receiver) = mpsc::channel(config.command_capacity);
let resources = TaskResources {
status: status_sender,
events: Some(event_sender),
@@ -1490,6 +1947,7 @@ fn build_task(
lifecycle: ManagedLifecycle::NotStarted,
config,
targets: runtimes,
commands: command_receiver,
history,
resources: Some(resources),
wake: Some(arm_wake(wake_receiver)),
@@ -1503,7 +1961,8 @@ fn build_task(
last_stagger_send: None,
stop_observed: false,
final_outcome: None,
_next_generation: next_generation,
next_generation,
applied_command_sequence: 0,
#[cfg(test)]
event_observations: None,
#[cfg(test)]
@@ -1513,6 +1972,7 @@ fn build_task(
};
let handle = ManagedClientHandle {
stop,
commands: command_sender,
status: status_receiver,
events: weak_events,
};
+416 -15
View File
@@ -465,27 +465,428 @@ fn explicit_empty_waits_for_stop() {
}
#[test]
fn duplicate_and_limit_rejections_are_transactional() {
let mut limited = config(ManagedPacing::Staggered);
limited.max_live_target_generations = 1;
fn dynamic_addition_opens_and_operates() {
let server = start_server(ServerBehavior::Echo, None);
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
managed.client.duration = None;
let (task, handle) = ManagedClient::task(managed, vec![]).unwrap();
let receipt = handle
.update_targets(vec![target("added", server.addr)])
.unwrap();
let runtime = runtime();
let mut task = Box::pin(task);
drive_task_until(&runtime, &mut task, |task| {
task.targets[0].counters.packets_sent != 0
});
let acknowledgement = runtime.block_on(receipt).unwrap();
assert_eq!(acknowledgement.sequence, 1);
assert_eq!(acknowledgement.status.targets[0].target.generation, 1);
assert!(acknowledgement.status.targets[0].desired);
let stop = handle.stop();
let outcome = runtime.block_on(task);
runtime.block_on(stop);
assert_eq!(outcome.applied_command_sequence, 1);
assert!(!probes(&server.finish()).is_empty());
}
#[test]
fn dynamic_identical_open_target_preserves_session_without_pending_event() {
let server = start_server(ServerBehavior::Echo, None);
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
managed.client.duration = None;
let configured = target("same", server.addr);
let (mut task, handle) = ManagedClient::task(managed, vec![configured.clone()]).unwrap();
let observations = Arc::new(Mutex::new(Vec::new()));
task.event_observations = Some(Arc::clone(&observations));
let runtime = runtime();
let mut task = Box::pin(task);
drive_task_until(&runtime, &mut task, |task| {
task.targets[0].counters.packets_sent >= 2
});
let instance = task.targets[0].instance.clone();
let sent = task.targets[0].counters.packets_sent;
let receipt = handle.update_targets(vec![configured]).unwrap();
drive_task_until(&runtime, &mut task, |task| {
task.applied_command_sequence == 1
});
let acknowledgement = runtime.block_on(receipt).unwrap();
assert_eq!(acknowledgement.status.targets[0].target, instance);
drive_task_until(&runtime, &mut task, |task| {
task.targets[0].counters.packets_sent > sent
});
assert!(!observations.lock().unwrap().iter().any(|(event, _)| {
matches!(event, ManagedEvent::TargetStateChanged { target, lifecycle: ManagedTargetLifecycle::Pending } if *target == instance)
}));
let stop = handle.stop();
runtime.block_on(task);
runtime.block_on(stop);
assert_eq!(
server
.finish()
.iter()
.filter(|record| matches!(record.kind, PacketKind::Open))
.count(),
1
);
}
#[test]
fn dynamic_retiring_active_is_ineligible_before_target_cleanup() {
let server = start_server(ServerBehavior::Echo, None);
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
managed.client.duration = None;
let (task, handle) = ManagedClient::task(managed, vec![target("old", server.addr)]).unwrap();
let runtime = runtime();
let mut task = Box::pin(task);
drive_task_until(&runtime, &mut task, |task| {
task.targets[0].counters.packets_sent != 0
});
let old = task.targets[0].instance.clone();
let sent = task.targets[0].counters.packets_sent;
let acknowledgement = task
.as_mut()
.get_mut()
.apply_targets(vec![ManagedTargetConfig::new("old", "127.0.0.1:9")])
.unwrap();
assert!(acknowledgement
.status
.targets
.iter()
.any(|target| target.target == old && !target.desired));
assert!(matches!(task.targets[0].state, TargetState::Active { .. }));
assert_eq!(
task.targets[0].retirement,
Some(ManagedTargetEndReason::Replaced)
);
let result = runtime.block_on(poll_fn(|cx| {
Poll::Ready(
task.as_mut()
.get_mut()
.poll_one_send(0, cx, Instant::now(), None),
)
}));
assert!(matches!(result, SendResult::NotAttempted));
assert_eq!(task.targets[0].counters.packets_sent, sent);
let stop = handle.stop();
runtime.block_on(task);
runtime.block_on(stop);
server.finish();
}
#[test]
fn dynamic_replacement_counts_sync_and_async_live_generations() {
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
managed.max_live_target_generations = 1;
let (mut pending, _) = ManagedClient::task(
managed.clone(),
vec![ManagedTargetConfig::new("same", "127.0.0.1:9")],
)
.unwrap();
let acknowledgement = pending
.apply_targets(vec![ManagedTargetConfig::new("same", "127.0.0.1:10")])
.unwrap();
assert_eq!(acknowledgement.status.targets.len(), 1);
assert_eq!(acknowledgement.status.targets[0].target.generation, 2);
assert!(matches!(
ManagedClient::task(
limited,
acknowledgement.status.recent_target_outcomes[0].end_reason,
ManagedTargetEndReason::Replaced
));
let server = start_server(ServerBehavior::Echo, None);
let replacement = start_server(ServerBehavior::Echo, None);
managed.client.duration = None;
let (task, handle) = ManagedClient::task(managed, vec![target("same", server.addr)]).unwrap();
let runtime = runtime();
let mut task = Box::pin(task);
drive_task_until(&runtime, &mut task, |task| {
task.targets[0].counters.packets_sent != 0
});
let before = task.targets[0].instance.clone();
let before_status = task.snapshot();
let before_generation = task.next_generation;
let before_sequence = task.applied_command_sequence;
assert!(matches!(
task.as_mut()
.get_mut()
.apply_targets(vec![target("same", replacement.addr)]),
Err(ManagedCommandApplyError::LiveGenerationLimitExceeded {
required: 2,
limit: 1
})
));
assert_eq!(task.snapshot(), before_status);
assert_eq!(task.next_generation, before_generation);
assert_eq!(task.applied_command_sequence, before_sequence);
assert_eq!(task.targets.len(), 1);
task.config.max_live_target_generations = 2;
let acknowledgement = task
.as_mut()
.get_mut()
.apply_targets(vec![target("same", replacement.addr)])
.unwrap();
assert!(acknowledgement
.status
.targets
.iter()
.any(|target| target.target == before && !target.desired));
assert_eq!(acknowledgement.status.targets.len(), 2);
drive_task_until(&runtime, &mut task, |task| {
task.targets
.iter()
.any(|target| target.desired && target.counters.packets_sent != 0)
});
let stop = handle.stop();
let outcome = runtime.block_on(task);
runtime.block_on(stop);
assert!(outcome.recent_target_outcomes.iter().any(|outcome| {
outcome.target == before && matches!(outcome.end_reason, ManagedTargetEndReason::Replaced)
}));
server.finish();
replacement.finish();
}
#[test]
fn terminal_targets_count_toward_retained_generation_limit() {
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
managed.max_live_target_generations = 1;
let (mut task, _) = ManagedClient::task(
managed.clone(),
vec![ManagedTargetConfig::new("same", "127.0.0.1:9")],
)
.unwrap();
task.targets[0].state = TargetState::Terminal;
assert!(matches!(
task.apply_targets(vec![
ManagedTargetConfig::new("same", "127.0.0.1:9"),
ManagedTargetConfig::new("other", "127.0.0.1:10"),
]),
Err(ManagedCommandApplyError::LiveGenerationLimitExceeded {
required: 2,
limit: 1,
})
));
let (mut task, _) = ManagedClient::task(
managed.clone(),
vec![ManagedTargetConfig::new("same", "127.0.0.1:9")],
)
.unwrap();
task.targets[0].state = TargetState::Terminal;
let acknowledgement = task
.apply_targets(vec![ManagedTargetConfig::new("same", "127.0.0.1:10")])
.unwrap();
assert_eq!(task.targets.len(), 1);
assert_eq!(acknowledgement.status.targets.len(), 1);
assert_eq!(acknowledgement.status.targets[0].target.generation, 2);
let (mut task, _) = ManagedClient::task(
managed,
vec![ManagedTargetConfig::new("same", "127.0.0.1:9")],
)
.unwrap();
task.targets[0].state = TargetState::Terminal;
task.apply_targets(Vec::new()).unwrap();
assert!(task.targets.is_empty());
assert!(task.snapshot().targets.is_empty());
}
#[test]
fn pruning_preserves_target_cursor_progress() {
let (mut task, _) = ManagedClient::task(
config(ManagedPacing::Staggered),
vec![
target("a", "127.0.0.1:1".parse().unwrap()),
target("b", "127.0.0.1:2".parse().unwrap())
]
),
Err(ManagedConfigError::TooManyTargets { .. })
));
ManagedTargetConfig::new("one", "127.0.0.1:9"),
ManagedTargetConfig::new("two", "127.0.0.1:10"),
ManagedTargetConfig::new("three", "127.0.0.1:11"),
],
)
.unwrap();
task.cursor = 5;
task.send_cursor = 5;
task.scan_remaining = 3;
task.burst_remaining = 3;
task.stagger_remaining = 3;
task.targets.truncate(2);
task.rebase_target_cursors();
assert_eq!(task.cursor, 1);
assert_eq!(task.send_cursor, 1);
assert_eq!(task.scan_remaining, 2);
assert_eq!(task.burst_remaining, 2);
assert_eq!(task.stagger_remaining, 2);
task.targets.clear();
task.rebase_target_cursors();
assert_eq!(task.cursor, 0);
assert_eq!(task.send_cursor, 0);
assert_eq!(task.scan_remaining, 0);
assert_eq!(task.burst_remaining, 0);
assert_eq!(task.stagger_remaining, 0);
}
#[test]
fn dynamic_rejections_leave_transaction_state_unchanged() {
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
let (mut task, _) = ManagedClient::task(
managed,
vec![ManagedTargetConfig::new("same", "127.0.0.1:9")],
)
.unwrap();
let status = task.snapshot();
let generation = task.next_generation;
let sequence = task.applied_command_sequence;
let instances = task
.targets
.iter()
.map(|target| (target.instance.clone(), target.desired))
.collect::<Vec<_>>();
let duplicate = vec![
target("same", "127.0.0.1:1".parse().unwrap()),
target("same", "127.0.0.1:2".parse().unwrap()),
ManagedTargetConfig::new("same", "127.0.0.1:10"),
ManagedTargetConfig::new("same", "127.0.0.1:11"),
];
assert!(task.apply_targets(duplicate).is_err());
task.config.client.open_timeouts.clear();
assert!(task
.apply_targets(vec![ManagedTargetConfig::new("same", "127.0.0.1:10")])
.is_err());
assert_eq!(task.next_generation, generation);
assert_eq!(task.applied_command_sequence, sequence);
assert_eq!(task.snapshot(), status);
assert_eq!(
task.targets
.iter()
.map(|target| (target.instance.clone(), target.desired))
.collect::<Vec<_>>(),
instances
);
}
#[test]
fn dynamic_empty_update_orders_stopping_before_acknowledgement() {
let mut managed = config(ManagedPacing::Staggered);
managed.command_capacity = 2;
let (mut task, handle) = ManagedClient::task(
managed,
vec![ManagedTargetConfig::new("gone", "127.0.0.1:9")],
)
.unwrap();
let observations = Arc::new(Mutex::new(Vec::new()));
task.event_observations = Some(Arc::clone(&observations));
let removal = handle.update_targets(vec![]).unwrap();
let later = handle
.update_targets(vec![ManagedTargetConfig::new("later", "127.0.0.1:9")])
.unwrap();
let outcome = runtime().block_on(task);
let acknowledgement = runtime().block_on(removal).unwrap();
assert_eq!(acknowledgement.status.lifecycle, ManagedLifecycle::Stopping);
assert_eq!(acknowledgement.status.desired_target_count, 0);
assert!(matches!(
ManagedClient::task(config(ManagedPacing::Staggered), duplicate),
Err(ManagedConfigError::DuplicateTargetId { .. })
runtime().block_on(later),
Err(ManagedCommandApplyError::Stopping)
));
assert_eq!(outcome.end_reason, ManagedEndReason::TargetsComplete);
let events = observations.lock().unwrap();
let stopping = events
.iter()
.position(|(event, _)| matches!(event, ManagedEvent::Stopping))
.unwrap();
let finished = events
.iter()
.position(|(event, _)| matches!(event, ManagedEvent::TargetFinished { .. }))
.unwrap();
assert!(stopping < finished);
assert_eq!(events[stopping].1.lifecycle, ManagedLifecycle::Stopping);
assert!(Arc::ptr_eq(&acknowledgement.status, &events[stopping].1));
}
#[test]
fn dynamic_explicit_remove_and_readd_uses_new_generation() {
let first = start_server(ServerBehavior::Echo, None);
let second = start_server(ServerBehavior::Echo, None);
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
managed.client.duration = None;
let (task, handle) = ManagedClient::task(managed, vec![target("same", first.addr)]).unwrap();
let runtime = runtime();
let mut task = Box::pin(task);
drive_task_until(&runtime, &mut task, |task| task.active_count() == 1);
let first_instance = task.targets[0].instance.clone();
let removal = handle.update_targets(vec![]).unwrap();
drive_task_until(&runtime, &mut task, |task| task.targets.is_empty());
let acknowledgement = runtime.block_on(removal).unwrap();
assert_eq!(acknowledgement.status.lifecycle, ManagedLifecycle::Running);
assert!(acknowledgement
.status
.targets
.iter()
.any(|target| target.target == first_instance && !target.desired));
assert!(matches!(
task.history.recent.back().unwrap().end_reason,
ManagedTargetEndReason::Removed
));
let addition = handle
.update_targets(vec![target("same", second.addr)])
.unwrap();
drive_task_until(&runtime, &mut task, |task| {
task.targets[0].counters.packets_sent != 0
});
let acknowledgement = runtime.block_on(addition).unwrap();
assert!(acknowledgement.status.targets[0].target.generation > first_instance.generation);
let stop = handle.stop();
runtime.block_on(task);
runtime.block_on(stop);
assert!(has_close(&first.finish()));
second.finish();
}
#[test]
fn dynamic_queue_stop_and_dropped_receipt_are_linearized() {
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
managed.command_capacity = 1;
let (task, handle) = ManagedClient::task(managed.clone(), vec![]).unwrap();
let receipt = handle
.update_targets(vec![ManagedTargetConfig::new("queued", "127.0.0.1:9")])
.unwrap();
assert!(matches!(
handle.update_targets(vec![]),
Err(ManagedCommandError::QueueFull)
));
let stop = handle.stop();
let outcome = runtime().block_on(task);
runtime().block_on(stop);
assert_eq!(outcome.applied_command_sequence, 0);
assert!(matches!(
runtime().block_on(receipt),
Err(ManagedCommandApplyError::Stopping)
));
let server = start_server(ServerBehavior::Echo, None);
managed.command_capacity = 2;
managed.client.duration = None;
let (task, handle) = ManagedClient::task(managed, vec![]).unwrap();
drop(
handle
.update_targets(vec![target("dropped", server.addr)])
.unwrap(),
);
let runtime = runtime();
let mut task = Box::pin(task);
drive_task_until(&runtime, &mut task, |task| {
task.targets[0].counters.packets_sent != 0
});
assert_eq!(handle.status().applied_command_sequence, 1);
assert_eq!(handle.status().targets[0].target.generation, 1);
let stop = handle.stop();
runtime.block_on(task);
runtime.block_on(stop);
assert!(!probes(&server.finish()).is_empty());
}
#[test]