zero_pool/
scope.rs

1use crate::queue::Queue;
2use crate::task_batch::TaskBatch;
3use std::marker::PhantomData;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::thread::{self, Thread};
6
7/// A scope for spawning concurrent tasks that borrow from the local stack.
8///
9/// All submitted tasks are guaranteed to complete before the scope exits.
10pub struct Scope<'scope, 'env: 'scope> {
11    counter: AtomicUsize,
12    queue: &'env Queue,
13    thread: Thread,
14    _marker: PhantomData<(&'scope mut &'scope (), &'env mut &'env ())>,
15}
16
17impl<'scope, 'env> Scope<'scope, 'env> {
18    pub(crate) fn new(queue: &'env Queue) -> Self {
19        Self {
20            counter: AtomicUsize::new(0),
21            queue,
22            thread: thread::current(),
23            _marker: PhantomData,
24        }
25    }
26
27    /// Submits tasks to the pool within this scope.
28    #[inline]
29    pub fn run<T: 'scope>(&self, task_fn: fn(&T), params: &'scope [T]) {
30        if !params.is_empty() {
31            self.counter.fetch_add(params.len(), Ordering::Relaxed);
32            let batch = unsafe {
33                TaskBatch::new(
34                    task_fn,
35                    params,
36                    Some((&raw const self.counter, self.thread.clone())),
37                )
38            };
39            self.queue.enqueue(batch, params.len());
40        }
41    }
42
43    /// Waits for all tasks currently submitted to this scope to complete.
44    ///
45    /// Can be called multiple times to synchronize between phases of work.
46    #[inline]
47    pub fn wait(&self) {
48        while self.counter.load(Ordering::Acquire) != 0 {
49            thread::park();
50        }
51    }
52
53    /// Checks if all submitted tasks have finished.
54    #[must_use]
55    #[inline]
56    pub fn is_complete(&self) -> bool {
57        self.counter.load(Ordering::Acquire) == 0
58    }
59}
60
61/// Guard ensuring all scoped tasks complete even if unwinding due to a panic.
62pub(crate) struct ScopeGuard<'s, 'scope, 'env>(pub &'s Scope<'scope, 'env>);
63
64impl Drop for ScopeGuard<'_, '_, '_> {
65    fn drop(&mut self) {
66        self.0.wait();
67    }
68}
zero_pool/scope.rs

zero_pool/
scope.rs
