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    /// Registers a new handler which persists for the entire process lifetime.
72    ///
73    /// Only use this for handlers which should exist for the lifetime of the program.
74    /// You should prefer to use `register` with a `HandlerGuard` when possible.
75    pub fn register_unguarded(&self, handler: Handler) -> Result<(), ShellError> {
76        let id = self.next_id.next()?;
77        if let Ok(mut handlers) = self.handlers.lock() {
78            handlers.push((id, handler));
79        }
80
81        Ok(())
82    }
83
84    /// Runs all registered handlers.
85    pub fn run(&self, action: SignalAction) {
86        if let Ok(handlers) = self.handlers.lock() {
87            for (_, handler) in handlers.iter() {
88                handler(action);
89            }
90        }
91    }
92}
93
94impl Default for Handlers {
95    fn default() -> Self {
96        Self::new()
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use std::sync::atomic::{AtomicBool, Ordering};
104
105    #[test]
106    /// Tests registering and running multiple handlers.
107    fn test_multiple_handlers() {
108        let handlers = Handlers::new();
109        let called1 = Arc::new(AtomicBool::new(false));
110        let called2 = Arc::new(AtomicBool::new(false));
111
112        let called1_clone = Arc::clone(&called1);
113        let called2_clone = Arc::clone(&called2);
114
115        let _guard1 = handlers.register(Box::new(move |_| {
116            called1_clone.store(true, Ordering::SeqCst);
117        }));
118        let _guard2 = handlers.register(Box::new(move |_| {
119            called2_clone.store(true, Ordering::SeqCst);
120        }));
121
122        handlers.run(SignalAction::Interrupt);
123
124        assert!(called1.load(Ordering::SeqCst));
125        assert!(called2.load(Ordering::SeqCst));
126    }
127
128    #[test]
129    /// Tests the dropping of a guard and ensuring the handler is unregistered.
130    fn test_guard_drop() {
131        let handlers = Handlers::new();
132        let called = Arc::new(AtomicBool::new(false));
133        let called_clone = Arc::clone(&called);
134
135        let guard = handlers.register(Box::new(move |_| {
136            called_clone.store(true, Ordering::Relaxed);
137        }));
138
139        // Ensure the handler is registered
140        assert_eq!(handlers.handlers.lock().unwrap().len(), 1);
141
142        drop(guard);
143
144        // Ensure the handler is removed after dropping the guard
145        assert_eq!(handlers.handlers.lock().unwrap().len(), 0);
146
147        handlers.run(SignalAction::Interrupt);
148
149        // Ensure the handler is not called after being dropped
150        assert!(!called.load(Ordering::Relaxed));
151    }
152}
nu_protocol/pipeline/handlers.rs

nu_protocol/pipeline/
handlers.rs
