zero_pool/
pool.rs

1use crate::{queue::Queue, task_future::TaskFuture, worker::spawn_worker};
2use std::{
3    num::NonZeroUsize,
4    sync::Arc,
5    thread::{self, JoinHandle},
6};
7
8pub struct ZeroPool {
9    queue: Arc<Queue>,
10    workers: Box<[JoinHandle<()>]>,
11}
12
13impl ZeroPool {
14    /// Creates a new thread pool with worker count equal to available parallelism
15    ///
16    /// Worker count is determined by `std::thread::available_parallelism()`,
17    /// falling back to 1 if unavailable. This is usually the optimal choice
18    /// for CPU-bound workloads.
19    ///
20    /// # Examples
21    ///
22    /// ```rust
23    /// use zero_pool::ZeroPool;
24    /// let pool = ZeroPool::new();
25    /// ```
26    #[must_use]
27    pub fn new() -> Self {
28        let worker_count = thread::available_parallelism().unwrap_or(NonZeroUsize::MIN);
29        Self::with_workers(worker_count)
30    }
31
32    /// Creates a new thread pool with the specified number of workers
33    ///
34    /// Use this when you need precise control over the worker count,
35    /// for example when coordinating with other thread pools or
36    /// when you know the optimal count for your specific workload.
37    ///
38    /// # Examples
39    ///
40    /// ```rust
41    /// use std::num::NonZeroUsize;
42    /// use zero_pool::ZeroPool;
43    /// let pool = ZeroPool::with_workers(NonZeroUsize::new(4).unwrap());
44    /// ```
45    #[must_use]
46    pub fn with_workers(worker_count: NonZeroUsize) -> Self {
47        let worker_count = worker_count.get();
48
49        let queue = Arc::new(Queue::new(worker_count));
50
51        let workers = (0..worker_count)
52            .map(|id| {
53                let handle = spawn_worker(id, queue.clone());
54                queue.register_thread(id, handle.thread().clone());
55                handle
56            })
57            .collect();
58
59        ZeroPool { queue, workers }
60    }
61
62    /// Submit a single typed task
63    ///
64    /// The parameter struct must remain valid until the future completes.
65    /// This is the recommended method for submitting individual tasks.
66    ///
67    /// # Examples
68    ///
69    /// ```rust
70    /// use zero_pool::ZeroPool;
71    ///
72    /// struct MyTaskParams { value: u64, result: *mut u64 }
73    ///
74    /// fn my_task_fn(params: &MyTaskParams) {
75    ///     unsafe { *params.result = params.value * 2; }
76    /// }
77    ///
78    /// let pool = ZeroPool::new();
79    /// let mut result = 0u64;
80    /// let task_params = MyTaskParams { value: 42, result: &mut result };
81    /// let future = pool.submit_task(my_task_fn, &task_params);
82    /// future.wait();
83    /// assert_eq!(result, 84);
84    /// ```
85    #[inline]
86    pub fn submit_task<T>(&self, task_fn: fn(&T), params: &T) -> TaskFuture {
87        let slice = std::slice::from_ref(params);
88        self.queue.push_task_batch(task_fn, slice)
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_vec: &[T]) -> TaskFuture {
123        self.queue.push_task_batch(task_fn, params_vec)
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
