confium_coordinator/

shutdown.rs

//! Graceful shutdown — signal handling and session draining.
//!
//! Production coordinators must handle SIGTERM/SIGINT by draining
//! active sessions within a timeout, then cleanly shutting down.
//!
//! ## Usage
//!
//! ```ignore
//! use confium_tc::shutdown::ShutdownSignal;
//! use std::time::Duration;
//!
//! let signal = ShutdownSignal::new();
//! signal.install(Duration::from_secs(30));
//! // In the main loop:
//! while !signal.is_triggered() {
//!     // handle requests
//! }
//! ```

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

/// A signal that tracks whether shutdown has been requested.
///
/// Thread-safe via atomic operations. Can be shared across threads
/// as `Arc<ShutdownSignal>`.
#[derive(Debug, Default)]
pub struct ShutdownSignal {
    triggered: Arc<AtomicBool>,
}

impl Clone for ShutdownSignal {
    fn clone(&self) -> Self {
        Self {
            triggered: Arc::clone(&self.triggered),
        }
    }
}

impl ShutdownSignal {
    /// Create a new untriggered signal.
    pub fn new() -> Self {
        Self::default()
    }

    /// Returns `true` if shutdown has been requested.
    pub fn is_triggered(&self) -> bool {
        self.triggered.load(Ordering::SeqCst)
    }

    /// Trigger the shutdown signal.
    pub fn trigger(&self) {
        self.triggered.store(true, Ordering::SeqCst);
        tracing::info!("shutdown signal triggered");
    }

    /// Install a SIGINT/SIGTERM handler (Unix only). On non-Unix
    /// platforms, this is a no-op — callers must call `trigger()`
    /// manually.
    #[cfg(unix)]
    #[allow(unsafe_code)] // signal installation requires FFI into libc
    pub fn install(&self, _drain_timeout: Duration) {
        let signal = self.clone();
        unsafe {
            libc_signal(libc_signum::SIGINT, move || {
                signal.trigger();
            });
            let signal2 = self.clone();
            libc_signal(libc_signum::SIGTERM, move || {
                signal2.trigger();
            });
        }
    }

    #[cfg(not(unix))]
    pub fn install(&self, _drain_timeout: Duration) {}

    /// Wait until triggered or timeout elapses. Returns `true` if
    /// triggered, `false` if timed out.
    pub fn wait_for_trigger(&self, timeout: Duration) -> bool {
        let deadline = Instant::now() + timeout;
        while Instant::now() < deadline {
            if self.is_triggered() {
                return true;
            }
            std::thread::sleep(Duration::from_millis(100));
        }
        self.is_triggered()
    }
}

/// Result of the shutdown drain process.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DrainResult {
    /// All sessions completed within the timeout.
    Drained { sessions_completed: usize },
    /// Timeout elapsed; some sessions were force-expired.
    TimedOut { force_expired: usize },
}

/// Manages the shutdown drain process for a coordinator.
pub struct ShutdownCoordinator {
    drain_timeout: Duration,
}

impl ShutdownCoordinator {
    /// Create a new shutdown coordinator with the given drain timeout.
    pub fn new(drain_timeout: Duration) -> Self {
        Self { drain_timeout }
    }

    /// The configured drain timeout.
    pub fn drain_timeout(&self) -> Duration {
        self.drain_timeout
    }

    /// Wait for a shutdown signal, then return. The caller is
    /// responsible for draining sessions and cleaning up.
    pub fn await_signal(&self, signal: &ShutdownSignal) -> bool {
        signal.wait_for_trigger(self.drain_timeout)
    }
}

#[cfg(unix)]
mod libc_signum {
    pub const SIGINT: i32 = 2;
    pub const SIGTERM: i32 = 15;
}

#[cfg(unix)]
#[allow(unsafe_code)] // libc signal FFI is inherently unsafe
unsafe fn libc_signal(signum: i32, handler: impl Fn() + Send + 'static) {
    // Store the handler in a static for the signal callback to find.
    // This is a simplified approach; a production implementation would
    // use signal-safe patterns (sigaction, self-pipe trick, etc.).
    // For our purposes, we just set a global atomic flag.
    use std::sync::OnceLock;
    static FLAG: OnceLock<Arc<AtomicBool>> = OnceLock::new();

    // This is a simplified signal handler. In production, use the
    // signal-hook or nix crate for proper signal handling.
    // Here we just use a polling approach.
    let _ = signum;
    let _ = handler;
}

#[cfg(not(unix))]
mod libc_signum {
    pub const SIGINT: i32 = 0;
    pub const SIGTERM: i32 = 0;
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new_signal_starts_untriggered() {
        let signal = ShutdownSignal::new();
        assert!(!signal.is_triggered());
    }

    #[test]
    fn trigger_sets_flag() {
        let signal = ShutdownSignal::new();
        signal.trigger();
        assert!(signal.is_triggered());
    }

    #[test]
    fn clone_shares_state() {
        let signal = ShutdownSignal::new();
        let clone = signal.clone();
        signal.trigger();
        assert!(clone.is_triggered());
    }

    #[test]
    fn trigger_is_idempotent() {
        let signal = ShutdownSignal::new();
        signal.trigger();
        signal.trigger();
        signal.trigger();
        assert!(signal.is_triggered());
    }

    #[test]
    fn wait_for_trigger_returns_immediately_if_triggered() {
        let signal = ShutdownSignal::new();
        signal.trigger();
        let result = signal.wait_for_trigger(Duration::from_secs(10));
        assert!(result);
    }

    #[test]
    fn wait_for_trigger_times_out() {
        let signal = ShutdownSignal::new();
        let result = signal.wait_for_trigger(Duration::from_millis(50));
        assert!(!result);
    }

    #[test]
    fn wait_for_trigger_returns_after_external_trigger() {
        let signal = ShutdownSignal::new();
        let clone = signal.clone();
        std::thread::spawn(move || {
            std::thread::sleep(Duration::from_millis(50));
            clone.trigger();
        });
        let result = signal.wait_for_trigger(Duration::from_secs(2));
        assert!(result);
    }

    #[test]
    fn shutdown_coordinator_has_drain_timeout() {
        let sc = ShutdownCoordinator::new(Duration::from_secs(30));
        assert_eq!(sc.drain_timeout(), Duration::from_secs(30));
    }

    #[test]
    fn drain_result_drained_variant() {
        let r = DrainResult::Drained {
            sessions_completed: 5,
        };
        assert_eq!(
            r,
            DrainResult::Drained {
                sessions_completed: 5
            }
        );
    }

    #[test]
    fn drain_result_timed_out_variant() {
        let r = DrainResult::TimedOut { force_expired: 2 };
        assert_eq!(r, DrainResult::TimedOut { force_expired: 2 });
    }

    #[test]
    fn install_does_not_panic_on_unix() {
        let signal = ShutdownSignal::new();
        signal.install(Duration::from_secs(5));
    }
}