use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
use std::time::Duration;
use tokio::sync::watch;
use crate::config::Shutdown as ShutdownConfig;
use crate::telemetry::metrics;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Phase {
Serving,
Draining,
Closing,
}
impl Phase {
pub fn as_str(self) -> &'static str {
match self {
Self::Serving => "serving",
Self::Draining => "draining",
Self::Closing => "closing",
}
}
fn from_code(code: u8) -> Self {
match code {
0 => Self::Serving,
1 => Self::Draining,
_ => Self::Closing,
}
}
fn code(self) -> u8 {
match self {
Self::Serving => 0,
Self::Draining => 1,
Self::Closing => 2,
}
}
}
pub struct Lifecycle {
phase: AtomicU8,
in_flight: AtomicU64,
abandoning: AtomicBool,
changes: watch::Sender<u64>,
}
impl Default for Lifecycle {
fn default() -> Self {
Self::new()
}
}
impl Lifecycle {
pub fn new() -> Self {
Self {
phase: AtomicU8::new(Phase::Serving.code()),
in_flight: AtomicU64::new(0),
abandoning: AtomicBool::new(false),
changes: watch::Sender::new(0),
}
}
pub fn phase(&self) -> Phase {
Phase::from_code(self.phase.load(Ordering::Acquire))
}
pub fn in_flight(&self) -> u64 {
self.in_flight.load(Ordering::Relaxed)
}
pub fn admit(self: &Arc<Self>) -> Option<Admitted> {
if self.phase() == Phase::Closing {
return None;
}
self.in_flight.fetch_add(1, Ordering::Relaxed);
Some(Admitted {
lifecycle: Arc::clone(self),
})
}
pub fn begin_drain(&self) {
let _ = self.phase.compare_exchange(
Phase::Serving.code(),
Phase::Draining.code(),
Ordering::AcqRel,
Ordering::Acquire,
);
}
pub fn close(&self) {
self.phase.store(Phase::Closing.code(), Ordering::Release);
self.announce();
}
pub async fn closed(&self) {
self.wait_for(|| self.phase() == Phase::Closing).await;
}
fn announce(&self) {
self.changes.send_modify(|version| *version += 1);
}
async fn wait_for(&self, mut reached: impl FnMut() -> bool) {
let mut changes = self.changes.subscribe();
while !reached() {
if changes.changed().await.is_err() {
return;
}
}
}
pub fn abandon(&self) {
self.abandoning.store(true, Ordering::Release);
self.announce();
}
pub async fn abandoned(&self) {
self.wait_for(|| self.abandoning.load(Ordering::Acquire))
.await;
}
pub async fn quiesce(&self, bound: Duration) -> u64 {
let _ = tokio::time::timeout(bound, self.wait_for(|| self.in_flight() == 0)).await;
self.in_flight()
}
}
pub struct Admitted {
lifecycle: Arc<Lifecycle>,
}
impl Drop for Admitted {
fn drop(&mut self) {
if self.lifecycle.in_flight.fetch_sub(1, Ordering::AcqRel) == 1 {
self.lifecycle.announce();
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Plan {
pub drain_grace: Duration,
pub deadline: Duration,
pub flush_timeout: Duration,
}
impl From<&ShutdownConfig> for Plan {
fn from(config: &ShutdownConfig) -> Self {
Self {
drain_grace: Duration::from_millis(config.drain_grace_ms),
deadline: Duration::from_millis(config.deadline_ms),
flush_timeout: Duration::from_millis(config.flush_timeout_ms),
}
}
}
impl Plan {
pub fn settle_share(self) -> Duration {
self.flush_timeout / 2
}
}
#[derive(Clone, Default)]
pub struct ResolvedPlan(Arc<std::sync::OnceLock<Plan>>);
impl ResolvedPlan {
pub fn new() -> Self {
Self::default()
}
pub fn or(&self, fallback: Plan) -> Plan {
self.0.get().copied().unwrap_or(fallback)
}
fn publish(&self, plan: Plan) {
let _ = self.0.set(plan);
}
}
#[derive(Debug, PartialEq, Eq)]
pub enum Outcome {
Completed,
Abandoned { in_flight: u64 },
Failed(String),
}
impl Outcome {
pub fn as_str(&self) -> &'static str {
match self {
Self::Completed => "completed",
Self::Abandoned { .. } => "abandoned",
Self::Failed(_) => "failed",
}
}
}
pub async fn drain(
lifecycle: Arc<Lifecycle>,
signals: Signals,
plan: impl Fn() -> Plan,
resolved: ResolvedPlan,
) {
let mut signals = signals;
let signal = signals.recv().await;
let plan = plan();
resolved.publish(plan);
lifecycle.begin_drain();
metrics::record_shutdown_phase(Phase::Draining);
tracing::warn!(
signal,
drain_grace_ms = plan.drain_grace.as_millis() as u64,
deadline_ms = plan.deadline.as_millis() as u64,
in_flight = lifecycle.in_flight(),
"shutdown requested: readiness now fails while admitted requests keep being served"
);
if !plan.drain_grace.is_zero() {
tokio::select! {
() = tokio::time::sleep(plan.drain_grace) => {}
second = signals.recv() => {
tracing::warn!(signal = second, "second termination signal: closing admission now");
}
}
}
lifecycle.close();
metrics::record_shutdown_phase(Phase::Closing);
tracing::info!(
in_flight = lifecycle.in_flight(),
deadline_ms = plan.deadline.as_millis() as u64,
"admission closed: new requests are refused with `draining`"
);
tokio::spawn(async move {
loop {
let signal = signals.recv().await;
tracing::warn!(
signal,
"termination signal ignored: admission is already closed and the remaining \
waits are bounded"
);
}
});
}
pub async fn serve_bounded<S, E>(
served: S,
lifecycle: &Lifecycle,
resolved: &ResolvedPlan,
boot: Plan,
) -> Outcome
where
S: std::future::IntoFuture<Output = Result<(), E>>,
E: std::fmt::Display,
{
let mut served = std::pin::pin!(served.into_future());
{
let closed = std::pin::pin!(lifecycle.closed());
tokio::select! {
result = &mut served => return finish(result),
() = closed => {}
}
}
let plan = resolved.or(boot);
match tokio::time::timeout(plan.deadline, served).await {
Ok(result) => finish(result),
Err(_) => {
let in_flight = lifecycle.in_flight();
metrics::record_shutdown_abandoned(in_flight);
tracing::warn!(
in_flight,
deadline_ms = plan.deadline.as_millis() as u64,
"shutdown deadline expired: ending still-open responses; streamed spend is \
settled as `client_cancelled` up to the last relayed token"
);
lifecycle.abandon();
Outcome::Abandoned { in_flight }
}
}
}
fn finish<E: std::fmt::Display>(result: Result<(), E>) -> Outcome {
match result {
Ok(()) => Outcome::Completed,
Err(error) => Outcome::Failed(error.to_string()),
}
}
pub struct Signals(Source);
enum Source {
#[cfg(unix)]
Os {
terminate: tokio::signal::unix::Signal,
interrupt: tokio::signal::unix::Signal,
},
#[cfg(not(unix))]
CtrlC,
#[cfg(test)]
Scripted(tokio::sync::mpsc::UnboundedReceiver<&'static str>),
}
impl Signals {
#[cfg(unix)]
pub fn install() -> std::io::Result<Self> {
use tokio::signal::unix::{SignalKind, signal};
Ok(Self(Source::Os {
terminate: signal(SignalKind::terminate())?,
interrupt: signal(SignalKind::interrupt())?,
}))
}
#[cfg(not(unix))]
pub fn install() -> std::io::Result<Self> {
Ok(Self(Source::CtrlC))
}
#[cfg(test)]
fn once() -> Self {
let (deliver, signals) = Self::scripted();
deliver("SIGTERM");
signals
}
#[cfg(test)]
fn scripted() -> (impl Fn(&'static str), Self) {
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
(
move |signal| {
let _ = tx.send(signal);
},
Self(Source::Scripted(rx)),
)
}
pub async fn recv(&mut self) -> &'static str {
match &mut self.0 {
#[cfg(unix)]
Source::Os {
terminate,
interrupt,
} => {
tokio::select! {
_ = terminate.recv() => "SIGTERM",
_ = interrupt.recv() => "SIGINT",
}
}
#[cfg(not(unix))]
Source::CtrlC => {
let _ = tokio::signal::ctrl_c().await;
"ctrl-c"
}
#[cfg(test)]
Source::Scripted(signals) => match signals.recv().await {
Some(signal) => signal,
None => std::future::pending().await,
},
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_settle_waits_can_never_take_the_whole_flush_budget() {
for flush_timeout_ms in [0, 1, 500, 5_000, 60_000] {
let plan = Plan::from(&ShutdownConfig {
flush_timeout_ms,
..ShutdownConfig::default()
});
let reserved = plan.flush_timeout - plan.settle_share();
assert!(
plan.settle_share() <= plan.flush_timeout / 2,
"settle share must be capped at half of {flush_timeout_ms}ms"
);
assert!(
reserved >= plan.settle_share(),
"the flush must keep at least the share the waits get"
);
}
}
#[test]
fn draining_fails_readiness_but_still_admits() {
let lifecycle = Arc::new(Lifecycle::new());
assert_eq!(lifecycle.phase(), Phase::Serving);
lifecycle.begin_drain();
assert_eq!(lifecycle.phase(), Phase::Draining);
let admitted = lifecycle.admit().expect("draining still admits");
assert_eq!(lifecycle.in_flight(), 1);
drop(admitted);
assert_eq!(lifecycle.in_flight(), 0);
}
#[test]
fn closing_refuses_new_work_and_never_reopens() {
let lifecycle = Arc::new(Lifecycle::new());
let admitted = lifecycle.admit().expect("serving admits");
lifecycle.close();
assert!(lifecycle.admit().is_none(), "admission must be closed");
assert_eq!(lifecycle.in_flight(), 1);
lifecycle.begin_drain();
assert_eq!(lifecycle.phase(), Phase::Closing);
drop(admitted);
assert_eq!(lifecycle.in_flight(), 0);
}
#[tokio::test]
async fn closed_resolves_for_a_waiter_that_arrives_late() {
let lifecycle = Arc::new(Lifecycle::new());
lifecycle.close();
tokio::time::timeout(Duration::from_secs(1), lifecycle.closed())
.await
.expect("a late waiter still resolves");
}
fn plan() -> Plan {
Plan {
drain_grace: Duration::from_millis(10),
deadline: Duration::from_millis(50),
flush_timeout: Duration::from_millis(200),
}
}
#[tokio::test]
async fn a_finished_server_completes_without_waiting_for_the_deadline() {
let lifecycle = Lifecycle::new();
let outcome = serve_bounded(
async { Ok::<(), std::io::Error>(()) },
&lifecycle,
&ResolvedPlan::new(),
plan(),
)
.await;
assert_eq!(outcome, Outcome::Completed);
}
#[tokio::test]
async fn the_deadline_enforced_is_the_one_read_when_the_signal_arrived() {
let lifecycle = Arc::new(Lifecycle::new());
let _admitted = lifecycle.admit().expect("serving admits");
let resolved = ResolvedPlan::new();
let reloaded = Plan {
drain_grace: Duration::ZERO,
deadline: Duration::from_secs(30),
flush_timeout: Duration::from_millis(200),
};
let draining = {
let lifecycle = Arc::clone(&lifecycle);
let resolved = resolved.clone();
tokio::spawn(async move {
drain(
lifecycle,
Signals::once(),
move || reloaded,
resolved.clone(),
)
.await;
})
};
let served = {
let lifecycle = Arc::clone(&lifecycle);
async move {
lifecycle.closed().await;
tokio::time::sleep(Duration::from_millis(150)).await;
Ok::<(), std::io::Error>(())
}
};
let outcome = serve_bounded(served, &lifecycle, &resolved, plan()).await;
assert_eq!(
outcome,
Outcome::Completed,
"the reloaded deadline must be the one enforced"
);
assert_eq!(resolved.or(plan()).deadline, reloaded.deadline);
assert_eq!(
resolved.or(plan()).flush_timeout,
reloaded.flush_timeout,
"the flush budget must come from the same snapshot as the deadline"
);
draining.await.expect("the drain finished");
}
#[test]
fn the_boot_plan_stands_when_no_signal_was_ever_handled() {
assert_eq!(ResolvedPlan::new().or(plan()), plan());
}
#[tokio::test]
async fn work_still_in_flight_at_the_deadline_is_told_to_end() {
let lifecycle = Arc::new(Lifecycle::new());
let admitted = lifecycle.admit().expect("serving admits");
lifecycle.close();
let body = {
let lifecycle = Arc::clone(&lifecycle);
tokio::spawn(async move {
lifecycle.abandoned().await;
drop(admitted);
})
};
let outcome = serve_bounded(
std::future::pending::<Result<(), std::io::Error>>(),
&lifecycle,
&ResolvedPlan::new(),
plan(),
)
.await;
assert_eq!(outcome, Outcome::Abandoned { in_flight: 1 });
body.await.expect("the response ended");
assert_eq!(
lifecycle.in_flight(),
0,
"the abandoned response must have released its slot"
);
}
#[tokio::test]
async fn a_server_error_is_reported_rather_than_swallowed() {
let lifecycle = Lifecycle::new();
let outcome = serve_bounded(
async { Err::<(), std::io::Error>(std::io::Error::other("listener failed")) },
&lifecycle,
&ResolvedPlan::new(),
plan(),
)
.await;
assert!(matches!(outcome, Outcome::Failed(message) if message.contains("listener failed")));
}
#[tokio::test]
async fn quiesce_returns_what_is_still_in_flight_when_the_bound_expires() {
let lifecycle = Arc::new(Lifecycle::new());
let _admitted = lifecycle.admit().expect("serving admits");
assert_eq!(lifecycle.quiesce(Duration::from_millis(20)).await, 1);
}
#[test]
fn a_change_published_before_the_first_poll_is_still_observable() {
let lifecycle = Arc::new(Lifecycle::new());
let changes = lifecycle.changes.subscribe();
assert!(!changes.has_changed().expect("the sender is alive"));
let admitted = lifecycle.admit().expect("serving admits");
lifecycle.close();
drop(admitted);
assert!(
changes.has_changed().expect("the sender is alive"),
"a waiter armed before the change must not have to be woken again"
);
}
#[tokio::test]
async fn a_wait_created_before_the_change_resolves_without_a_wake_up() {
let lifecycle = Arc::new(Lifecycle::new());
let admitted = lifecycle.admit().expect("serving admits");
let abandoned = lifecycle.abandoned();
let quiesced = lifecycle.quiesce(Duration::from_secs(30));
lifecycle.abandon();
drop(admitted);
tokio::time::timeout(Duration::from_secs(1), abandoned)
.await
.expect("`abandoned` must not wait for a notification that already fired");
assert_eq!(
tokio::time::timeout(Duration::from_secs(1), quiesced)
.await
.expect("`quiesce` must not sleep out its bound once nothing is in flight"),
0
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn quiesce_sees_a_guard_released_while_it_is_arming() {
let started = std::time::Instant::now();
for _ in 0..1_000 {
let lifecycle = Arc::new(Lifecycle::new());
let admitted = lifecycle.admit().expect("serving admits");
let releasing = tokio::task::spawn_blocking(move || drop(admitted));
assert_eq!(lifecycle.quiesce(Duration::from_secs(30)).await, 0);
releasing.await.expect("the guard was released");
assert!(
started.elapsed() < Duration::from_secs(20),
"quiesce stalled: a release was missed and the bound is being slept out"
);
}
}
#[tokio::test]
async fn a_second_signal_cuts_the_grace_window_short() {
let lifecycle = Arc::new(Lifecycle::new());
let (deliver, signals) = Signals::scripted();
let long_grace = Plan {
drain_grace: Duration::from_secs(30),
..plan()
};
let draining = {
let lifecycle = Arc::clone(&lifecycle);
tokio::spawn(async move {
drain(lifecycle, signals, move || long_grace, ResolvedPlan::new()).await;
})
};
deliver("SIGTERM");
tokio::time::timeout(Duration::from_secs(5), async {
while lifecycle.phase() != Phase::Draining {
tokio::task::yield_now().await;
}
})
.await
.expect("the first signal begins the drain");
assert!(
lifecycle.admit().is_some(),
"the grace window keeps admitting"
);
deliver("SIGINT");
tokio::time::timeout(Duration::from_secs(5), draining)
.await
.expect("the second signal must not wait out the grace window")
.expect("the drain finished");
assert_eq!(lifecycle.phase(), Phase::Closing);
}
#[tokio::test]
async fn a_zero_grace_window_closes_admission_on_the_first_signal() {
let lifecycle = Arc::new(Lifecycle::new());
let zero_grace = Plan {
drain_grace: Duration::ZERO,
..plan()
};
tokio::time::timeout(
Duration::from_secs(5),
drain(
Arc::clone(&lifecycle),
Signals::once(),
move || zero_grace,
ResolvedPlan::new(),
),
)
.await
.expect("a zero grace window needs no second signal");
assert_eq!(lifecycle.phase(), Phase::Closing);
assert!(lifecycle.admit().is_none());
}
#[tokio::test]
async fn a_signal_after_the_close_leaves_the_sequence_alone() {
let lifecycle = Arc::new(Lifecycle::new());
let (deliver, signals) = Signals::scripted();
let zero_grace = Plan {
drain_grace: Duration::ZERO,
..plan()
};
deliver("SIGTERM");
drain(
Arc::clone(&lifecycle),
signals,
move || zero_grace,
ResolvedPlan::new(),
)
.await;
deliver("SIGTERM");
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(lifecycle.phase(), Phase::Closing);
assert!(lifecycle.admit().is_none());
}
#[test]
fn the_plan_comes_from_the_shutdown_section() {
let plan = Plan::from(&ShutdownConfig::default());
assert!(
!plan.deadline.is_zero(),
"waits must be bounded, not absent"
);
assert!(!plan.flush_timeout.is_zero());
}
}