rs_zero/core/
shutdown.rs

1use std::{future::Future, pin::Pin, sync::Arc};
2
3use tokio::sync::Notify;
4
5/// Cloneable shutdown token used by [`crate::core::ServiceGroup`].
6#[derive(Debug, Clone)]
7pub struct ShutdownToken {
8    inner: Arc<ShutdownState>,
9}
10
11#[derive(Debug)]
12struct ShutdownState {
13    notify: Notify,
14    cancelled: std::sync::atomic::AtomicBool,
15}
16
17impl Default for ShutdownToken {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl ShutdownToken {
24    /// Creates a token in the running state.
25    pub fn new() -> Self {
26        Self {
27            inner: Arc::new(ShutdownState {
28                notify: Notify::new(),
29                cancelled: std::sync::atomic::AtomicBool::new(false),
30            }),
31        }
32    }
33
34    /// Requests shutdown. Calling this method multiple times is safe.
35    pub fn cancel(&self) {
36        if !self
37            .inner
38            .cancelled
39            .swap(true, std::sync::atomic::Ordering::SeqCst)
40        {
41            self.inner.notify.notify_waiters();
42        }
43    }
44
45    /// Returns whether shutdown has been requested.
46    pub fn is_cancelled(&self) -> bool {
47        self.inner
48            .cancelled
49            .load(std::sync::atomic::Ordering::SeqCst)
50    }
51
52    /// Waits until shutdown is requested.
53    pub fn cancelled(&self) -> ShutdownFuture<'_> {
54        Box::pin(async move {
55            loop {
56                let notified = self.inner.notify.notified();
57                if self.is_cancelled() {
58                    return;
59                }
60                notified.await;
61            }
62        })
63    }
64}
65
66/// Future returned by [`ShutdownToken::cancelled`].
67pub type ShutdownFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
68
69/// Waits until the process receives Ctrl-C.
70///
71/// Services can pass this future to server helpers that support graceful
72/// shutdown.
73pub async fn shutdown_signal() {
74    if let Err(error) = tokio::signal::ctrl_c().await {
75        tracing::warn!(%error, "failed to listen for ctrl-c");
76    }
77}
rs_zero/core/shutdown.rs

rs_zero/core/
shutdown.rs
