weavatrix_scan/
runtime.rs

1#![allow(clippy::missing_const_for_thread_local)]
2
3use crate::pool::{Job, ThreadPool};
4use std::cell::RefCell;
5use std::fmt;
6use std::io;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::{Arc, OnceLock};
9#[cfg(feature = "rayon")]
10use std::sync::{Mutex, mpsc};
11use std::time::Duration;
12
13thread_local! {
14    static ACTIVE_RUNTIMES: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
15}
16
17/// A job accepted by an embeddable [`ParallelExecutor`].
18pub type ParallelJob = Box<dyn FnOnce() + Send + 'static>;
19
20/// Adapter contract for an application-owned thread pool.
21///
22/// Implementations must either accept `job` exactly once or return an error
23/// without retaining it. `busy_timeout` lets bounded pools reject work instead
24/// of indefinitely waiting for capacity.
25pub trait ParallelExecutor: Send + Sync + 'static {
26    /// Maximum useful concurrent jobs for this executor.
27    fn parallelism(&self) -> usize;
28
29    /// Attempts to schedule one job.
30    ///
31    /// # Errors
32    ///
33    /// Returns an I/O error when the pool is closed, saturated past
34    /// `busy_timeout`, or otherwise cannot accept the job.
35    fn try_execute(&self, job: ParallelJob, busy_timeout: Option<Duration>) -> io::Result<()>;
36}
37
38/// Ready-to-use adapter for an application-owned Rayon pool.
39#[cfg(feature = "rayon")]
40#[derive(Clone)]
41pub struct RayonExecutor {
42    pool: Arc<rayon::ThreadPool>,
43}
44
45#[cfg(feature = "rayon")]
46impl RayonExecutor {
47    /// Wraps an existing Rayon pool.
48    #[must_use]
49    pub fn new(pool: Arc<rayon::ThreadPool>) -> Self {
50        Self { pool }
51    }
52
53    /// Returns the wrapped pool.
54    #[must_use]
55    pub fn pool(&self) -> &Arc<rayon::ThreadPool> {
56        &self.pool
57    }
58}
59
60#[cfg(feature = "rayon")]
61impl ParallelExecutor for RayonExecutor {
62    fn parallelism(&self) -> usize {
63        self.pool.current_num_threads().max(1)
64    }
65
66    fn try_execute(&self, job: ParallelJob, busy_timeout: Option<Duration>) -> io::Result<()> {
67        let Some(timeout) = busy_timeout else {
68            self.pool.spawn(job);
69            return Ok(());
70        };
71        let pending = Arc::new(Mutex::new(Some(job)));
72        let worker_pending = Arc::clone(&pending);
73        let (started_sender, started_receiver) = mpsc::sync_channel(0);
74        self.pool.spawn(move || {
75            let job = worker_pending
76                .lock()
77                .unwrap_or_else(std::sync::PoisonError::into_inner)
78                .take();
79            if let Some(job) = job {
80                let _ = started_sender.send(());
81                job();
82            }
83        });
84        match started_receiver.recv_timeout(timeout) {
85            Ok(()) => Ok(()),
86            Err(mpsc::RecvTimeoutError::Timeout) => {
87                let cancelled = pending
88                    .lock()
89                    .unwrap_or_else(std::sync::PoisonError::into_inner)
90                    .take()
91                    .is_some();
92                if cancelled {
93                    Err(io::Error::new(
94                        io::ErrorKind::TimedOut,
95                        "Rayon pool did not start the job before busy timeout",
96                    ))
97                } else {
98                    Ok(())
99                }
100            }
101            Err(mpsc::RecvTimeoutError::Disconnected) => Err(io::Error::new(
102                io::ErrorKind::BrokenPipe,
103                "Rayon pool dropped the scheduled job",
104            )),
105        }
106    }
107}
108
109enum Executor {
110    Global,
111    Dedicated(Arc<ThreadPool>),
112    External(Arc<dyn ParallelExecutor>),
113}
114
115struct RuntimeInner {
116    id: u64,
117    executor: Executor,
118    busy_timeout: Option<Duration>,
119}
120
121/// Selects where parallel traversal jobs execute.
122///
123/// The default uses the process-wide Weavatrix pool. Dedicated pools are
124/// joined on last drop. External executors receive the configured busy timeout
125/// and may reject submission without leaving a traversal waiting for a worker.
126#[derive(Clone)]
127pub struct ParallelRuntime {
128    inner: Arc<RuntimeInner>,
129}
130
131impl ParallelRuntime {
132    /// Returns the process-wide default runtime.
133    #[must_use]
134    pub fn global() -> Self {
135        static RUNTIME: OnceLock<ParallelRuntime> = OnceLock::new();
136        RUNTIME
137            .get_or_init(|| Self::new(Executor::Global, None))
138            .clone()
139    }
140
141    /// Creates an owned pool with exactly `parallelism.max(1)` workers.
142    ///
143    /// # Errors
144    ///
145    /// Returns an operating-system thread creation error.
146    pub fn dedicated(parallelism: usize) -> io::Result<Self> {
147        let pool = ThreadPool::with_workers(parallelism.max(1))?;
148        Ok(Self::new(Executor::Dedicated(Arc::new(pool)), None))
149    }
150
151    /// Uses an application-owned executor.
152    #[must_use]
153    pub fn external(executor: Arc<dyn ParallelExecutor>) -> Self {
154        Self::new(Executor::External(executor), None)
155    }
156
157    /// Uses an existing application-owned Rayon pool.
158    ///
159    /// A one-second busy timeout is enabled by default to reject nested
160    /// single-thread pool starvation. Override it with
161    /// [`Self::with_busy_timeout`] when the application has a stronger
162    /// scheduling guarantee.
163    #[cfg(feature = "rayon")]
164    #[must_use]
165    pub fn rayon_existing(pool: Arc<rayon::ThreadPool>) -> Self {
166        Self::external(Arc::new(RayonExecutor::new(pool)))
167            .with_busy_timeout(Some(Duration::from_secs(1)))
168    }
169
170    /// Creates a dedicated Rayon pool with exactly `parallelism.max(1)`
171    /// workers.
172    ///
173    /// # Errors
174    ///
175    /// Returns an error when Rayon cannot create the requested worker pool.
176    #[cfg(feature = "rayon")]
177    pub fn rayon_new(parallelism: usize) -> io::Result<Self> {
178        let pool = rayon::ThreadPoolBuilder::new()
179            .num_threads(parallelism.max(1))
180            .build()
181            .map_err(io::Error::other)?;
182        Ok(Self::rayon_existing(Arc::new(pool)))
183    }
184
185    /// Supplies the maximum wait an external executor may use to accept work.
186    #[must_use]
187    pub fn with_busy_timeout(mut self, busy_timeout: Option<Duration>) -> Self {
188        Arc::make_mut(&mut self.inner).busy_timeout = busy_timeout;
189        self
190    }
191
192    /// Maximum useful worker count advertised by this runtime.
193    #[must_use]
194    pub fn parallelism(&self) -> usize {
195        match &self.inner.executor {
196            Executor::Global => ThreadPool::global().workers(),
197            Executor::Dedicated(pool) => pool.workers(),
198            Executor::External(executor) => executor.parallelism().max(1),
199        }
200    }
201
202    pub(crate) fn is_worker_thread(&self) -> bool {
203        ACTIVE_RUNTIMES.with(|active| active.borrow().contains(&self.inner.id))
204    }
205
206    pub(crate) fn try_execute<F>(&self, job: F) -> io::Result<()>
207    where
208        F: FnOnce() + Send + 'static,
209    {
210        let id = self.inner.id;
211        let wrapped: Job = Box::new(move || {
212            let _guard = ActiveRuntimeGuard::enter(id);
213            job();
214        });
215        match &self.inner.executor {
216            Executor::Global => ThreadPool::global().execute(wrapped),
217            Executor::Dedicated(pool) => pool.execute(wrapped),
218            Executor::External(executor) => executor.try_execute(wrapped, self.inner.busy_timeout),
219        }
220    }
221
222    fn new(executor: Executor, busy_timeout: Option<Duration>) -> Self {
223        static NEXT_RUNTIME_ID: AtomicU64 = AtomicU64::new(1);
224        Self {
225            inner: Arc::new(RuntimeInner {
226                id: NEXT_RUNTIME_ID.fetch_add(1, Ordering::Relaxed),
227                executor,
228                busy_timeout,
229            }),
230        }
231    }
232}
233
234impl Default for ParallelRuntime {
235    fn default() -> Self {
236        Self::global()
237    }
238}
239
240impl fmt::Debug for ParallelRuntime {
241    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
242        let kind = match &self.inner.executor {
243            Executor::Global => "global",
244            Executor::Dedicated(_) => "dedicated",
245            Executor::External(_) => "external",
246        };
247        formatter
248            .debug_struct("ParallelRuntime")
249            .field("kind", &kind)
250            .field("parallelism", &self.parallelism())
251            .field("busy_timeout", &self.inner.busy_timeout)
252            .finish()
253    }
254}
255
256impl Clone for RuntimeInner {
257    fn clone(&self) -> Self {
258        Self {
259            id: self.id,
260            executor: match &self.executor {
261                Executor::Global => Executor::Global,
262                Executor::Dedicated(pool) => Executor::Dedicated(Arc::clone(pool)),
263                Executor::External(executor) => Executor::External(Arc::clone(executor)),
264            },
265            busy_timeout: self.busy_timeout,
266        }
267    }
268}
269
270struct ActiveRuntimeGuard;
271
272impl ActiveRuntimeGuard {
273    fn enter(id: u64) -> Self {
274        ACTIVE_RUNTIMES.with(|active| active.borrow_mut().push(id));
275        Self
276    }
277}
278
279impl Drop for ActiveRuntimeGuard {
280    fn drop(&mut self) {
281        ACTIVE_RUNTIMES.with(|active| {
282            active.borrow_mut().pop();
283        });
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::{ParallelExecutor, ParallelJob, ParallelRuntime};
290    use std::io;
291    use std::sync::{Arc, mpsc};
292    use std::time::Duration;
293
294    struct Inline;
295
296    impl ParallelExecutor for Inline {
297        fn parallelism(&self) -> usize {
298            1
299        }
300
301        fn try_execute(&self, job: ParallelJob, _busy_timeout: Option<Duration>) -> io::Result<()> {
302            job();
303            Ok(())
304        }
305    }
306
307    #[test]
308    fn dedicated_runtime_executes_and_joins() {
309        let runtime = ParallelRuntime::dedicated(2).unwrap();
310        let (sender, receiver) = mpsc::channel();
311        runtime
312            .try_execute(move || sender.send(9).unwrap())
313            .unwrap();
314        assert_eq!(receiver.recv().unwrap(), 9);
315    }
316
317    #[test]
318    fn external_runtime_marks_nested_execution() {
319        let runtime = ParallelRuntime::external(Arc::new(Inline));
320        let nested = runtime.clone();
321        let (sender, receiver) = mpsc::channel();
322        runtime
323            .try_execute(move || sender.send(nested.is_worker_thread()).unwrap())
324            .unwrap();
325        assert!(receiver.recv().unwrap());
326    }
327
328    #[cfg(feature = "rayon")]
329    #[test]
330    fn rayon_runtime_uses_existing_pool() {
331        let pool = Arc::new(
332            rayon::ThreadPoolBuilder::new()
333                .num_threads(2)
334                .build()
335                .unwrap(),
336        );
337        let runtime = ParallelRuntime::rayon_existing(Arc::clone(&pool));
338        let (sender, receiver) = mpsc::channel();
339        runtime
340            .try_execute(move || sender.send(11).unwrap())
341            .unwrap();
342        assert_eq!(receiver.recv().unwrap(), 11);
343        assert_eq!(runtime.parallelism(), 2);
344    }
345
346    #[cfg(feature = "rayon")]
347    #[test]
348    fn rayon_busy_timeout_cancels_unstarted_job() {
349        let pool = Arc::new(
350            rayon::ThreadPoolBuilder::new()
351                .num_threads(1)
352                .build()
353                .unwrap(),
354        );
355        let (block_sender, block_receiver) = mpsc::sync_channel(0);
356        let (release_sender, release_receiver) = mpsc::sync_channel(0);
357        pool.spawn(move || {
358            block_sender.send(()).unwrap();
359            release_receiver.recv().unwrap();
360        });
361        block_receiver.recv().unwrap();
362
363        let runtime = ParallelRuntime::rayon_existing(pool)
364            .with_busy_timeout(Some(Duration::from_millis(10)));
365        let executed = Arc::new(std::sync::atomic::AtomicBool::new(false));
366        let worker_executed = Arc::clone(&executed);
367        let error = runtime
368            .try_execute(move || {
369                worker_executed.store(true, std::sync::atomic::Ordering::SeqCst);
370            })
371            .unwrap_err();
372        assert_eq!(error.kind(), io::ErrorKind::TimedOut);
373        release_sender.send(()).unwrap();
374        std::thread::sleep(Duration::from_millis(10));
375        assert!(!executed.load(std::sync::atomic::Ordering::SeqCst));
376    }
377}
weavatrix_scan/runtime.rs

weavatrix_scan/
runtime.rs
