rskit_hook/
registry.rs

1use std::any::{Any, TypeId};
2use std::collections::HashMap;
3use std::sync::Arc;
4use std::sync::atomic::{AtomicUsize, Ordering};
5
6use parking_lot::RwLock;
7use tokio_util::sync::CancellationToken;
8
9use crate::types::{Event, EventType, HookError, HookResult};
10
11type ErasedEvent = dyn Any + Send + Sync;
12type ErasedHandler = Arc<dyn Fn(CancellationToken, &ErasedEvent) -> HookResult + Send + Sync>;
13type HandlerKey = (TypeId, EventType);
14type HandlerMap = Arc<RwLock<HashMap<HandlerKey, Vec<(usize, ErasedHandler)>>>>;
15
16/// A thread-safe registry that maps typed events to ordered observe-only handlers.
17///
18/// Registration and emission are constructor-injected through registry values;
19/// there is no global hook registry. Handlers for the same event type run in
20/// registration order. A fatal [`HookError`] short-circuits the remaining
21/// handlers for that event type; non-fatal errors are collected and the first
22/// non-fatal error is returned after all handlers run.
23pub struct HookRegistry {
24    handlers: HandlerMap,
25    event_types: Arc<RwLock<HashMap<TypeId, EventType>>>,
26    next_id: AtomicUsize,
27}
28
29impl std::fmt::Debug for HookRegistry {
30    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        let event_types = self.event_types.read();
32        let handlers = self.handlers.read();
33        let counts = handlers
34            .iter()
35            .map(|((type_id, registered_type), handlers)| {
36                let event_type = event_types
37                    .get(type_id)
38                    .unwrap_or(registered_type)
39                    .to_string();
40                (event_type, handlers.len())
41            })
42            .collect::<HashMap<_, _>>();
43        f.debug_struct("HookRegistry")
44            .field("handler_counts", &counts)
45            .finish()
46    }
47}
48
49impl Default for HookRegistry {
50    fn default() -> Self {
51        Self::new()
52    }
53}
54
55impl HookRegistry {
56    /// Create a new empty registry.
57    #[must_use]
58    pub fn new() -> Self {
59        Self {
60            handlers: Arc::new(RwLock::new(HashMap::new())),
61            event_types: Arc::new(RwLock::new(HashMap::new())),
62            next_id: AtomicUsize::new(0),
63        }
64    }
65
66    /// Register a typed handler for event `E`.
67    ///
68    /// The returned function unregisters this handler when invoked.
69    pub fn on<E>(
70        &self,
71        event_type: EventType,
72        handler: impl Fn(CancellationToken, &E) -> HookResult + Send + Sync + 'static,
73    ) -> Box<dyn FnOnce() + Send>
74    where
75        E: Event,
76    {
77        let type_id = TypeId::of::<E>();
78        let key = (type_id, event_type.clone());
79        self.event_types.write().insert(type_id, event_type);
80        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
81        self.handlers.write().entry(key.clone()).or_default().push((
82            id,
83            Arc::new(move |cancel, event| {
84                let Some(event) = event.downcast_ref::<E>() else {
85                    return Err(HookError::fatal("hook event type mismatch"));
86                };
87                handler(cancel, event)
88            }),
89        ));
90        let handlers = Arc::clone(&self.handlers);
91        Box::new(move || {
92            if let Some(items) = handlers.write().get_mut(&key) {
93                items.retain(|(handler_id, _)| *handler_id != id);
94            }
95        })
96    }
97
98    /// Emit an event to all registered handlers for its concrete type.
99    ///
100    /// Dispatch uses snapshot semantics: handlers registered or removed during an
101    /// in-flight emit do not affect that emit, but apply to subsequent emits.
102    pub fn emit<E>(&self, event: &E, cancel: CancellationToken) -> HookResult
103    where
104        E: Event,
105    {
106        let key = (TypeId::of::<E>(), event.event_type());
107        let handler_list = {
108            let handlers = self.handlers.read();
109            let Some(handler_list) = handlers.get(&key) else {
110                return Ok(());
111            };
112            handler_list
113                .iter()
114                .map(|(_, handler)| Arc::clone(handler))
115                .collect::<Vec<_>>()
116        };
117        let mut first_error: Option<HookError> = None;
118        for handler in handler_list {
119            if cancel.is_cancelled() {
120                let err = HookError::fatal("hook dispatch cancelled");
121                return Err(err);
122            }
123            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
124                handler(cancel.clone(), event)
125            }))
126            .unwrap_or_else(|payload| {
127                Err(HookError::new(format!(
128                    "hook handler panicked: {}",
129                    panic_message(payload)
130                )))
131            });
132            if let Err(err) = result {
133                if err.is_fatal() {
134                    return Err(err);
135                }
136                if first_error.is_none() {
137                    first_error = Some(err);
138                }
139            }
140        }
141        first_error.map_or(Ok(()), Err)
142    }
143
144    /// Check whether any handlers are registered for event `E` and `event_type`.
145    #[must_use]
146    pub fn has_handlers<E>(&self, event_type: &EventType) -> bool
147    where
148        E: Event,
149    {
150        self.handlers
151            .read()
152            .get(&(TypeId::of::<E>(), event_type.clone()))
153            .is_some_and(|handlers| !handlers.is_empty())
154    }
155
156    /// Remove all handlers for event `E` and `event_type`.
157    pub fn clear<E>(&self, event_type: &EventType)
158    where
159        E: Event,
160    {
161        self.handlers
162            .write()
163            .remove(&(TypeId::of::<E>(), event_type.clone()));
164    }
165
166    /// Remove all handlers from the registry.
167    pub fn clear_all(&self) {
168        let mut handlers = self.handlers.write();
169        handlers.clear();
170    }
171}
172
173/// Typed lifecycle hook registry.
174///
175/// Lifecycle hooks use the same ordered synchronous dispatch semantics as
176/// [`HookRegistry`]; the alias exists to make lifecycle-specific injection
177/// points self-documenting.
178pub type LifecycleHookRegistry = HookRegistry;
179
180fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
181    if let Some(message) = payload.downcast_ref::<&str>() {
182        (*message).to_string()
183    } else if let Some(message) = payload.downcast_ref::<String>() {
184        message.clone()
185    } else {
186        "unknown panic payload".to_string()
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use std::sync::Arc;
193    use std::sync::atomic::{AtomicU32, Ordering};
194
195    use tokio_util::sync::CancellationToken;
196
197    use super::HookRegistry;
198    use crate::{Event, EventType, HookError};
199
200    struct Ping;
201    impl Event for Ping {
202        fn event_type(&self) -> EventType {
203            EventType::new("ping")
204        }
205    }
206    struct Pong;
207    impl Event for Pong {
208        fn event_type(&self) -> EventType {
209            EventType::new("pong")
210        }
211    }
212    fn ping_type() -> EventType {
213        EventType::new("ping")
214    }
215    fn pong_type() -> EventType {
216        EventType::new("pong")
217    }
218
219    #[test]
220    fn register_and_emit() {
221        let registry = HookRegistry::new();
222        let counter = Arc::new(AtomicU32::new(0));
223        let cloned = Arc::clone(&counter);
224        let _unsubscribe = registry.on::<Ping>(ping_type(), move |_, _| {
225            cloned.fetch_add(1, Ordering::SeqCst);
226            Ok(())
227        });
228        assert!(registry.emit(&Ping, CancellationToken::new()).is_ok());
229        assert_eq!(counter.load(Ordering::SeqCst), 1);
230    }
231
232    #[test]
233    fn emit_no_handlers_returns_ok() {
234        assert!(
235            HookRegistry::new()
236                .emit(&Pong, CancellationToken::new())
237                .is_ok()
238        );
239    }
240
241    #[test]
242    fn non_fatal_error_does_not_short_circuit() {
243        let registry = HookRegistry::new();
244        let counter = Arc::new(AtomicU32::new(0));
245        let _unsubscribe1 = registry.on::<Ping>(ping_type(), |_, _| Err(HookError::new("warn")));
246        let cloned = Arc::clone(&counter);
247        let _unsubscribe2 = registry.on::<Ping>(ping_type(), move |_, _| {
248            cloned.fetch_add(1, Ordering::SeqCst);
249            Ok(())
250        });
251        let result = registry.emit(&Ping, CancellationToken::new());
252        assert_eq!(result.expect_err("error").message(), "warn");
253        assert_eq!(counter.load(Ordering::SeqCst), 1);
254    }
255
256    #[test]
257    fn fatal_error_short_circuits() {
258        let registry = HookRegistry::new();
259        let counter = Arc::new(AtomicU32::new(0));
260        let _unsubscribe1 =
261            registry.on::<Ping>(ping_type(), |_, _| Err(HookError::fatal("blocked")));
262        let cloned = Arc::clone(&counter);
263        let _unsubscribe2 = registry.on::<Ping>(ping_type(), move |_, _| {
264            cloned.fetch_add(1, Ordering::SeqCst);
265            Ok(())
266        });
267        assert!(
268            registry
269                .emit(&Ping, CancellationToken::new())
270                .expect_err("fatal")
271                .is_fatal()
272        );
273        assert_eq!(counter.load(Ordering::SeqCst), 0);
274    }
275
276    #[test]
277    fn cancelled_emit_errors_before_dispatch() {
278        let registry = HookRegistry::new();
279        let counter = Arc::new(AtomicU32::new(0));
280        let cloned = Arc::clone(&counter);
281        let _unsubscribe = registry.on::<Ping>(ping_type(), move |_, _| {
282            cloned.fetch_add(1, Ordering::SeqCst);
283            Ok(())
284        });
285        let token = CancellationToken::new();
286        token.cancel();
287        assert!(
288            registry
289                .emit(&Ping, token)
290                .expect_err("cancelled")
291                .is_fatal()
292        );
293        assert_eq!(counter.load(Ordering::SeqCst), 0);
294    }
295
296    #[test]
297    fn panicking_handler_becomes_error_and_continues() {
298        let registry = HookRegistry::new();
299        let counter = Arc::new(AtomicU32::new(0));
300        let _unsubscribe1 = registry.on::<Ping>(ping_type(), |_, _| panic!("boom"));
301        let cloned = Arc::clone(&counter);
302        let _unsubscribe2 = registry.on::<Ping>(ping_type(), move |_, _| {
303            cloned.fetch_add(1, Ordering::SeqCst);
304            Ok(())
305        });
306        assert!(registry.emit(&Ping, CancellationToken::new()).is_err());
307        assert_eq!(counter.load(Ordering::SeqCst), 1);
308    }
309
310    #[test]
311    fn unsubscribe_removes_handler() {
312        let registry = HookRegistry::new();
313        let counter = Arc::new(AtomicU32::new(0));
314        let cloned = Arc::clone(&counter);
315        struct ErrorEvent;
316        impl Event for ErrorEvent {
317            fn event_type(&self) -> EventType {
318                EventType::new("error")
319            }
320        }
321        let unsubscribe = registry.on::<ErrorEvent>(EventType::new("error"), move |_, _| {
322            cloned.fetch_add(1, Ordering::SeqCst);
323            Ok(())
324        });
325        unsubscribe();
326        let _ = registry.emit(&ErrorEvent, CancellationToken::new());
327        assert_eq!(counter.load(Ordering::SeqCst), 0);
328    }
329
330    #[test]
331    fn clear_removes_selected_event_types() {
332        let registry = HookRegistry::new();
333        struct A;
334        impl Event for A {
335            fn event_type(&self) -> EventType {
336                EventType::new("a")
337            }
338        }
339        struct B;
340        impl Event for B {
341            fn event_type(&self) -> EventType {
342                EventType::new("b")
343            }
344        }
345        let a = EventType::new("a");
346        let b = EventType::new("b");
347        let _u1 = registry.on::<A>(a.clone(), |_, _| Ok(()));
348        let _u2 = registry.on::<B>(b.clone(), |_, _| Ok(()));
349        registry.clear::<A>(&a);
350        assert!(!registry.has_handlers::<A>(&a));
351        assert!(registry.has_handlers::<B>(&b));
352        registry.clear_all();
353        assert!(!registry.has_handlers::<B>(&b));
354    }
355
356    #[test]
357    fn multiple_handlers_same_event_all_run() {
358        let registry = HookRegistry::new();
359        let counter = Arc::new(AtomicU32::new(0));
360        for _ in 0..3 {
361            let cloned = Arc::clone(&counter);
362            let _unsubscribe = registry.on::<Pong>(pong_type(), move |_, _| {
363                cloned.fetch_add(1, Ordering::SeqCst);
364                Ok(())
365            });
366        }
367        let _ = registry.emit(&Pong, CancellationToken::new());
368        assert_eq!(counter.load(Ordering::SeqCst), 3);
369    }
370
371    #[test]
372    fn handler_can_mutate_registry_during_emit() {
373        let registry = Arc::new(HookRegistry::new());
374        let counter = Arc::new(AtomicU32::new(0));
375
376        let registry_for_handler = Arc::clone(&registry);
377        let _unsubscribe1 = registry.on::<Ping>(ping_type(), move |_, _| {
378            registry_for_handler.clear::<Ping>(&ping_type());
379            Ok(())
380        });
381
382        let cloned = Arc::clone(&counter);
383        let _unsubscribe2 = registry.on::<Ping>(ping_type(), move |_, _| {
384            cloned.fetch_add(1, Ordering::SeqCst);
385            Ok(())
386        });
387
388        registry
389            .emit(&Ping, CancellationToken::new())
390            .expect("emit");
391
392        assert_eq!(counter.load(Ordering::SeqCst), 1);
393        assert!(!registry.has_handlers::<Ping>(&ping_type()));
394    }
395}
rskit_hook/registry.rs

rskit_hook/
registry.rs
