botkit_core/
shutdown.rs

1use std::future::pending;
2
3/// Signals a running bot to stop
4///
5/// Cloneable and cheap; every clone signals the same bot. Dropping every clone
6/// does *not* stop the bot — call [`ShutdownSignal::shutdown`] for that.
7#[derive(Clone, Debug)]
8pub struct ShutdownSignal {
9    tx: async_channel::Sender<()>,
10}
11
12impl ShutdownSignal {
13    /// Ask the bot to stop
14    ///
15    /// Returns immediately; the bot stops once it reaches its next await point.
16    /// Calling this more than once is harmless.
17    pub fn shutdown(&self) {
18        self.tx.close();
19    }
20
21    /// Whether shutdown has already been requested
22    pub fn is_shutdown(&self) -> bool {
23        self.tx.is_closed()
24    }
25}
26
27/// The receiving half of a shutdown signal, handed to [`crate::Bot::run_until`]
28#[derive(Clone, Debug)]
29pub struct Shutdown {
30    rx: Option<async_channel::Receiver<()>>,
31}
32
33impl Shutdown {
34    /// Create a linked signal/receiver pair
35    pub fn channel() -> (ShutdownSignal, Self) {
36        let (tx, rx) = async_channel::bounded(1);
37        (ShutdownSignal { tx }, Self { rx: Some(rx) })
38    }
39
40    /// A shutdown that never fires, for bots meant to run until the process exits
41    pub fn never() -> Self {
42        Self { rx: None }
43    }
44
45    /// Whether shutdown has already been requested
46    pub fn is_shutdown(&self) -> bool {
47        self.rx.as_ref().is_some_and(|rx| rx.is_closed())
48    }
49
50    /// Resolve once shutdown is requested
51    ///
52    /// Never resolves for [`Shutdown::never`], so it is safe to select on
53    /// unconditionally.
54    pub async fn wait(&self) {
55        match &self.rx {
56            // `recv` resolves with `Err` as soon as the sender is closed.
57            Some(rx) => while rx.recv().await.is_ok() {},
58            None => pending().await,
59        }
60    }
61}
62
63impl Default for Shutdown {
64    fn default() -> Self {
65        Self::never()
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72    use futures_lite::future::{block_on, poll_once};
73
74    #[test]
75    fn wait_resolves_after_shutdown() {
76        let (signal, shutdown) = Shutdown::channel();
77        assert!(!shutdown.is_shutdown());
78        assert!(block_on(poll_once(shutdown.wait())).is_none());
79
80        signal.shutdown();
81
82        assert!(signal.is_shutdown());
83        assert!(shutdown.is_shutdown());
84        assert!(block_on(poll_once(shutdown.wait())).is_some());
85    }
86
87    #[test]
88    fn never_stays_pending() {
89        let shutdown = Shutdown::never();
90        assert!(!shutdown.is_shutdown());
91        assert!(block_on(poll_once(shutdown.wait())).is_none());
92    }
93
94    #[test]
95    fn shutdown_is_idempotent() {
96        let (signal, shutdown) = Shutdown::channel();
97        signal.shutdown();
98        signal.clone().shutdown();
99        assert!(block_on(poll_once(shutdown.wait())).is_some());
100    }
101}
botkit_core/shutdown.rs

botkit_core/
shutdown.rs
