moirai_executor/schedule/runtime/
idle_hooks.rs

1//! Worker-side idle hooks for quiescent resource reclamation.
2//!
3//! Executor workers are long-lived: they outlive individual data-parallel
4//! operations, so thread-local scratch buffers and allocator arenas can remain
5//! resident for the process lifetime. A consumer that can name a cheap,
6//! owner-thread-only reclamation step registers it here; workers run the hooks
7//! after exhausting their spin budget and before parking for more work.
8//!
9//! Registration is bounded and allocation-free. The fixed capacity makes an
10//! accepted registration a guarantee that the hook appears in every later
11//! snapshot, instead of accepting entries that the old snapshot bound could
12//! silently discard.
13
14use std::fmt;
15use std::sync::{Mutex, MutexGuard, OnceLock};
16
17/// Maximum number of process-wide worker idle hooks.
18///
19/// A registration beyond this capacity returns
20/// [`IdleHookRegistrationError::CapacityExhausted`] without changing the
21/// registry.
22pub const MAX_IDLE_HOOKS: usize = 16;
23
24/// A worker idle hook: a plain function pointer called on the worker thread
25/// immediately before it parks for work.
26pub type IdleHook = fn();
27
28/// Failure returned when a worker idle hook cannot be admitted.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30#[non_exhaustive]
31pub enum IdleHookRegistrationError {
32    /// The fixed registry has no remaining registration slot.
33    CapacityExhausted,
34}
35
36impl fmt::Display for IdleHookRegistrationError {
37    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
38        match self {
39            Self::CapacityExhausted => write!(
40                formatter,
41                "worker idle-hook registry capacity ({MAX_IDLE_HOOKS}) is exhausted"
42            ),
43        }
44    }
45}
46
47impl std::error::Error for IdleHookRegistrationError {}
48
49/// Fixed-capacity registry used by the process-wide API and isolated tests.
50struct HookRegistry {
51    hooks: Mutex<[Option<IdleHook>; MAX_IDLE_HOOKS]>,
52}
53
54impl HookRegistry {
55    const fn new() -> Self {
56        Self {
57            hooks: Mutex::new([None; MAX_IDLE_HOOKS]),
58        }
59    }
60
61    fn lock(&self) -> MutexGuard<'_, [Option<IdleHook>; MAX_IDLE_HOOKS]> {
62        self.hooks
63            .lock()
64            .unwrap_or_else(|poisoned| poisoned.into_inner())
65    }
66
67    fn register(&self, hook: IdleHook) -> Result<(), IdleHookRegistrationError> {
68        let mut hooks = self.lock();
69        let Some(slot) = hooks.iter_mut().find(|entry| entry.is_none()) else {
70            return Err(IdleHookRegistrationError::CapacityExhausted);
71        };
72        *slot = Some(hook);
73        Ok(())
74    }
75
76    fn run(&self) {
77        let snapshot = {
78            let hooks = self.lock();
79            *hooks
80        };
81        for hook in snapshot.into_iter().flatten() {
82            hook();
83        }
84    }
85}
86
87/// Process-wide worker idle-hook registry.
88static HOOKS: OnceLock<HookRegistry> = OnceLock::new();
89
90fn registry() -> &'static HookRegistry {
91    HOOKS.get_or_init(HookRegistry::new)
92}
93
94/// Registers `hook` to run on every worker thread before it parks for work.
95///
96/// Registration fills the next fixed slot and never allocates. Hooks run in
97/// registration order, and duplicate function pointers are retained as
98/// separate registrations. The registry lock is released before callbacks
99/// execute, so a callback may register a later hook without deadlocking.
100///
101/// # Errors
102/// Returns [`IdleHookRegistrationError::CapacityExhausted`] when all
103/// [`MAX_IDLE_HOOKS`] slots are occupied. A rejected registration leaves every
104/// existing slot unchanged.
105pub fn register_idle_hook(hook: IdleHook) -> Result<(), IdleHookRegistrationError> {
106    registry().register(hook)
107}
108
109/// Runs the hooks registered for the calling worker.
110///
111/// The fixed function-pointer snapshot is copied while the registry lock is
112/// held and callbacks run after that lock is released. A callback panic
113/// propagates to its worker and stops the current snapshot; it does not poison
114/// the registry because callbacks never run while the lock is held.
115pub fn run_idle_hooks() {
116    registry().run();
117}
118
119#[cfg(test)]
120mod tests;
moirai_executor/schedule/runtime/idle_hooks.rs

moirai_executor/schedule/runtime/
idle_hooks.rs
