Skip to main content

moonpool_core/
task.rs

1//! Task spawning abstraction for single-threaded simulation environments.
2//!
3//! This module provides task provider abstractions for spawning local tasks
4//! that work with both simulation and real Tokio execution.
5
6use std::future::Future;
7
8/// Error returned by [`TaskProvider::JoinHandle`] when a task did not complete
9/// normally.
10///
11/// This is the runtime-agnostic error surfaced by the [`TaskProvider`] trait.
12/// Implementations convert their runtime-specific join error into one of these
13/// variants.
14#[derive(Debug, thiserror::Error)]
15pub enum JoinError {
16    /// The task was cancelled (for example, the runtime aborted it).
17    #[error("task was cancelled")]
18    Cancelled,
19    /// The task panicked.
20    #[error("task panicked")]
21    Panicked,
22}
23
24/// Provider for spawning tasks.
25///
26/// This trait abstracts task spawning to enable both real tokio tasks
27/// and simulation-controlled task scheduling. The simulation runtime
28/// runs on a single OS thread, but the spawned futures are Send-bounded
29/// so customer call graphs can use `Arc<RwLock<…>>`, `DashMap`, and other
30/// `Send + Sync` primitives without contortion.
31pub trait TaskProvider: Clone + Send + Sync + 'static {
32    /// Future returned by [`Self::spawn_task`].
33    ///
34    /// Resolves with `Ok(())` on normal completion, or a [`JoinError`] if the
35    /// task was cancelled or panicked.
36    type JoinHandle: Future<Output = Result<(), JoinError>> + Send + Sync + 'static;
37
38    /// Spawn a named task.
39    fn spawn_task<F>(&self, name: &str, future: F) -> Self::JoinHandle
40    where
41        F: Future<Output = ()> + Send + 'static;
42
43    /// Yield control to allow other tasks to run.
44    ///
45    /// This is equivalent to `tokio::task::yield_now()` but abstracted
46    /// to enable simulation control and deterministic behavior.
47    fn yield_now(&self) -> impl Future<Output = ()> + Send;
48}
49
50/// Tokio-based task provider.
51///
52/// This provider creates tasks via `tokio::spawn`. The task name is emitted
53/// in `tracing::trace!` spans around the future rather than attached to the
54/// tokio task itself.
55#[cfg(feature = "tokio-task")]
56#[derive(Clone, Debug)]
57pub struct TokioTaskProvider;
58
59/// `JoinHandle` produced by [`TokioTaskProvider`].
60///
61/// Wraps tokio's `JoinHandle<()>` and converts the runtime-specific
62/// `tokio::task::JoinError` into the runtime-agnostic [`JoinError`] variants
63/// when polled.
64#[cfg(feature = "tokio-task")]
65#[derive(Debug)]
66pub struct TokioJoinHandle(tokio::task::JoinHandle<()>);
67
68#[cfg(feature = "tokio-task")]
69impl Future for TokioJoinHandle {
70    type Output = Result<(), JoinError>;
71
72    fn poll(
73        mut self: std::pin::Pin<&mut Self>,
74        cx: &mut std::task::Context<'_>,
75    ) -> std::task::Poll<Self::Output> {
76        use std::task::Poll;
77        match std::pin::Pin::new(&mut self.0).poll(cx) {
78            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
79            Poll::Ready(Err(e)) if e.is_cancelled() => Poll::Ready(Err(JoinError::Cancelled)),
80            Poll::Ready(Err(_)) => Poll::Ready(Err(JoinError::Panicked)),
81            Poll::Pending => Poll::Pending,
82        }
83    }
84}
85
86#[cfg(feature = "tokio-task")]
87impl TaskProvider for TokioTaskProvider {
88    type JoinHandle = TokioJoinHandle;
89
90    fn spawn_task<F>(&self, name: &str, future: F) -> Self::JoinHandle
91    where
92        F: Future<Output = ()> + Send + 'static,
93    {
94        let task_name = name.to_string();
95        let fut = async move {
96            tracing::trace!("Task {} starting", task_name);
97            future.await;
98            tracing::trace!("Task {} completed", task_name);
99        };
100        TokioJoinHandle(tokio::spawn(fut))
101    }
102
103    async fn yield_now(&self) {
104        tokio::task::yield_now().await;
105    }
106}