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    pub fn is_complete(&self) -> bool {
42        self.count.load(Ordering::Acquire) == 0
43    }
44
45    /// Wait for all tasks to complete.
46    ///
47    /// First checks completion with an atomic load; if incomplete, parks the thread.
48    ///
49    /// **Warning:** This method must be called from the same thread that created this `TaskFuture`.
50    pub fn wait(&self) {
51        debug_assert_eq!(
52            self.owner_thread.id(),
53            thread::current().id(),
54            "TaskFuture::wait() must be called from the thread that created it."
55        );
56
57        while !self.is_complete() {
58            thread::park();
59        }
60    }
61
62    /// Wait for all tasks to complete with a timeout.
63    ///
64    /// First checks completion with an atomic load; if incomplete, parks the thread.
65    /// Returns `true` if all tasks completed within the timeout,
66    /// `false` if the timeout was reached first.
67    ///
68    /// **Warning:** This method must be called from the same thread that created this `TaskFuture`.
69    pub fn wait_timeout(&self, timeout: Duration) -> bool {
70        debug_assert_eq!(
71            self.owner_thread.id(),
72            thread::current().id(),
73            "TaskFuture::wait_timeout() must be called from the thread that created it."
74        );
75
76        let start = Instant::now();
77        loop {
78            if self.is_complete() {
79                return true;
80            }
81            let elapsed = start.elapsed();
82            if elapsed >= timeout {
83                return false;
84            }
85            thread::park_timeout(timeout - elapsed);
86        }
87    }
88
89    // completes multiple tasks, decrements counter and notifies if all done
90    pub(crate) fn complete_many(&self, count: usize) {
91        if self.count.fetch_sub(count, Ordering::Release) == count {
92            self.owner_thread.unpark();
93        }
94    }
95}
zero_pool/task_future.rs

zero_pool/
task_future.rs
