faucet_server/
shutdown.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2
3use tokio::sync::Notify;
4
5use crate::leak;
6
7const WAIT_STOP_PRINT: std::time::Duration = std::time::Duration::from_secs(5);
8
9pub struct ShutdownSignal {
10    is_shutdown: AtomicBool,
11    notify: Notify,
12}
13
14impl Default for ShutdownSignal {
15    fn default() -> Self {
16        Self::new()
17    }
18}
19
20impl ShutdownSignal {
21    pub fn new() -> Self {
22        ShutdownSignal {
23            is_shutdown: AtomicBool::new(false),
24            notify: Notify::new(),
25        }
26    }
27
28    pub fn shutdown(&self) {
29        self.is_shutdown.store(true, Ordering::Relaxed);
30        self.notify.notify_waiters();
31    }
32
33    pub async fn wait(&self) {
34        if self.is_shutdown.load(Ordering::Relaxed) {
35            return;
36        }
37        // Wait for notification
38        self.notify.notified().await;
39    }
40}
41
42pub fn graceful() -> &'static ShutdownSignal {
43    use crate::global_conn::current_connections;
44
45    let signal = leak!(ShutdownSignal::new()) as &'static ShutdownSignal;
46
47    {
48        ctrlc::set_handler(move || {
49        log::info!(target: "faucet", "Received stop signal, waiting for all users to disconnect");
50        let mut last_5_sec = std::time::Instant::now();
51        while current_connections() > 0 {
52            std::thread::yield_now();
53            if last_5_sec.elapsed() > WAIT_STOP_PRINT {
54                log::info!(
55                    target: "faucet",
56                    "Active connections = {}, waiting for all connections to stop.",
57                    current_connections()
58                );
59                last_5_sec = std::time::Instant::now();
60            }
61        }
62        signal.shutdown();
63    })
64    .expect("Unable to set term handler. This is a bug");
65    }
66
67    signal
68}
69
70pub fn immediate() -> &'static ShutdownSignal {
71    let signal = leak!(ShutdownSignal::new()) as &'static ShutdownSignal;
72    {
73        ctrlc::set_handler(move || {
74            log::info!(target: "faucet", "Starting immediate shutdown handle");
75            signal.shutdown()
76        })
77        .expect("Unable to set term handler. This is a bug");
78    }
79    signal
80}
faucet_server/shutdown.rs

faucet_server/
shutdown.rs
