use std::fmt;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
use signal_hook::consts::signal::{SIGINT, SIGTERM};
use signal_hook::iterator::{Handle as SignalIteratorHandle, Signals};
use crate::ServerError;
use crate::server::connection::{ConnectionSupervisor, WebSocketListener};
use crate::server::listener::ServerListener;
pub(crate) const WEDGED_CONNECTION_STOP: Duration = Duration::from_millis(500);
#[derive(Clone)]
pub struct ShutdownHandle {
inner: Arc<ShutdownState>,
}
impl ShutdownHandle {
#[must_use]
pub fn new() -> Self {
Self {
inner: Arc::new(ShutdownState::new()),
}
}
pub fn initiate(&self) -> bool {
if self.inner.initiated.swap(true, Ordering::SeqCst) {
tracing::debug!("shutdown request ignored because shutdown is already active");
return false;
}
tracing::info!("shutdown requested");
self.inner.notify();
true
}
pub fn wait(&self) {
#[cfg(test)]
let _returned = Returned(&self.inner);
if self.is_initiated() {
return;
}
let Ok(mut guard) = self.inner.wait_lock.lock() else {
tracing::error!(
"shutdown wait lock poisoned; returning from wait as if shutdown was requested"
);
return;
};
while !self.is_initiated() {
#[cfg(test)]
self.inner.signal_parked();
let Ok(next_guard) = self.inner.waiter.wait(guard) else {
tracing::error!(
"shutdown wait condvar reacquired a poisoned lock; returning from wait as if shutdown was requested"
);
return;
};
guard = next_guard;
}
}
#[must_use]
pub fn is_initiated(&self) -> bool {
self.inner.initiated.load(Ordering::SeqCst)
}
}
impl Default for ShutdownHandle {
fn default() -> Self {
Self::new()
}
}
impl fmt::Debug for ShutdownHandle {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ShutdownHandle")
.field("initiated", &self.is_initiated())
.finish()
}
}
#[derive(Debug)]
struct ShutdownState {
initiated: AtomicBool,
wait_lock: Mutex<()>,
waiter: Condvar,
#[cfg(test)]
parked: std::sync::atomic::AtomicUsize,
#[cfg(test)]
returned: std::sync::atomic::AtomicUsize,
#[cfg(test)]
notified: std::sync::atomic::AtomicUsize,
#[cfg(test)]
signal: Condvar,
#[cfg(test)]
signal_lock: Mutex<()>,
}
#[cfg(test)]
struct Returned<'a>(&'a ShutdownState);
#[cfg(test)]
impl Drop for Returned<'_> {
fn drop(&mut self) {
self.0.count_and_signal(&self.0.returned);
}
}
impl ShutdownState {
const fn new() -> Self {
Self {
initiated: AtomicBool::new(false),
wait_lock: Mutex::new(()),
waiter: Condvar::new(),
#[cfg(test)]
parked: std::sync::atomic::AtomicUsize::new(0),
#[cfg(test)]
returned: std::sync::atomic::AtomicUsize::new(0),
#[cfg(test)]
notified: std::sync::atomic::AtomicUsize::new(0),
#[cfg(test)]
signal: Condvar::new(),
#[cfg(test)]
signal_lock: Mutex::new(()),
}
}
#[cfg(test)]
fn signal_parked(&self) {
self.count_and_signal(&self.parked);
}
#[cfg(test)]
fn count_and_signal(&self, count: &std::sync::atomic::AtomicUsize) {
let guard = self.signal_lock.lock();
count.fetch_add(1, Ordering::SeqCst);
self.signal.notify_all();
drop(guard);
}
fn notify(&self) {
match self.wait_lock.lock() {
Ok(_guard) => self.waiter.notify_all(),
Err(poisoned) => {
tracing::error!("shutdown wait lock poisoned; notifying waiters anyway");
let _guard = poisoned.into_inner();
self.waiter.notify_all();
}
}
#[cfg(test)]
self.notified.fetch_add(1, Ordering::SeqCst);
}
}
#[derive(Debug)]
pub struct SignalShutdownRegistration {
signal_handle: SignalIteratorHandle,
worker: Option<JoinHandle<()>>,
}
impl SignalShutdownRegistration {
const fn new(signal_handle: SignalIteratorHandle, worker: JoinHandle<()>) -> Self {
Self {
signal_handle,
worker: Some(worker),
}
}
}
impl Drop for SignalShutdownRegistration {
fn drop(&mut self) {
self.signal_handle.close();
let Some(worker) = self.worker.take() else {
return;
};
if worker.join().is_err() {
tracing::warn!("shutdown signal worker terminated unexpectedly");
}
}
}
pub fn register_signal_handlers(
handle: ShutdownHandle,
) -> Result<SignalShutdownRegistration, ServerError> {
let mut signals =
Signals::new([SIGTERM, SIGINT]).map_err(|error| ServerError::ListenerAccept {
message: format!("failed to register shutdown signal handlers: {error}"),
})?;
let signal_handle = signals.handle();
let worker = thread::spawn(move || {
for signal in signals.forever() {
tracing::info!(signal, "received shutdown signal");
handle.initiate();
}
});
Ok(SignalShutdownRegistration::new(signal_handle, worker))
}
pub fn run_shutdown_sequence(
listener: &mut ServerListener,
websocket_listener: Option<&mut WebSocketListener>,
supervisor: &ConnectionSupervisor,
drain_timeout: Duration,
) -> Result<(), ServerError> {
run_shutdown_sequence_before_flush(
listener,
websocket_listener,
supervisor,
drain_timeout,
|| Ok(()),
)
}
pub(in crate::server) fn run_shutdown_sequence_before_flush(
listener: &mut ServerListener,
websocket_listener: Option<&mut WebSocketListener>,
supervisor: &ConnectionSupervisor,
drain_timeout: Duration,
finalize: impl FnOnce() -> Result<(), ServerError>,
) -> Result<(), ServerError> {
let started = Instant::now();
tracing::info!(
ignored_drain_timeout = ?drain_timeout,
"starting shutdown sequence; the configured drain timeout has been ignored since 0.14.3 \
because every write is durable before it is acknowledged, so no request is in flight to \
drain"
);
if let Some(websocket_listener) = websocket_listener {
websocket_listener.stop_accepting()?;
}
listener.stop_accepting()?;
if !supervisor.wait_for_delivery_quiesced(Instant::now() + WEDGED_CONNECTION_STOP) {
tracing::warn!(
stop = ?WEDGED_CONNECTION_STOP,
"delivery flush barrier completion was not observed before the wedged-connection stop; \
proceeding to shutdown notification"
);
}
supervisor.notify_shutdown_subscribers();
supervisor.force_close_active_connections();
wait_after_force_close(supervisor);
let settled = supervisor.settle_parked_connection_fates();
supervisor.shutdown();
let finalized = finalize();
let flushed = flush_durable_state(supervisor);
tracing::info!(
elapsed_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX),
clean = settled.is_ok() && finalized.is_ok() && flushed.is_ok(),
"shutdown sequence complete"
);
let mut failure = None;
for result in [finalized, flushed, settled] {
if let Err(error) = result {
tracing::error!(%error, "shutdown phase failed");
if failure.is_none() {
failure = Some(error);
}
}
}
failure.map_or(Ok(()), Err)
}
pub(crate) fn wait_after_force_close(supervisor: &ConnectionSupervisor) {
let deadline = Instant::now() + WEDGED_CONNECTION_STOP;
if supervisor.wait_for_connections_drained(deadline) {
return;
}
let remaining = supervisor.active_connection_count();
if remaining > 0 {
tracing::warn!(
active_connections = remaining,
stop = ?WEDGED_CONNECTION_STOP,
"connections remained active after the wedged-connection stop"
);
}
}
fn flush_durable_state(supervisor: &ConnectionSupervisor) -> Result<(), ServerError> {
tracing::info!("flushing durable channel state");
supervisor.flush_durable_state().map_err(|error| {
tracing::error!(%error, "durable state flush failed during shutdown");
match error {
ServerError::ShutdownFlush { .. } => error,
other => ServerError::ShutdownFlush {
message: other.to_string(),
},
}
})?;
tracing::info!("durable channel state flushed");
Ok(())
}
#[cfg(test)]
#[path = "shutdown_inflight_tests.rs"]
mod shutdown_inflight_tests;
#[cfg(test)]
mod tests {
use std::error::Error;
use std::sync::atomic::Ordering;
use std::thread;
use std::time::Duration;
use super::{ShutdownHandle, wait_after_force_close};
use crate::server::connection::ConnectionSupervisor;
#[test]
fn shutdown_handle_initiates_once() {
let handle = ShutdownHandle::new();
assert!(!handle.is_initiated());
assert!(handle.initiate());
assert!(handle.is_initiated());
assert!(!handle.initiate());
}
const SIGNAL_LOCK_POISONED: &str =
"the test signal lock was poisoned: another test thread panicked holding it";
fn await_parked(handle: &ShutdownHandle) -> Result<(), Box<dyn Error>> {
let state = &handle.inner;
let mut guard = state.signal_lock.lock().map_err(|_| SIGNAL_LOCK_POISONED)?;
loop {
if state.returned.load(Ordering::SeqCst) > 0 {
return Err("the waiter returned from wait() without parking".into());
}
if state.parked.load(Ordering::SeqCst) > 0 {
drop(guard);
return Ok(());
}
guard = state.signal.wait(guard).map_err(|_| SIGNAL_LOCK_POISONED)?;
}
}
fn await_returned(handle: &ShutdownHandle, parked_before: usize) -> Result<(), Box<dyn Error>> {
let state = &handle.inner;
if state.notified.load(Ordering::SeqCst) == 0 {
return Err("initiate() never notified the parked waiter".into());
}
let mut guard = state.signal_lock.lock().map_err(|_| SIGNAL_LOCK_POISONED)?;
while state.returned.load(Ordering::SeqCst) == 0 {
if state.parked.load(Ordering::SeqCst) > parked_before {
return Err("the waiter woke and parked again: it did not see the shutdown".into());
}
guard = state.signal.wait(guard).map_err(|_| SIGNAL_LOCK_POISONED)?;
}
drop(guard);
Ok(())
}
#[test]
fn shutdown_handle_wait_unblocks_on_initiate() -> Result<(), Box<dyn std::error::Error>> {
let handle = ShutdownHandle::new();
let waiter = handle.clone();
let worker = thread::spawn(move || {
waiter.wait();
waiter.is_initiated()
});
await_parked(&handle)?;
let parked = handle.inner.parked.load(Ordering::SeqCst);
assert!(handle.initiate());
await_returned(&handle, parked)?;
let observed = worker.join().map_err(|_| "wait worker panicked")?;
assert!(observed);
Ok(())
}
fn wait_lock_of(state: &super::ShutdownState) -> &std::sync::Mutex<()> {
&state.wait_lock
}
fn signal_lock_of(state: &super::ShutdownState) -> &std::sync::Mutex<()> {
&state.signal_lock
}
#[allow(clippy::panic)]
fn poison(handle: &ShutdownHandle, lock: fn(&super::ShutdownState) -> &std::sync::Mutex<()>) {
let poisoner = handle.clone();
let outcome = thread::spawn(move || {
let _guard = lock(&poisoner.inner).lock();
panic!("poison a shutdown lock");
})
.join();
assert!(outcome.is_err(), "the poisoner thread must have panicked");
}
fn poison_wait_lock(handle: &ShutdownHandle) {
poison(handle, wait_lock_of);
}
#[test]
fn a_poisoned_test_signal_lock_is_refused_by_name() {
let handle = ShutdownHandle::new();
poison(&handle, signal_lock_of);
assert!(handle.initiate());
let refused =
await_returned(&handle, 0).map_or_else(|error| error.to_string(), |()| String::new());
assert_eq!(refused, SIGNAL_LOCK_POISONED);
let refused =
await_parked(&handle).map_or_else(|error| error.to_string(), |()| String::new());
assert_eq!(refused, SIGNAL_LOCK_POISONED);
}
#[test]
fn wait_on_poisoned_lock_returns_and_logs_error() {
let handle = ShutdownHandle::new();
poison_wait_lock(&handle);
let log = crate::test_log::CapturedLog::default();
log.capture(|| handle.wait());
let output = log.contents();
assert!(
output.contains("ERROR"),
"a poisoned wait lock must be logged at error level, got: {output}"
);
}
#[test]
fn notify_on_poisoned_lock_still_wakes_parked_waiter() -> Result<(), Box<dyn std::error::Error>>
{
let handle = ShutdownHandle::new();
let waiter = handle.clone();
let worker = thread::spawn(move || waiter.wait());
await_parked(&handle)?;
poison_wait_lock(&handle);
let parked = handle.inner.parked.load(Ordering::SeqCst);
assert!(handle.initiate());
await_returned(&handle, parked)?;
worker.join().map_err(|_| "wait worker panicked")?;
Ok(())
}
#[test]
fn the_close_settle_returns_immediately_when_no_connections_are_active()
-> Result<(), Box<dyn std::error::Error>> {
let supervisor = ConnectionSupervisor::new()?;
let started = std::time::Instant::now();
wait_after_force_close(&supervisor);
let elapsed = started.elapsed();
assert!(
elapsed < Duration::from_millis(50),
"the close settle took {elapsed:?} with no connections tracked"
);
supervisor.shutdown();
Ok(())
}
#[test]
fn shutdown_source_has_no_drain_and_no_reap_count_sleep_loop() {
let source = include_str!("shutdown.rs");
let implementation = source.split("mod tests").next().unwrap_or(source);
for forbidden in [
"DRAIN_PROGRESS_INTERVAL",
"FORCE_CLOSE_SETTLE_TIMEOUT",
"FORCE_CLOSE_POLL_INTERVAL",
"reap_crashed_connections",
"fn drain_connections",
] {
assert!(
!implementation.contains(forbidden),
"retired poll/reap/drain token `{forbidden}` must not appear in the shutdown implementation"
);
}
}
#[test]
fn the_configured_drain_timeout_is_never_turned_into_a_deadline() {
let source = include_str!("shutdown.rs");
let implementation = source.split("mod tests").next().unwrap_or(source);
assert!(
!implementation.contains("+ drain_timeout"),
"drain_timeout must never be added to an Instant to form a shutdown deadline"
);
assert!(
implementation.contains("ignored_drain_timeout = ?drain_timeout"),
"drain_timeout must still be named in the shutdown log line that reports it ignored"
);
}
}