zero_pool/
pool.rs

1use crate::{
2    queue::Queue,
3    scope::{Scope, ScopeGuard},
4    task_batch::TaskBatch,
5    worker::spawn_worker,
6};
7use std::{
8    num::NonZeroUsize,
9    sync::Arc,
10    thread::{self, JoinHandle},
11};
12
13/// A high-performance thread pool with lock-free task dispatch.
14///
15/// `ZeroPool` coordinates tasks across worker threads using a single lock-free
16/// FIFO queue and cooperative memory reclamation. Tasks write their results
17/// directly to caller-provided memory via function pointers without virtual dispatch.
18pub struct ZeroPool {
19    queue: Arc<Queue>,
20    workers: Box<[JoinHandle<()>]>,
21}
22
23impl ZeroPool {
24    /// Creates a new thread pool with worker count equal to available parallelism.
25    ///
26    /// Worker count is determined by `std::thread::available_parallelism()`,
27    /// falling back to 1 if unavailable. This is usually the optimal choice
28    /// for CPU-bound workloads.
29    ///
30    /// # Examples
31    ///
32    /// ```rust
33    /// use zero_pool::ZeroPool;
34    /// let pool = ZeroPool::new();
35    /// ```
36    #[must_use]
37    pub fn new() -> Self {
38        let worker_count = thread::available_parallelism().unwrap_or(NonZeroUsize::MIN);
39        Self::with_workers(worker_count)
40    }
41
42    /// Creates a new thread pool with the specified number of workers.
43    ///
44    /// Use this when you need precise control over the worker count,
45    /// for example when coordinating with other thread pools or
46    /// when you know the optimal count for your specific workload.
47    ///
48    /// # Examples
49    ///
50    /// ```rust
51    /// use std::num::NonZeroUsize;
52    /// use zero_pool::ZeroPool;
53    /// let pool = ZeroPool::with_workers(NonZeroUsize::new(4).unwrap());
54    /// ```
55    #[must_use]
56    pub fn with_workers(worker_count: NonZeroUsize) -> Self {
57        let worker_count = worker_count.get();
58        let queue = Arc::new(Queue::new(worker_count));
59
60        let workers = (0..worker_count)
61            .map(|id| {
62                let handle = spawn_worker(id, queue.clone());
63                queue.register_thread(id, handle.thread().clone());
64                handle
65            })
66            .collect();
67
68        ZeroPool { queue, workers }
69    }
70
71    /// Creates a scope for executing concurrent tasks that can borrow from the stack.
72    ///
73    /// All tasks are guaranteed to complete before `scope` returns.
74    ///
75    /// # Examples
76    ///
77    /// ```rust
78    /// use zero_pool::ZeroPool;
79    ///
80    /// struct Params { value: u64, result: *mut u64 }
81    /// fn compute(p: &Params) { unsafe { *p.result = p.value * 2; } }
82    ///
83    /// let pool = ZeroPool::new();
84    /// let mut r1 = 0;
85    /// let mut r2 = 0;
86    ///
87    /// let p1 = [Params { value: 10, result: &raw mut r1 }];
88    /// let p2 = [Params { value: 20, result: &raw mut r2 }];
89    ///
90    /// pool.scope(|s| {
91    ///     s.run(compute, &p1);
92    ///     s.run(compute, &p2);
93    /// });
94    ///
95    /// assert_eq!(r1, 20);
96    /// assert_eq!(r2, 40);
97    /// ```
98    #[inline]
99    pub fn scope<'env, F, R>(&'env self, f: F) -> R
100    where
101        F: for<'scope> FnOnce(&'scope Scope<'scope, 'env>) -> R,
102    {
103        let scope = Scope::new(&self.queue);
104        let _guard = ScopeGuard(&scope);
105        f(&scope)
106    }
107
108    /// Submits tasks and waits for all to complete.
109    ///
110    /// ```rust
111    /// use zero_pool::ZeroPool;
112    ///
113    /// struct Params { value: u64, result: *mut u64 }
114    /// fn compute(p: &Params) { unsafe { *p.result = p.value * 2; } }
115    ///
116    /// let pool = ZeroPool::new();
117    /// let mut result = 0;
118    /// pool.run(compute, &[Params { value: 42, result: &raw mut result }]);
119    /// assert_eq!(result, 84);
120    /// ```
121    #[inline]
122    pub fn run<T>(&self, task_fn: fn(&T), params: &[T]) {
123        self.scope(|s| s.run(task_fn, params));
124    }
125
126    /// Submits tasks without waiting for completion.
127    ///
128    /// # Safety
129    ///
130    /// The caller must ensure that `params` remains valid until all tasks finish.
131    #[inline]
132    pub unsafe fn run_detached<T>(&self, task_fn: fn(&T), params: *const [T]) {
133        if !params.is_empty() {
134            let batch = unsafe { TaskBatch::new(task_fn, params, None) };
135            self.queue.enqueue(batch, params.len());
136        }
137    }
138}
139
140impl Default for ZeroPool {
141    /// Creates a new thread pool with default settings
142    ///
143    /// Equivalent to calling `ZeroPool::new()`. Worker count is determined by
144    /// `std::thread::available_parallelism()`, falling back to 1 if unavailable.
145    #[inline]
146    fn default() -> Self {
147        Self::new()
148    }
149}
150
151impl Drop for ZeroPool {
152    fn drop(&mut self) {
153        self.queue.shutdown();
154
155        for handle in std::mem::take(&mut self.workers) {
156            let _ = handle.join();
157        }
158    }
159}
zero_pool/pool.rs

zero_pool/
pool.rs
