zero_pool/
pool.rs

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

zero_pool/
pool.rs
