nu_protocol/pipeline/
handlers.rs

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