Skip to main content

nu_protocol/pipeline/
handlers.rs

1use core::fmt;
2use std::fmt::Display;
3use std::sync::{Arc, Mutex};
4
5use crate::{ShellError, SignalAction, engine::Sequence};
6
7/// Handler is a closure that can be sent across threads and shared.
8pub type Handler = Box<dyn Fn(SignalAction) + Send + Sync>;
9
10/// Manages a collection of handlers.
11#[derive(Clone, derive_more::Debug, Default)]
12pub struct Handlers {
13    /// List of handler tuples containing an ID and the handler itself.
14    #[debug("{}", debug_fmt_handlers(&self.handlers))]
15    handlers: Arc<Mutex<Vec<(usize, Handler)>>>,
16    /// Sequence generator for unique IDs.
17    next_id: Arc<Sequence>,
18}
19
20/// HandlerGuard that unregisters a handler when dropped.
21#[derive(Clone, derive_more::Debug)]
22pub struct HandlerGuard {
23    /// Unique ID of the handler.
24    id: usize,
25    /// Reference to the handlers list.
26    #[debug("{}", debug_fmt_handlers(&self.handlers))]
27    handlers: Arc<Mutex<Vec<(usize, Handler)>>>,
28}
29
30impl Drop for HandlerGuard {
31    /// Drops the `Guard`, removing the associated handler from the list.
32    fn drop(&mut self) {
33        if let Ok(mut handlers) = self.handlers.lock() {
34            handlers.retain(|(id, _)| *id != self.id);
35        }
36    }
37}
38
39impl Handlers {
40    pub fn new() -> Handlers {
41        Self::default()
42    }
43
44    /// Registers a new handler and returns an RAII guard which will unregister the handler when
45    /// dropped.
46    pub fn register(&self, handler: Handler) -> Result<HandlerGuard, ShellError> {
47        let id = self.next_id.next()?;
48        if let Ok(mut handlers) = self.handlers.lock() {
49            handlers.push((id, handler));
50        }
51
52        Ok(HandlerGuard {
53            id,
54            handlers: Arc::clone(&self.handlers),
55        })
56    }
57
58    /// Registers a new handler which persists for the entire process lifetime.
59    ///
60    /// Only use this for handlers which should exist for the lifetime of the program.
61    /// You should prefer to use `register` with a `HandlerGuard` when possible.
62    pub fn register_unguarded(&self, handler: Handler) -> Result<(), ShellError> {
63        let id = self.next_id.next()?;
64        if let Ok(mut handlers) = self.handlers.lock() {
65            handlers.push((id, handler));
66        }
67
68        Ok(())
69    }
70
71    /// Runs all registered handlers.
72    pub fn run(&self, action: SignalAction) {
73        if let Ok(handlers) = self.handlers.lock() {
74            for (_, handler) in handlers.iter() {
75                handler(action);
76            }
77        }
78    }
79}
80
81#[inline]
82#[expect(unused, reason = "used in `Debug` impls")]
83fn debug_fmt_handlers(handlers: &Mutex<Vec<(usize, Handler)>>) -> impl Display {
84    fmt::from_fn(|f| match handlers.try_lock() {
85        Err(err) => write!(f, "{err:?}"),
86        Ok(handlers) => {
87            let ids: Vec<_> = handlers.iter().map(|(id, _)| id).collect();
88            write!(f, "{ids:?}")
89        }
90    })
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use std::sync::atomic::{AtomicBool, Ordering};
97
98    #[test]
99    /// Tests registering and running multiple handlers.
100    fn test_multiple_handlers() {
101        let handlers = Handlers::new();
102        let called1 = Arc::new(AtomicBool::new(false));
103        let called2 = Arc::new(AtomicBool::new(false));
104
105        let called1_clone = Arc::clone(&called1);
106        let called2_clone = Arc::clone(&called2);
107
108        let _guard1 = handlers.register(Box::new(move |_| {
109            called1_clone.store(true, Ordering::SeqCst);
110        }));
111        let _guard2 = handlers.register(Box::new(move |_| {
112            called2_clone.store(true, Ordering::SeqCst);
113        }));
114
115        handlers.run(SignalAction::Interrupt);
116
117        assert!(called1.load(Ordering::SeqCst));
118        assert!(called2.load(Ordering::SeqCst));
119    }
120
121    #[test]
122    /// Tests the dropping of a guard and ensuring the handler is unregistered.
123    fn test_guard_drop() {
124        let handlers = Handlers::new();
125        let called = Arc::new(AtomicBool::new(false));
126        let called_clone = Arc::clone(&called);
127
128        let guard = handlers.register(Box::new(move |_| {
129            called_clone.store(true, Ordering::Relaxed);
130        }));
131
132        // Ensure the handler is registered
133        assert_eq!(handlers.handlers.lock().unwrap().len(), 1);
134
135        drop(guard);
136
137        // Ensure the handler is removed after dropping the guard
138        assert_eq!(handlers.handlers.lock().unwrap().len(), 0);
139
140        handlers.run(SignalAction::Interrupt);
141
142        // Ensure the handler is not called after being dropped
143        assert!(!called.load(Ordering::Relaxed));
144    }
145}