Skip to main content

rustbridge_runtime/
shutdown.rs

1//! Graceful shutdown support
2
3use std::sync::Arc;
4use std::sync::atomic::{AtomicBool, Ordering};
5use tokio::sync::broadcast;
6
7/// Handle for triggering shutdown
8#[derive(Clone)]
9pub struct ShutdownHandle {
10    triggered: Arc<AtomicBool>,
11    sender: broadcast::Sender<()>,
12}
13
14impl ShutdownHandle {
15    /// Create a new shutdown handle
16    pub fn new() -> Self {
17        let (sender, _) = broadcast::channel(1);
18        Self {
19            triggered: Arc::new(AtomicBool::new(false)),
20            sender,
21        }
22    }
23
24    /// Trigger shutdown
25    pub fn trigger(&self) {
26        if !self.triggered.swap(true, Ordering::SeqCst) {
27            // Only send if this is the first trigger
28            let _ = self.sender.send(());
29        }
30    }
31
32    /// Check if shutdown has been triggered
33    pub fn is_triggered(&self) -> bool {
34        self.triggered.load(Ordering::SeqCst)
35    }
36
37    /// Get a signal that can be used to detect shutdown
38    pub fn signal(&self) -> ShutdownSignal {
39        ShutdownSignal {
40            triggered: self.triggered.clone(),
41            receiver: self.sender.subscribe(),
42        }
43    }
44}
45
46impl Default for ShutdownHandle {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52/// Signal for detecting shutdown (cloneable, can be passed to tasks)
53pub struct ShutdownSignal {
54    triggered: Arc<AtomicBool>,
55    receiver: broadcast::Receiver<()>,
56}
57
58impl ShutdownSignal {
59    /// Check if shutdown has been triggered (non-blocking)
60    pub fn is_triggered(&self) -> bool {
61        self.triggered.load(Ordering::SeqCst)
62    }
63
64    /// Wait for shutdown to be triggered
65    ///
66    /// Returns immediately if already triggered.
67    pub async fn wait(&mut self) {
68        if self.is_triggered() {
69            return;
70        }
71        let _ = self.receiver.recv().await;
72    }
73
74    /// Create a future that completes when shutdown is triggered
75    ///
76    /// This is useful for select! statements.
77    pub async fn notified(&mut self) {
78        self.wait().await;
79    }
80}
81
82impl Clone for ShutdownSignal {
83    fn clone(&self) -> Self {
84        Self {
85            triggered: self.triggered.clone(),
86            receiver: self.receiver.resubscribe(),
87        }
88    }
89}
90
91#[cfg(test)]
92#[path = "shutdown/shutdown_tests.rs"]
93mod shutdown_tests;