Line data    Source code

       1              : pub use signal_hook::consts::TERM_SIGNALS;
       2              : pub use signal_hook::consts::signal::*;
       3              : use signal_hook::iterator::Signals;
       4              : use tokio::signal::unix::{SignalKind, signal};
       5              : use tracing::info;
       6              : 
       7              : pub enum Signal {
       8              :     Quit,
       9              :     Interrupt,
      10              :     Terminate,
      11              : }
      12              : 
      13              : impl Signal {
      14            0 :     pub fn name(&self) -> &'static str {
      15            0 :         match self {
      16            0 :             Signal::Quit => "SIGQUIT",
      17            0 :             Signal::Interrupt => "SIGINT",
      18            0 :             Signal::Terminate => "SIGTERM",
      19              :         }
      20            0 :     }
      21              : }
      22              : 
      23              : pub struct ShutdownSignals;
      24              : 
      25              : impl ShutdownSignals {
      26            0 :     pub fn handle(mut handler: impl FnMut(Signal) -> anyhow::Result<()>) -> anyhow::Result<()> {
      27            0 :         for raw_signal in Signals::new(TERM_SIGNALS)?.into_iter() {
      28            0 :             let signal = match raw_signal {
      29            0 :                 SIGINT => Signal::Interrupt,
      30            0 :                 SIGTERM => Signal::Terminate,
      31            0 :                 SIGQUIT => Signal::Quit,
      32            0 :                 other => panic!("unknown signal: {}", other),
      33              :             };
      34              : 
      35            0 :             handler(signal)?;
      36              :         }
      37              : 
      38            0 :         Ok(())
      39            0 :     }
      40              : }
      41              : 
      42              : /// Runs in a loop since we want to be responsive to multiple signals
      43              : /// even after triggering shutdown (e.g. a SIGQUIT after a slow SIGTERM shutdown)
      44              : /// <https://github.com/neondatabase/neon/issues/9740>
      45            0 : pub async fn signal_handler(token: tokio_util::sync::CancellationToken) {
      46            0 :     let mut sigint = signal(SignalKind::interrupt()).unwrap();
      47            0 :     let mut sigterm = signal(SignalKind::terminate()).unwrap();
      48            0 :     let mut sigquit = signal(SignalKind::quit()).unwrap();
      49              : 
      50              :     loop {
      51            0 :         let signal = tokio::select! {
      52            0 :             _ = sigquit.recv() => {
      53            0 :                 info!("Got signal SIGQUIT. Terminating in immediate shutdown mode.");
      54            0 :                 std::process::exit(111);
      55              :             }
      56            0 :             _ = sigint.recv() => "SIGINT",
      57            0 :             _ = sigterm.recv() => "SIGTERM",
      58              :         };
      59              : 
      60            0 :         if !token.is_cancelled() {
      61            0 :             info!("Got signal {signal}. Terminating gracefully in fast shutdown mode.");
      62            0 :             token.cancel();
      63              :         } else {
      64            0 :             info!("Got signal {signal}. Already shutting down.");
      65              :         }
      66              :     }
      67              : }