faucet_server/
shutdown.rs

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

faucet_server/
shutdown.rs
