zero_pool/
pool.rs

1use crate::{queue::Queue, task_future::TaskFuture, worker::spawn_worker};
2use std::{
3    num::NonZeroUsize,
4    slice,
5    sync::Arc,
6    thread::{self, JoinHandle},
7};
8
9pub struct ZeroPool {
10    queue: Arc<Queue>,
11    workers: Box<[JoinHandle<()>]>,
12}
13
14impl ZeroPool {
15    /// Creates a new thread pool with worker count equal to available parallelism
16    ///
17    /// Worker count is determined by `std::thread::available_parallelism()`,
18    /// falling back to 1 if unavailable. This is usually the optimal choice
19    /// for CPU-bound workloads.
20    ///
21    /// # Examples
22    ///
23    /// ```rust
24    /// use zero_pool::ZeroPool;
25    /// let pool = ZeroPool::new();
26    /// ```
27    #[must_use]
28    pub fn new() -> Self {
29        let worker_count = thread::available_parallelism().unwrap_or(NonZeroUsize::MIN);
30        Self::with_workers(worker_count)
31    }
32
33    /// Creates a new thread pool with the specified number of workers
34    ///
35    /// Use this when you need precise control over the worker count,
36    /// for example when coordinating with other thread pools or
37    /// when you know the optimal count for your specific workload.
38    ///
39    /// # Examples
40    ///
41    /// ```rust
42    /// use std::num::NonZeroUsize;
43    /// use zero_pool::ZeroPool;
44    /// let pool = ZeroPool::with_workers(NonZeroUsize::new(4).unwrap());
45    /// ```
46    #[must_use]
47    pub fn with_workers(worker_count: NonZeroUsize) -> Self {
48        let worker_count = worker_count.get();
49
50        let queue = Arc::new(Queue::new(worker_count));
51
52        let workers = (0..worker_count)
53            .map(|id| {
54                let handle = spawn_worker(id, queue.clone());
55                queue.register_thread(id, handle.thread().clone());
56                handle
57            })
58            .collect();
59
60        ZeroPool { queue, workers }
61    }
62
63    /// Submit a single typed task
64    ///
65    /// The parameter struct must remain valid until the future completes.
66    /// This is the recommended method for submitting individual tasks.
67    ///
68    /// # Examples
69    ///
70    /// ```rust
71    /// use zero_pool::ZeroPool;
72    ///
73    /// struct MyTaskParams { value: u64, result: *mut u64 }
74    ///
75    /// fn my_task_fn(params: &MyTaskParams) {
76    ///     unsafe { *params.result = params.value * 2; }
77    /// }
78    ///
79    /// let pool = ZeroPool::new();
80    /// let mut result = 0u64;
81    /// let task_params = MyTaskParams { value: 42, result: &mut result };
82    /// let future = pool.submit_task(my_task_fn, &task_params);
83    /// future.wait();
84    /// assert_eq!(result, 84);
85    /// ```
86    #[inline]
87    pub fn submit_task<T>(&self, task_fn: fn(&T), param: &T) -> TaskFuture {
88        self.queue.push_task_batch(task_fn, slice::from_ref(param))
89    }
90
91    /// Submit a batch of uniform tasks
92    ///
93    /// All tasks in the batch must be the same type and use the same task function.
94    /// This method handles the pointer conversion automatically and is the most
95    /// convenient way to submit large batches of similar work.
96    ///
97    /// # Examples
98    ///
99    /// ```rust
100    /// use zero_pool::ZeroPool;
101    ///
102    /// struct MyTaskParams { value: u64, result: *mut u64 }
103    ///
104    /// fn my_task_fn(params: &MyTaskParams) {
105    ///     unsafe { *params.result = params.value * 2; }
106    /// }
107    ///
108    /// let pool = ZeroPool::new();
109    /// let mut results = vec![0u64; 1000];
110    /// let task_params: Vec<_> = results
111    ///     .iter_mut()
112    ///     .enumerate()
113    ///     .map(|(i, res)| MyTaskParams { value: i as u64, result: res })
114    ///     .collect();
115    /// let future = pool.submit_batch(my_task_fn, &task_params);
116    /// future.wait();
117    /// assert_eq!(results[0], 0);
118    /// assert_eq!(results[1], 2);
119    /// assert_eq!(results[999], 1998);
120    /// ```
121    #[inline]
122    pub fn submit_batch<T>(&self, task_fn: fn(&T), params: &[T]) -> TaskFuture {
123        self.queue.push_task_batch(task_fn, params)
124    }
125}
126
127impl Default for ZeroPool {
128    /// Creates a new thread pool with default settings
129    ///
130    /// Equivalent to calling `ZeroPool::new()`. Worker count is determined by
131    /// `std::thread::available_parallelism()`, falling back to 1 if unavailable.
132    fn default() -> Self {
133        Self::new()
134    }
135}
136
137impl Drop for ZeroPool {
138    fn drop(&mut self) {
139        self.queue.shutdown();
140
141        for handle in std::mem::take(&mut self.workers) {
142            let _ = handle.join();
143        }
144    }
145}
zero_pool/pool.rs

zero_pool/
pool.rs
