Ferrit Explore
中文·繁體·EN·日本語 Sign in Register
cielxl / veld / src / util / signal.rs
use tokio::sync::broadcast;

/// Signal types the server handles
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ServerSignal {
    /// Reload configuration (SIGHUP on Unix, service control on Windows)
    Reload,
    /// Graceful shutdown (SIGTERM/SIGINT on Unix)
    Shutdown,
    /// Reopen log files (SIGUSR1 on Unix)
    ReopenLogs,
}

/// Cross-platform signal handler
pub struct SignalHandler {
    tx: broadcast::Sender<ServerSignal>,
}

impl Default for SignalHandler {
    fn default() -> Self {
        Self::new()
    }
}

impl SignalHandler {
    pub fn new() -> Self {
        let (tx, _) = broadcast::channel(16);
        Self { tx }
    }

    pub fn subscribe(&self) -> broadcast::Receiver<ServerSignal> {
        self.tx.subscribe()
    }

    /// Start listening for OS signals
    pub async fn run(&self) {
        #[cfg(unix)]
        self.run_unix().await;

        #[cfg(windows)]
        self.run_windows().await;
    }

    #[cfg(unix)]
    async fn run_unix(&self) {
        use signal_hook::consts::{SIGHUP, SIGINT, SIGTERM, SIGUSR1};
        use signal_hook::iterator::Signals;

        let mut signals = Signals::new([SIGHUP, SIGINT, SIGTERM, SIGUSR1])
            .expect("Failed to register signal handlers");

        let tx = self.tx.clone();
        tokio::spawn(async move {
            loop {
                // Use tokio's spawn_blocking for signal iteration
                let signal = tokio::task::spawn_blocking({
                    // Create a new signals handle for each iteration
                    let mut sigs = Signals::new([SIGHUP, SIGINT, SIGTERM, SIGUSR1])
                        .expect("Failed to register signal handlers");
                    move || sigs.forever().next()
                })
                .await
                .ok()
                .flatten();

                match signal {
                    Some(SIGHUP) => {
                        tracing::info!("Received SIGHUP, reloading configuration");
                        let _ = tx.send(ServerSignal::Reload);
                    }
                    Some(SIGTERM) | Some(SIGINT) => {
                        tracing::info!("Received shutdown signal, initiating graceful shutdown");
                        let _ = tx.send(ServerSignal::Shutdown);
                        break;
                    }
                    Some(SIGUSR1) => {
                        tracing::info!("Received SIGUSR1, reopening log files");
                        let _ = tx.send(ServerSignal::ReopenLogs);
                    }
                    _ => {}
                }
            }
        });
    }

    #[cfg(windows)]
    async fn run_windows(&self) {
        use tokio::signal::windows;

        let mut ctrl_c = windows::ctrl_c().expect("Failed to register Ctrl+C handler");
        let mut ctrl_break = windows::ctrl_break().expect("Failed to register Ctrl+Break handler");
        let mut ctrl_close = windows::ctrl_close().expect("Failed to register Ctrl+Close handler");
        let mut ctrl_shutdown =
            windows::ctrl_shutdown().expect("Failed to register Ctrl+Shutdown handler");

        let tx = self.tx.clone();
        tokio::spawn(async move {
            // Wait for whichever shutdown signal arrives first, then exit.
            tokio::select! {
                _ = ctrl_c.recv() => {
                    tracing::info!("Received Ctrl+C, initiating graceful shutdown");
                    let _ = tx.send(ServerSignal::Shutdown);
                }
                _ = ctrl_break.recv() => {
                    tracing::info!("Received Ctrl+Break, initiating graceful shutdown");
                    let _ = tx.send(ServerSignal::Shutdown);
                }
                _ = ctrl_close.recv() => {
                    tracing::info!("Received Ctrl+Close, initiating graceful shutdown");
                    let _ = tx.send(ServerSignal::Shutdown);
                }
                _ = ctrl_shutdown.recv() => {
                    tracing::info!("Received shutdown signal, initiating graceful shutdown");
                    let _ = tx.send(ServerSignal::Shutdown);
                }
            }
        });
    }
}