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}