lib

Core libraries for Radroots
git clone https://radroots.dev/git/lib.git
Log | Files | Refs | README | LICENSE

signals.rs (3291B)

      1 use core::future::Future;
      2 use tokio::signal;
      3 use tracing::info;
      4 
      5 pub async fn shutdown_signal() {
      6     let ctrl_c = async {
      7         signal::ctrl_c()
      8             .await
      9             .expect("failed to install Ctrl+C handler");
     10     };
     11 
     12     #[cfg(unix)]
     13     let terminate = async {
     14         signal::unix::signal(signal::unix::SignalKind::terminate())
     15             .expect("failed to install signal handler")
     16             .recv()
     17             .await;
     18     };
     19 
     20     #[cfg(not(unix))]
     21     let terminate = std::future::pending::<()>();
     22 
     23     wait_for_shutdown(ctrl_c, terminate).await;
     24 }
     25 
     26 async fn wait_for_shutdown<C, T>(ctrl_c: C, terminate: T)
     27 where
     28     C: Future<Output = ()>,
     29     T: Future<Output = ()>,
     30 {
     31     tokio::select! {
     32         _ = ctrl_c => {},
     33         _ = terminate => {},
     34     }
     35 
     36     info!("Shutdown signal received, terminating...");
     37 }
     38 
     39 #[cfg(test)]
     40 mod tests {
     41     use super::{shutdown_signal, wait_for_shutdown};
     42     use core::future::Future;
     43     use core::pin::Pin;
     44     use core::task::{Context, Poll};
     45     use std::process::Command;
     46     use std::time::Duration;
     47 
     48     struct TestFuture {
     49         ready: bool,
     50     }
     51 
     52     impl Future for TestFuture {
     53         type Output = ();
     54 
     55         fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
     56             if self.ready {
     57                 Poll::Ready(())
     58             } else {
     59                 Poll::Pending
     60             }
     61         }
     62     }
     63 
     64     #[tokio::test]
     65     async fn wait_for_shutdown_returns_when_ctrl_completes() {
     66         wait_for_shutdown(TestFuture { ready: true }, TestFuture { ready: false }).await;
     67     }
     68 
     69     #[tokio::test]
     70     async fn wait_for_shutdown_returns_when_terminate_completes() {
     71         wait_for_shutdown(TestFuture { ready: false }, TestFuture { ready: true }).await;
     72     }
     73 
     74     #[tokio::test]
     75     async fn wait_for_shutdown_polls_pending_paths() {
     76         let handle = tokio::task::spawn(wait_for_shutdown(
     77             TestFuture { ready: false },
     78             TestFuture { ready: false },
     79         ));
     80         tokio::task::yield_now().await;
     81         handle.abort();
     82         let _ = handle.await;
     83     }
     84 
     85     #[cfg(unix)]
     86     #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
     87     async fn shutdown_signal_returns_on_sigterm() {
     88         let pid = std::process::id();
     89         let sender = std::thread::spawn(move || {
     90             std::thread::sleep(Duration::from_millis(50));
     91             let status = Command::new("kill")
     92                 .args(["-TERM", pid.to_string().as_str()])
     93                 .status()
     94                 .expect("run kill");
     95             assert!(status.success());
     96         });
     97         shutdown_signal().await;
     98         sender.join().expect("signal sender should join");
     99     }
    100 
    101     #[cfg(unix)]
    102     #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    103     async fn shutdown_signal_returns_on_sigint() {
    104         let pid = std::process::id();
    105         let sender = std::thread::spawn(move || {
    106             std::thread::sleep(Duration::from_millis(50));
    107             let status = Command::new("kill")
    108                 .args(["-INT", pid.to_string().as_str()])
    109                 .status()
    110                 .expect("run kill");
    111             assert!(status.success());
    112         });
    113         shutdown_signal().await;
    114         sender.join().expect("signal sender should join");
    115     }
    116 }