fix(managed): preserve drain and pacing invariants

This commit is contained in:
Nils Andreas Svee committed 2026-08-05 00:09:46 +07:00
1 parent 8445c0cd99
commit 99b6ddfead
2 files changed
+317 -83

No files matched your search

+162 -48
View File
@@ -171,6 +171,7 @@ enum TargetState {
},
Draining {
client: AsyncClient,
drain_started_at: Instant,
deadline: Instant,
primary_end: ManagedTargetEndReason,
cleanup_failure: Option<ManagedTargetFailure>,
@@ -213,7 +214,6 @@ struct TargetRuntime {
server_addr: Arc<str>,
remote: Option<std::net::SocketAddr>,
counters: TargetCounters,
latest_committed_timeout: Option<Instant>,
send_waiting: bool,
active: bool,
state: TargetState,
@@ -302,6 +302,12 @@ enum DriverState {
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 {
@@ -416,11 +422,14 @@ impl ManagedClientTask {
fn install_target_state(&mut self, index: usize, state: TargetState) {
let lifecycle = state.lifecycle();
let active = matches!(state, TargetState::Active { .. });
let membership_changed = self.targets[index].active != active;
let was_active = self.targets[index].active;
self.targets[index].active = active;
self.targets[index].state = state;
if membership_changed {
self.recompute_stagger_gate(Instant::now());
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 {
@@ -493,6 +502,26 @@ impl ManagedClientTask {
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].counters.packets_sent = client.packets_sent();
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;
@@ -631,9 +660,14 @@ impl ManagedClientTask {
match client.poll_timeouts_at(now) {
Ok(events) => self.publish_client_events(index, events),
Err(error) => {
self.targets[index].counters.packets_sent = client.packets_sent();
self.fail_target(index, ManagedTargetFailurePhase::Timing, error);
return false;
return self.begin_open_session_failure(
index,
client,
ManagedTargetFailurePhase::Timing,
error,
now,
OpenSessionFailureCleanup::Close,
);
}
}
}
@@ -644,9 +678,14 @@ impl ManagedClientTask {
true
}
Poll::Ready(Err(error)) => {
self.targets[index].counters.packets_sent = client.packets_sent();
self.fail_target(index, ManagedTargetFailurePhase::Receiving, error);
return false;
return self.begin_open_session_failure(
index,
client,
ManagedTargetFailurePhase::Receiving,
error,
now,
OpenSessionFailureCleanup::Close,
);
}
};
if client.is_peer_closed() {
@@ -668,6 +707,7 @@ impl ManagedClientTask {
}
TargetState::Draining {
mut client,
drain_started_at,
deadline,
primary_end,
mut cleanup_failure,
@@ -680,6 +720,7 @@ impl ManagedClientTask {
if defer_work {
self.targets[index].state = TargetState::Draining {
client,
drain_started_at,
deadline,
primary_end,
cleanup_failure,
@@ -692,12 +733,16 @@ impl ManagedClientTask {
let injected_receive_failure = self.take_drain_receive_failure();
#[cfg(not(test))]
let injected_receive_failure = None;
let mut retained_state_changed = false;
if client
.next_probe_timeout_deadline()
.is_some_and(|timeout| timeout <= now)
{
match client.poll_timeouts_at(now) {
Ok(events) => self.publish_client_events(index, events),
Ok(events) => {
self.publish_client_events(index, events);
retained_state_changed = true;
}
Err(error) => {
self.targets[index].counters.packets_sent = client.packets_sent();
cleanup_failure.get_or_insert_with(|| {
@@ -717,6 +762,7 @@ impl ManagedClientTask {
Poll::Pending => false,
Poll::Ready(Ok(events)) => {
self.publish_client_events(index, events);
retained_state_changed = true;
true
}
Poll::Ready(Err(error)) => {
@@ -739,11 +785,29 @@ impl ManagedClientTask {
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,
@@ -756,6 +820,7 @@ impl ManagedClientTask {
} else {
self.targets[index].state = TargetState::Draining {
client,
drain_started_at,
deadline,
primary_end,
cleanup_failure,
@@ -824,23 +889,24 @@ impl ManagedClientTask {
index: usize,
client: AsyncClient,
primary_end: ManagedTargetEndReason,
now: Instant,
_now: Instant,
) -> bool {
let latest = self.targets[index]
.latest_committed_timeout
.into_iter()
.chain(client.latest_probe_timeout_deadline())
.max()
.unwrap_or(now)
.max(now);
let Some(deadline) = latest.checked_add(self.config.final_drain) else {
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, now);
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,
@@ -850,6 +916,14 @@ impl ManagedClientTask {
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,
@@ -914,15 +988,50 @@ impl ManagedClientTask {
minimum.map(|interval| (active, stagger_spacing(interval, active)))
}
fn recompute_stagger_gate(&mut self, now: Instant) {
self.send_gate = self
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 poll_one_send(&mut self, index: usize, cx: &mut Context<'_>, now: Instant) -> SendResult {
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 {
let state = mem::replace(&mut self.targets[index].state, TargetState::Terminal);
let TargetState::Active { mut client } = state else {
self.targets[index].state = state;
@@ -940,18 +1049,17 @@ impl ManagedClientTask {
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].counters.packets_sent = after;
if accepted {
self.targets[index].latest_committed_timeout = client
.latest_probe_timeout_deadline()
.into_iter()
.chain(self.targets[index].latest_committed_timeout)
.max();
}
match result {
Poll::Pending => {
self.targets[index].send_waiting = true;
@@ -963,7 +1071,19 @@ impl ManagedClientTask {
self.targets[index].state = TargetState::Active { client };
self.publish_client_events(index, events);
if let Some(error) = schedule_error {
self.fail_target(index, ManagedTargetFailurePhase::Timing, 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 }
@@ -971,7 +1091,14 @@ impl ManagedClientTask {
}
Poll::Ready(Err(error)) => {
self.targets[index].send_waiting = false;
self.fail_target(index, ManagedTargetFailurePhase::Sending, error);
self.begin_open_session_failure(
index,
client,
ManagedTargetFailurePhase::Sending,
error,
now,
OpenSessionFailureCleanup::Drain,
);
SendResult::Failed { accepted }
}
}
@@ -992,24 +1119,12 @@ impl ManagedClientTask {
if !matches!(self.targets[index].state, TargetState::Active { .. }) {
continue;
}
#[cfg(test)]
let stagger_spacing = self.active_stagger_spacing();
let result = self.poll_one_send(index, cx, now);
let result = self.poll_one_send(index, cx, now, stagger_spacing);
if result == SendResult::NotAttempted {
continue;
}
self.stagger_remaining = 0;
if result.accepted() {
#[cfg(test)]
if let Some((_active, spacing)) = stagger_spacing {
if let Some(observations) = &self.stagger_observations {
observations.lock().unwrap().push((_active, spacing));
}
}
let accepted_at = Instant::now();
self.last_stagger_send = Some(accepted_at);
self.recompute_stagger_gate(accepted_at);
}
return matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. });
}
self.stagger_remaining > 0
@@ -1028,7 +1143,7 @@ impl ManagedClientTask {
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);
let result = self.poll_one_send(index, cx, now, None);
immediate |= matches!(result, SendResult::Ready { .. } | SendResult::Failed { .. });
}
immediate || self.burst_remaining > 0
@@ -1319,7 +1434,6 @@ fn build_task(
server_addr: Arc::from(target.server_addr),
remote: None,
counters: TargetCounters::default(),
latest_committed_timeout: None,
send_waiting: false,
active: false,
state: TargetState::Pending { client_config },
+155 -35
View File
@@ -1,5 +1,6 @@
use std::{
future::{poll_fn, Future},
mem,
net::{SocketAddr, UdpSocket},
pin::Pin,
sync::{Arc, Condvar, Mutex},
@@ -524,14 +525,19 @@ fn pre_poll_stop() {
}
#[test]
fn finite_one_target_lifecycle() {
fn finite_immediate_replies_use_retained_drain_deadline() {
let server = start_server(ServerBehavior::Echo, None);
let (task, _) = ManagedClient::task(
config(ManagedPacing::Staggered),
vec![target("one", server.addr)],
)
.unwrap();
let outcome = runtime().block_on(task);
let mut managed = config(ManagedPacing::Staggered);
managed.client.probe_timeout = Duration::from_secs(3);
managed.final_drain = Duration::from_millis(30);
let (task, _) = ManagedClient::task(managed, vec![target("one", server.addr)]).unwrap();
let started_at = Instant::now();
let outcome = runtime().block_on(async {
tokio::time::timeout(Duration::from_secs(1), task)
.await
.expect("replied probes must not hold the drain until their obsolete timeout")
});
assert!(started_at.elapsed() < Duration::from_secs(1));
let records = server.finish();
assert_eq!(outcome.end_reason, ManagedEndReason::TargetsComplete);
assert!(matches!(
@@ -793,7 +799,8 @@ fn staggered_pacing_uses_all_active_negotiated_intervals() {
#[test]
fn stagger_gate_tracks_active_membership() {
let first = start_server_negotiating(ServerBehavior::Echo, None, Some(Duration::from_secs(1)));
let first =
start_server_negotiating(ServerBehavior::Echo, None, Some(Duration::from_millis(100)));
let second_gate = Arc::new(PacketGate::default());
let second = start_server_with_gates(
ServerBehavior::Echo,
@@ -803,7 +810,7 @@ fn stagger_gate_tracks_active_membership() {
None,
);
let mut managed = config(ManagedPacing::Staggered);
managed.client.duration = Some(Duration::from_secs(2));
managed.client.duration = Some(Duration::from_millis(1_250));
managed.client.interval = Duration::from_secs(1);
managed.client.negotiation_policy = NegotiationPolicy::Loose;
managed.client.open_timeouts = vec![Duration::from_secs(2)];
@@ -818,19 +825,44 @@ fn stagger_gate_tracks_active_membership() {
task.targets[0].counters.packets_sent == 1
});
let one_target_last = task.last_stagger_send.unwrap();
let one_target_gate = task.send_gate.unwrap();
assert_eq!(
task.send_gate,
one_target_last.checked_add(Duration::from_secs(1))
one_target_last.checked_add(Duration::from_millis(100))
);
second_gate.release();
drive_task_until(&runtime, &mut task, |task| task.active_count() == 2);
let last = task.last_stagger_send.unwrap();
let gate = task.send_gate.unwrap();
assert_eq!(gate, last.checked_add(Duration::from_millis(500)).unwrap());
assert!(gate > last && gate < last.checked_add(Duration::from_secs(1)).unwrap());
assert_eq!(gate, last.checked_add(Duration::from_millis(50)).unwrap());
assert!(gate <= one_target_gate);
assert!(gate > last && gate < last.checked_add(Duration::from_millis(100)).unwrap());
drive_task_until(&runtime, &mut task, |task| {
task.targets[1].counters.packets_sent == 1
});
let accepted_at = Instant::now();
task.record_stagger_acceptance(
SendResult::Failed { accepted: true },
Some((2, Duration::from_millis(50))),
accepted_at,
);
assert_eq!(
task.send_gate,
accepted_at.checked_add(Duration::from_millis(50))
);
let preserved_gate = task.send_gate;
let TargetState::Active { client } =
mem::replace(&mut task.targets[0].state, TargetState::Terminal)
else {
panic!("fast target was not active before removal");
};
assert!(task.begin_drain(0, client, ManagedTargetEndReason::Stopped, Instant::now()));
assert_eq!(task.send_gate, preserved_gate);
drive_task_until(&runtime, &mut task, |task| {
task.targets[1].counters.packets_sent == 2
});
drop(handle.stop());
let outcome = runtime.block_on(task);
assert_eq!(outcome.failed_target_outcomes, 0);
@@ -907,37 +939,125 @@ fn overdue_queued_reply_is_lost_then_late() {
}
#[test]
fn committed_probe_graceful_drain() {
let server = start_server(ServerBehavior::DelayedEcho(Duration::from_millis(25)), None);
let mut config = config(ManagedPacing::Staggered);
config.completion = ManagedCompletionPolicy::ExplicitStop;
config.client.duration = None;
config.client.interval = Duration::from_millis(100);
config.client.probe_timeout = Duration::from_millis(80);
config.final_drain = Duration::from_millis(5);
let (task, handle) = ManagedClient::task(config, vec![target("one", server.addr)]).unwrap();
let stop_handle = handle.clone();
let seen = Arc::clone(&server.probe_seen);
let stopper = thread::spawn(move || {
let (flag, ready) = &*seen;
let guard = flag.lock().unwrap();
let (guard, timeout) = ready
.wait_timeout_while(guard, Duration::from_secs(2), |seen| !*seen)
.unwrap();
let probe_seen = *guard;
drop(guard);
drop(stop_handle.stop());
assert!(probe_seen && !timeout.timed_out());
fn retained_probe_deadline_shortens_after_late_reply() {
let reply_gate = Arc::new(PacketGate::default());
let server = start_server_with_gates(
ServerBehavior::Echo,
None,
None,
None,
Some(Arc::clone(&reply_gate)),
);
let mut managed = config(ManagedPacing::Staggered);
managed.completion = ManagedCompletionPolicy::ExplicitStop;
managed.client.duration = None;
managed.client.interval = Duration::from_secs(1);
managed.client.probe_timeout = Duration::from_millis(100);
managed.final_drain = Duration::from_millis(200);
let final_drain = managed.final_drain;
let (task, handle) = ManagedClient::task(managed, vec![target("one", 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 == 1
});
let outcome = runtime().block_on(task);
stopper.join().unwrap();
drop(handle.stop());
drive_task_until(&runtime, &mut task, |task| {
matches!(task.targets[0].state, TargetState::Draining { .. })
});
let (drain_started_at, initial_deadline, timeout_at) = match &task.targets[0].state {
TargetState::Draining {
client,
drain_started_at,
deadline,
..
} => (
*drain_started_at,
*deadline,
client.next_probe_timeout_deadline().unwrap(),
),
_ => unreachable!(),
};
assert_eq!(
initial_deadline,
timeout_at.checked_add(final_drain).unwrap()
);
thread::sleep(timeout_at.saturating_duration_since(Instant::now()));
while Instant::now() <= timeout_at {
thread::yield_now();
}
drive_task_until(&runtime, &mut task, |task| match &task.targets[0].state {
TargetState::Draining { client, .. } => {
client.next_probe_timeout_deadline().is_none()
&& client.latest_probe_timeout_deadline() == Some(timeout_at)
}
_ => false,
});
let retained_deadline = match &task.targets[0].state {
TargetState::Draining { deadline, .. } => *deadline,
_ => unreachable!(),
};
assert_eq!(retained_deadline, initial_deadline);
reply_gate.release();
wait_flag(&server.reply_sent);
let shortened_deadline = drain_started_at.checked_add(final_drain).unwrap();
drive_task_until(&runtime, &mut task, |task| {
matches!(
task.targets[0].state,
TargetState::Draining { deadline, .. } if deadline == shortened_deadline
)
});
assert!(shortened_deadline < initial_deadline);
let outcome = runtime.block_on(task);
let records = server.finish();
let target = &outcome.recent_target_outcomes[0];
assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested);
assert!(matches!(target.end_reason, ManagedTargetEndReason::Stopped));
assert_eq!(target.packets_sent, 1);
assert_eq!(target.replies_received, 0);
assert_eq!(target.late, 1);
assert_eq!(probes(&records).len(), 1);
}
#[test]
fn pending_limit_failure_drains_reply_and_closes_session() {
let key = b"managed-send-failure-cleanup".to_vec();
let server = start_server(
ServerBehavior::DelayedEcho(Duration::from_millis(80)),
Some(key.clone()),
);
let mut managed = config(ManagedPacing::Staggered);
managed.client.duration = Some(Duration::from_millis(250));
managed.client.interval = Duration::from_millis(20);
managed.client.probe_timeout = Duration::from_secs(2);
managed.client.max_pending_probes = 1;
managed.final_drain = Duration::from_millis(50);
let mut configured = target("one", server.addr);
configured.auth = Some(ClientAuthConfig {
hmac_key: Some(key),
});
let (task, _) = ManagedClient::task(managed, vec![configured]).unwrap();
let outcome = runtime().block_on(async {
tokio::time::timeout(Duration::from_secs(1), task)
.await
.expect("send failure cleanup must shorten after draining the committed reply")
});
let records = server.finish();
let target = &outcome.recent_target_outcomes[0];
assert_eq!(outcome.failed_target_outcomes, 1);
assert!(matches!(
target.end_reason,
ManagedTargetEndReason::Failed(ManagedTargetFailure {
phase: ManagedTargetFailurePhase::Sending,
kind: ManagedTargetFailureKind::ResourceExhausted,
..
})
));
assert_eq!(target.cleanup_failure, None);
assert_eq!(target.packets_sent, 1);
assert_eq!(target.replies_received, 1);
assert_eq!(probes(&records).len(), 1);
assert!(has_close(&records));
}
#[test]