moonpool_sim/executor/
join.rs

1//! Join handles and yield primitives for the deterministic executor.
2
3use std::any::Any;
4use std::future::Future;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8use async_task::FallibleTask;
9use moonpool_core::JoinError;
10use parking_lot::Mutex;
11
12use super::TaskMeta;
13
14/// A task's raw output: the value, or the payload of the panic that killed it
15/// (captured by the `catch_unwind` wrapper installed at spawn).
16pub(crate) type TaskResult<T> = Result<T, Box<dyn Any + Send>>;
17
18/// Owned handle to a spawned task, mirroring `tokio::task::JoinHandle`
19/// semantics.
20///
21/// - `.await` resolves to `Ok(output)`, or `Err(JoinError::Cancelled)` after
22///   [`abort`](Self::abort), or `Err(JoinError::Panicked)` if the task
23///   panicked (siblings and the executor are unaffected).
24/// - Dropping the handle **detaches** the task: it keeps running to
25///   completion. Only [`abort`](Self::abort) or dropping the whole
26///   [`Executor`](super::Executor) cancels it.
27/// - [`is_finished`](Self::is_finished) is a non-blocking completion probe.
28#[derive(Debug)]
29pub struct JoinHandle<T> {
30    /// `None` after `abort()` consumed the inner task (async-task cancels a
31    /// task when its `Task`/`FallibleTask` handle is dropped). The `Mutex`
32    /// (uncontended, single-threaded) exists so `abort(&self)` can take the
33    /// inner value while the handle satisfies the `Sync` bound that
34    /// `TaskProvider::JoinHandle` requires.
35    inner: Mutex<Option<FallibleTask<TaskResult<T>, TaskMeta>>>,
36}
37
38impl<T> JoinHandle<T> {
39    pub(crate) fn new(task: FallibleTask<TaskResult<T>, TaskMeta>) -> Self {
40        Self {
41            inner: Mutex::new(Some(task)),
42        }
43    }
44
45    /// Report whether the task has finished running (completed, panicked, or
46    /// was aborted). Non-blocking; mirrors `tokio::task::JoinHandle::is_finished`.
47    #[must_use]
48    pub fn is_finished(&self) -> bool {
49        match self.inner.lock().as_ref() {
50            Some(task) => task.is_finished(),
51            None => true,
52        }
53    }
54
55    /// Abort the task: it will never be polled again and its future is
56    /// dropped. Awaiting the handle afterwards yields
57    /// `Err(JoinError::Cancelled)`.
58    ///
59    /// Aborting a task that already finished is a no-op and awaiting the
60    /// handle still yields the task's output (tokio parity: abort racing
61    /// completion loses).
62    pub fn abort(&self) {
63        let mut inner = self.inner.lock();
64        // A finished task keeps its output harvestable; only a live task is
65        // cancelled (dropping the FallibleTask cancels, async-task semantics).
66        if inner.as_ref().is_some_and(|task| !task.is_finished()) {
67            drop(inner.take());
68        }
69    }
70}
71
72impl<T> Future for JoinHandle<T> {
73    type Output = Result<T, JoinError>;
74
75    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
76        let mut inner = self.inner.lock();
77        match inner.as_mut() {
78            // Aborted through this handle.
79            None => Poll::Ready(Err(JoinError::Cancelled)),
80            Some(task) => match Pin::new(task).poll(cx) {
81                Poll::Pending => Poll::Pending,
82                // Cancelled elsewhere (executor drop) before completing.
83                Poll::Ready(None) => Poll::Ready(Err(JoinError::Cancelled)),
84                Poll::Ready(Some(Ok(output))) => Poll::Ready(Ok(output)),
85                Poll::Ready(Some(Err(_panic_payload))) => Poll::Ready(Err(JoinError::Panicked)),
86            },
87        }
88    }
89}
90
91impl<T> Drop for JoinHandle<T> {
92    fn drop(&mut self) {
93        // Detach-on-drop (tokio parity): the task keeps running.
94        if let Some(task) = self.inner.lock().take() {
95            task.detach();
96        }
97    }
98}
99
100/// Future returned by [`yield_now`] and [`until_stalled`]: pends once,
101/// waking its own waker, then completes on the next poll.
102#[derive(Debug)]
103pub struct YieldNow {
104    yielded: bool,
105}
106
107impl Future for YieldNow {
108    type Output = ();
109
110    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
111        if self.yielded {
112            Poll::Ready(())
113        } else {
114            self.yielded = true;
115            cx.waker().wake_by_ref();
116            Poll::Pending
117        }
118    }
119}
120
121/// Yield control from inside a task.
122///
123/// The task is rescheduled into the ready queue and resumes at a
124/// seeded-random position among the runnable tasks (there is no "back of
125/// the queue" under randomized scheduling).
126#[must_use = "futures do nothing unless awaited"]
127pub fn yield_now() -> YieldNow {
128    YieldNow { yielded: false }
129}
130
131/// Driver-only: resume after every currently-runnable task has been polled
132/// to `Pending` or completion.
133///
134/// Awaited by the orchestrator's step loops between `sim.step()` calls.
135/// Mechanically identical to [`yield_now`], but because
136/// [`Executor::block_on`](super::Executor::block_on) only re-polls the
137/// driver after a full `run_until_stalled` drain, awaiting this from the
138/// driver guarantees the "all runnable tasks ran" contract that the old
139/// tokio-FIFO `yield_now` dance only approximated.
140///
141/// # Panics
142///
143/// Panics when called from inside a task (all builds): a task awaiting
144/// "drain until nothing is runnable" would silently degrade to plain
145/// [`yield_now`] semantics, resuming mid-drain while believing every
146/// runnable task quiesced. That contract violation must not survive only
147/// in release. Tasks yield with [`yield_now`].
148#[must_use = "futures do nothing unless awaited"]
149pub fn until_stalled() -> YieldNow {
150    assert!(
151        !super::in_task(),
152        "until_stalled() is driver-only; tasks must use executor::yield_now()"
153    );
154    YieldNow { yielded: false }
155}
moonpool_sim/executor/join.rs

moonpool_sim/executor/
join.rs
