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> {
crate::server::runtime::runtime_links::shutdown(
supervisor,
drain_timeout,
|| stop_accepting(listener, websocket_listener),
(
&mut crate::server::runtime::runtime_links::ConfiguredLinks::default(),
finalize,
),
)
}
pub(in crate::server) fn stop_accepting(
listener: &mut ServerListener,
websocket_listener: Option<&mut WebSocketListener>,
) -> Result<(), ServerError> {
let websocket_result = websocket_listener.map_or(Ok(()), WebSocketListener::stop_accepting);
let wire_result = listener.stop_accepting();
match (websocket_result, wire_result) {
(Err(first), Err(error)) => {
tracing::error!(%error, "wire listeners also failed to stop accepting");
Err(first)
}
(Err(error), Ok(())) | (Ok(()), Err(error)) => Err(error),
(Ok(()), Ok(())) => Ok(()),
}
}
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)]
#[path = "shutdown_tests.rs"]
mod tests;
#[path = "shutdown_sequence.rs"]
pub(in crate::server) mod sequence;