turbovault_plugin_api/
lifecycle.rs

1use std::sync::Arc;
2
3use tokio::sync::watch;
4
5/// Cooperative shutdown signal handed to a plugin's background work.
6///
7/// A plugin that maintains derived state generally wants a worker: something
8/// that drains the hook bus, re-embeds changed notes, compacts an index. That
9/// work outlives any single tool call, so it needs a way to learn the host is
10/// going away — otherwise a detached task keeps running against state that is
11/// being torn down, and the process cannot exit cleanly.
12///
13/// Cheap to clone; every clone observes the same signal.
14#[derive(Clone)]
15pub struct ShutdownSignal {
16    receiver: watch::Receiver<bool>,
17}
18
19impl std::fmt::Debug for ShutdownSignal {
20    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        formatter
22            .debug_struct("ShutdownSignal")
23            .field("is_shutting_down", &self.is_shutting_down())
24            .finish()
25    }
26}
27
28impl ShutdownSignal {
29    /// Whether shutdown has already been requested.
30    ///
31    /// Check this at the top of a work loop; use [`Self::cancelled`] to wait.
32    pub fn is_shutting_down(&self) -> bool {
33        *self.receiver.borrow()
34    }
35
36    /// Resolve when shutdown is requested, immediately if it already was.
37    ///
38    /// Select on this alongside the work a loop would otherwise block on:
39    ///
40    /// ```ignore
41    /// loop {
42    ///     tokio::select! {
43    ///         _ = shutdown.cancelled() => break,
44    ///         event = events.recv() => { /* ... */ }
45    ///     }
46    /// }
47    /// ```
48    pub async fn cancelled(&self) {
49        let mut receiver = self.receiver.clone();
50        // `borrow_and_update` marks the current value seen, so the wait below
51        // reacts to the next change rather than an already-observed one.
52        if *receiver.borrow_and_update() {
53            return;
54        }
55        // A send error means the host dropped the trigger without signalling,
56        // which can only happen if the host is already gone — treat it as
57        // shutdown rather than waiting forever.
58        let _ = receiver.changed().await;
59    }
60}
61
62/// Host-side handle that requests shutdown.
63///
64/// Held by the server, never handed across the plugin boundary.
65pub struct ShutdownTrigger {
66    sender: Arc<watch::Sender<bool>>,
67    /// Keeps the channel alive when no plugin holds a signal.
68    ///
69    /// Without a live receiver the state would be lost, and a shutdown
70    /// requested before any worker started — an early failure during startup,
71    /// say — would leave a worker that starts afterwards waiting forever for a
72    /// signal that already fired.
73    _keepalive: Arc<watch::Receiver<bool>>,
74}
75
76impl std::fmt::Debug for ShutdownTrigger {
77    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        formatter
79            .debug_struct("ShutdownTrigger")
80            .finish_non_exhaustive()
81    }
82}
83
84impl Default for ShutdownTrigger {
85    fn default() -> Self {
86        Self::new()
87    }
88}
89
90impl ShutdownTrigger {
91    /// Create a trigger that has not yet fired.
92    pub fn new() -> Self {
93        let (sender, keepalive) = watch::channel(false);
94        Self {
95            sender: Arc::new(sender),
96            _keepalive: Arc::new(keepalive),
97        }
98    }
99
100    /// Hand out a signal observing this trigger.
101    pub fn signal(&self) -> ShutdownSignal {
102        ShutdownSignal {
103            receiver: self.sender.subscribe(),
104        }
105    }
106
107    /// Request shutdown. Idempotent.
108    pub fn shutdown(&self) {
109        // `send_replace` rather than `send`: the latter reports an error when
110        // no receiver exists, and this state must be recorded whether or not
111        // anyone is listening yet.
112        self.sender.send_replace(true);
113    }
114}
115
116impl Clone for ShutdownTrigger {
117    fn clone(&self) -> Self {
118        Self {
119            sender: Arc::clone(&self.sender),
120            _keepalive: Arc::clone(&self._keepalive),
121        }
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[tokio::test]
130    async fn waiting_workers_are_released_on_shutdown() {
131        let trigger = ShutdownTrigger::new();
132        let signal = trigger.signal();
133        assert!(!signal.is_shutting_down());
134
135        let waiter = tokio::spawn({
136            let signal = signal.clone();
137            async move {
138                signal.cancelled().await;
139                true
140            }
141        });
142
143        trigger.shutdown();
144        assert!(waiter.await.expect("waiter task"));
145        assert!(signal.is_shutting_down());
146    }
147
148    #[tokio::test]
149    async fn a_signal_taken_after_shutdown_resolves_immediately() {
150        let trigger = ShutdownTrigger::new();
151        trigger.shutdown();
152        let signal = trigger.signal();
153        assert!(signal.is_shutting_down());
154        // Must not hang: the worker started too late to see the transition.
155        signal.cancelled().await;
156    }
157}
turbovault_plugin_api/lifecycle.rs

turbovault_plugin_api/
lifecycle.rs
