zero_pool/
task_future.rs

1use std::sync::Arc;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::thread::{self, Thread};
4use std::time::{Duration, Instant};
5
6/// A future that tracks completion of submitted tasks
7///
8/// `TaskFuture` provides both blocking and non-blocking ways to wait for
9/// task completion. Tasks can be checked for completion, waited on
10/// indefinitely, or waited on with a timeout.
11///
12/// `TaskFuture` is cheaply cloneable. However, it captures the thread handle
13/// of the thread that created it.
14///
15/// If sharing the future with other threads, `is_complete()` is safe to call
16/// from anywhere.
17///
18/// **Important:** `wait()` and `wait_timeout()` **must** be called from the
19/// same thread that created the `TaskFuture`. Calling these methods from a
20/// different thread will panic in debug builds and may cause the calling thread
21/// to hang indefinitely in release builds.
22///
23#[derive(Clone)]
24pub struct TaskFuture {
25    count: Arc<AtomicUsize>,
26    owner_thread: Thread,
27}
28
29impl TaskFuture {
30    pub(crate) fn new(task_count: usize) -> Self {
31        TaskFuture {
32            count: Arc::new(AtomicUsize::new(task_count)),
33            owner_thread: thread::current(),
34        }
35    }
36
37    /// Check if all tasks are complete without blocking
38    ///
39    /// Returns `true` if all tasks have finished execution.
40    /// This is a non-blocking operation using an atomic load.
41    #[must_use]
42    #[inline]
43    pub fn is_complete(&self) -> bool {
44        self.count.load(Ordering::Acquire) == 0
45    }
46
47    /// Wait for all tasks to complete.
48    ///
49    /// First checks completion with an atomic load; if incomplete, parks the thread.
50    ///
51    /// **Warning:** This method must be called from the same thread that created this `TaskFuture`.
52    #[inline]
53    pub fn wait(&self) {
54        debug_assert_eq!(
55            self.owner_thread.id(),
56            thread::current().id(),
57            "TaskFuture::wait() must be called from the thread that created it."
58        );
59
60        while !self.is_complete() {
61            thread::park();
62        }
63    }
64
65    /// Wait for all tasks to complete with a timeout.
66    ///
67    /// First checks completion with an atomic load; if incomplete, parks the thread.
68    /// Returns `true` if all tasks completed within the timeout,
69    /// `false` if the timeout was reached first.
70    ///
71    /// **Warning:** This method must be called from the same thread that created this `TaskFuture`.
72    #[must_use]
73    #[inline]
74    pub fn wait_timeout(&self, timeout: Duration) -> bool {
75        debug_assert_eq!(
76            self.owner_thread.id(),
77            thread::current().id(),
78            "TaskFuture::wait_timeout() must be called from the thread that created it."
79        );
80
81        let start = Instant::now();
82        loop {
83            if self.is_complete() {
84                return true;
85            }
86            let elapsed = start.elapsed();
87            if elapsed >= timeout {
88                return false;
89            }
90            thread::park_timeout(timeout.saturating_sub(elapsed));
91        }
92    }
93
94    // completes multiple tasks, decrements counter and notifies if all done
95    pub(crate) fn complete_many(&self, count: usize) {
96        if self.count.fetch_sub(count, Ordering::Release) == count {
97            self.owner_thread.unpark();
98        }
99    }
100}
zero_pool/task_future.rs

zero_pool/
task_future.rs
