reth_tasks/
lib.rs

1//! Reth task management.
2//!
3//! # Feature Flags
4//!
5//! - `rayon`: Enable rayon thread pool for blocking tasks.
6
7#![doc(
8    html_logo_url = "https://raw.githubusercontent.com/paradigmxyz/reth/main/assets/reth-docs.png",
9    html_favicon_url = "https://avatars0.githubusercontent.com/u/97369466?s=256",
10    issue_tracker_base_url = "https://github.com/paradigmxyz/reth/issues/"
11)]
12#![cfg_attr(not(test), warn(unused_crate_dependencies))]
13#![cfg_attr(docsrs, feature(doc_cfg))]
14
15use crate::shutdown::{signal, Shutdown, Signal};
16use std::{
17    any::Any,
18    fmt::{Display, Formatter},
19    pin::Pin,
20    sync::{
21        atomic::{AtomicUsize, Ordering},
22        Arc,
23    },
24    task::{ready, Context, Poll},
25    thread,
26};
27use tokio::{
28    runtime::Handle,
29    sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
30    task::JoinHandle,
31};
32use tracing::debug;
33
34pub mod cancel;
35pub mod lazy;
36pub mod metrics;
37pub mod runtime;
38pub mod shutdown;
39pub mod utils;
40pub(crate) mod worker_map;
41
42#[cfg(feature = "rayon")]
43pub mod pool;
44#[cfg(feature = "rayon")]
45pub use pool::{build_pool_with_panic_handler, Worker, WorkerPool, WorkerStatefulPool};
46
47/// Lock-free ordered parallel iterator extension trait.
48#[cfg(feature = "rayon")]
49pub mod for_each_ordered;
50#[cfg(feature = "rayon")]
51pub use for_each_ordered::ForEachOrdered;
52
53pub use cancel::{CancelOnDrop, ManualCancel};
54pub use lazy::LazyHandle;
55#[cfg(feature = "rayon")]
56pub use runtime::RayonConfig;
57pub use runtime::{Runtime, RuntimeBuildError, RuntimeBuilder, RuntimeConfig, TokioConfig};
58
59/// A [`TaskExecutor`] is now an alias for [`Runtime`].
60pub type TaskExecutor = Runtime;
61
62/// Spawns an OS thread with the current tokio runtime context propagated.
63///
64/// This function captures the current tokio runtime handle (if available) and enters it
65/// in the newly spawned thread. This ensures that code running in the spawned thread can
66/// use [`Handle::current()`], [`Handle::spawn_blocking()`], and other tokio utilities that
67/// require a runtime context.
68#[track_caller]
69pub fn spawn_os_thread<F, T>(name: &str, f: F) -> thread::JoinHandle<T>
70where
71    F: FnOnce() -> T + Send + 'static,
72    T: Send + 'static,
73{
74    let handle = Handle::try_current().ok();
75    thread::Builder::new()
76        .name(name.to_string())
77        .spawn(move || {
78            let _guard = handle.as_ref().map(Handle::enter);
79            f()
80        })
81        .unwrap_or_else(|e| panic!("failed to spawn thread {name:?}: {e}"))
82}
83
84/// Spawns a scoped OS thread with the current tokio runtime context propagated.
85///
86/// This is the scoped thread version of [`spawn_os_thread`], for use with [`std::thread::scope`].
87#[track_caller]
88pub fn spawn_scoped_os_thread<'scope, 'env, F, T>(
89    scope: &'scope thread::Scope<'scope, 'env>,
90    name: &str,
91    f: F,
92) -> thread::ScopedJoinHandle<'scope, T>
93where
94    F: FnOnce() -> T + Send + 'scope,
95    T: Send + 'scope,
96{
97    let handle = Handle::try_current().ok();
98    thread::Builder::new()
99        .name(name.to_string())
100        .spawn_scoped(scope, move || {
101            let _guard = handle.as_ref().map(Handle::enter);
102            f()
103        })
104        .unwrap_or_else(|e| panic!("failed to spawn scoped thread {name:?}: {e}"))
105}
106
107/// Spawns a blocking closure on the current tokio runtime, like [`tokio::task::spawn_blocking`].
108///
109/// The closure is boxed first, so tokio's blocking task machinery is instantiated once per return
110/// type `R` instead of once per closure type.
111///
112/// # Panics
113///
114/// Panics if called outside of a tokio runtime context.
115pub fn spawn_blocking<F, R>(func: F) -> JoinHandle<R>
116where
117    F: FnOnce() -> R + Send + 'static,
118    R: Send + 'static,
119{
120    tokio::task::spawn_blocking(Box::new(func) as Box<dyn FnOnce() -> R + Send>)
121}
122
123/// Monitors critical tasks for panics and manages graceful shutdown.
124///
125/// The main purpose of this type is to be able to monitor if a critical task panicked, for
126/// diagnostic purposes, since tokio tasks essentially fail silently. Therefore, this type is a
127/// Future that resolves with the name of the panicked task. See [`Runtime::spawn_critical_task`].
128///
129/// Automatically spawned as a background task when building a [`Runtime`]. Use
130/// [`Runtime::take_task_manager_handle`] to extract the join handle if you need to poll for
131/// panic errors directly.
132#[derive(Debug)]
133#[must_use = "TaskManager must be polled to monitor critical tasks"]
134pub struct TaskManager {
135    /// Receiver for task events.
136    task_events_rx: UnboundedReceiver<TaskEvent>,
137    /// The [Signal] to fire when all tasks should be shutdown.
138    ///
139    /// This is fired when dropped.
140    signal: Option<Signal>,
141    /// How many [`GracefulShutdown`](crate::shutdown::GracefulShutdown) tasks are currently
142    /// active.
143    graceful_tasks: Arc<AtomicUsize>,
144}
145
146// === impl TaskManager ===
147
148impl TaskManager {
149    /// Create a new [`TaskManager`] without an associated [`Runtime`], returning
150    /// the shutdown/event primitives for [`RuntimeBuilder`] to wire up.
151    pub(crate) fn new_parts(
152        _handle: Handle,
153    ) -> (Self, Shutdown, UnboundedSender<TaskEvent>, Arc<AtomicUsize>) {
154        let (task_events_tx, task_events_rx) = unbounded_channel();
155        let (signal, on_shutdown) = signal();
156        let graceful_tasks = Arc::new(AtomicUsize::new(0));
157        let manager = Self {
158            task_events_rx,
159            signal: Some(signal),
160            graceful_tasks: Arc::clone(&graceful_tasks),
161        };
162        (manager, on_shutdown, task_events_tx, graceful_tasks)
163    }
164
165    /// Fires the shutdown signal and awaits until all tasks are shutdown.
166    pub fn graceful_shutdown(self) {
167        let _ = self.do_graceful_shutdown(None);
168    }
169
170    /// Fires the shutdown signal and awaits until all tasks are shutdown.
171    ///
172    /// Returns true if all tasks were shutdown before the timeout elapsed.
173    pub fn graceful_shutdown_with_timeout(self, timeout: std::time::Duration) -> bool {
174        self.do_graceful_shutdown(Some(timeout))
175    }
176
177    fn do_graceful_shutdown(self, timeout: Option<std::time::Duration>) -> bool {
178        drop(self.signal);
179        let deadline = timeout.map(|t| std::time::Instant::now() + t);
180        while self.graceful_tasks.load(Ordering::SeqCst) > 0 {
181            if deadline.is_some_and(|d| std::time::Instant::now() > d) {
182                debug!("graceful shutdown timed out");
183                return false;
184            }
185            thread::yield_now();
186        }
187        debug!("gracefully shut down");
188        true
189    }
190}
191
192/// An endless future that resolves if a critical task panicked.
193///
194/// See [`Runtime::spawn_critical_task`]
195impl std::future::Future for TaskManager {
196    type Output = Result<(), PanickedTaskError>;
197
198    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
199        match ready!(self.as_mut().get_mut().task_events_rx.poll_recv(cx)) {
200            Some(TaskEvent::Panic(err)) => Poll::Ready(Err(err)),
201            Some(TaskEvent::GracefulShutdown) | None => {
202                if let Some(signal) = self.get_mut().signal.take() {
203                    signal.fire();
204                }
205                Poll::Ready(Ok(()))
206            }
207        }
208    }
209}
210
211/// Error with the name of the task that panicked and an error downcasted to string, if possible.
212#[derive(Debug, thiserror::Error, PartialEq, Eq)]
213pub struct PanickedTaskError {
214    task_name: &'static str,
215    error: Option<String>,
216}
217
218impl Display for PanickedTaskError {
219    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
220        let task_name = self.task_name;
221        if let Some(error) = &self.error {
222            write!(f, "Critical task `{task_name}` panicked: `{error}`")
223        } else {
224            write!(f, "Critical task `{task_name}` panicked")
225        }
226    }
227}
228
229impl PanickedTaskError {
230    pub(crate) fn new(task_name: &'static str, error: Box<dyn Any>) -> Self {
231        let error = match error.downcast::<String>() {
232            Ok(value) => Some(*value),
233            Err(error) => match error.downcast::<&str>() {
234                Ok(value) => Some(value.to_string()),
235                Err(_) => None,
236            },
237        };
238
239        Self { task_name, error }
240    }
241}
242
243/// Represents the events that the `TaskManager`'s main future can receive.
244#[derive(Debug)]
245pub(crate) enum TaskEvent {
246    /// Indicates that a critical task has panicked.
247    Panic(PanickedTaskError),
248    /// A signal requesting a graceful shutdown of the `TaskManager`.
249    GracefulShutdown,
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use std::{
256        sync::atomic::{AtomicBool, AtomicUsize, Ordering},
257        time::Duration,
258    };
259
260    #[test]
261    fn test_critical() {
262        let rt = Runtime::test();
263        let handle = rt.take_task_manager_handle().unwrap();
264
265        rt.spawn_critical_task("this is a critical task", async { panic!("intentionally panic") });
266
267        rt.handle().block_on(async move {
268            let err_result = handle.await.unwrap();
269            assert!(err_result.is_err(), "Expected TaskManager to return an error due to panic");
270            let panicked_err = err_result.unwrap_err();
271
272            assert_eq!(panicked_err.task_name, "this is a critical task");
273            assert_eq!(panicked_err.error, Some("intentionally panic".to_string()));
274        })
275    }
276
277    #[test]
278    fn test_manager_shutdown_critical() {
279        let rt = Runtime::test();
280
281        let (signal, shutdown) = signal();
282
283        rt.spawn_critical_task("this is a critical task", async move {
284            tokio::time::sleep(Duration::from_millis(200)).await;
285            drop(signal);
286        });
287
288        rt.graceful_shutdown();
289
290        rt.handle().block_on(shutdown);
291    }
292
293    #[test]
294    fn test_manager_shutdown() {
295        let rt = Runtime::test();
296
297        let (signal, shutdown) = signal();
298
299        rt.spawn_task(async move {
300            tokio::time::sleep(Duration::from_millis(200)).await;
301            drop(signal);
302        });
303
304        rt.graceful_shutdown();
305
306        rt.handle().block_on(shutdown);
307    }
308
309    #[test]
310    fn test_manager_graceful_shutdown() {
311        let rt = Runtime::test();
312
313        let val = Arc::new(AtomicBool::new(false));
314        let c = val.clone();
315        rt.spawn_critical_with_graceful_shutdown_signal("grace", async move |shutdown| {
316            let _guard = shutdown.await;
317            tokio::time::sleep(Duration::from_millis(200)).await;
318            c.store(true, Ordering::Relaxed);
319        });
320
321        rt.graceful_shutdown();
322        assert!(val.load(Ordering::Relaxed));
323    }
324
325    #[test]
326    fn test_manager_graceful_shutdown_many() {
327        let rt = Runtime::test();
328
329        let counter = Arc::new(AtomicUsize::new(0));
330        let num = 10;
331        for _ in 0..num {
332            let c = counter.clone();
333            rt.spawn_critical_with_graceful_shutdown_signal("grace", async move |shutdown| {
334                let _guard = shutdown.await;
335                tokio::time::sleep(Duration::from_millis(200)).await;
336                c.fetch_add(1, Ordering::SeqCst);
337            });
338        }
339
340        rt.graceful_shutdown();
341        assert_eq!(counter.load(Ordering::Relaxed), num);
342    }
343
344    #[test]
345    fn test_manager_graceful_shutdown_timeout() {
346        let rt = Runtime::test();
347
348        let timeout = Duration::from_millis(500);
349        let val = Arc::new(AtomicBool::new(false));
350        let val2 = val.clone();
351        rt.spawn_critical_with_graceful_shutdown_signal("grace", async move |shutdown| {
352            let _guard = shutdown.await;
353            tokio::time::sleep(timeout * 3).await;
354            val2.store(true, Ordering::Relaxed);
355            unreachable!("should not be reached");
356        });
357
358        rt.graceful_shutdown_with_timeout(timeout);
359        assert!(!val.load(Ordering::Relaxed));
360    }
361
362    #[test]
363    fn can_build_runtime() {
364        let rt = Runtime::test();
365        let _handle = rt.handle();
366    }
367
368    #[test]
369    fn test_graceful_shutdown_triggered_by_executor() {
370        let rt = Runtime::test();
371        let task_manager_handle = rt.take_task_manager_handle().unwrap();
372
373        let task_did_shutdown_flag = Arc::new(AtomicBool::new(false));
374        let flag_clone = task_did_shutdown_flag.clone();
375
376        let spawned_task_handle = rt.spawn_with_signal(async move |shutdown_signal| {
377            shutdown_signal.await;
378            flag_clone.store(true, Ordering::SeqCst);
379        });
380
381        let send_result = rt.initiate_graceful_shutdown();
382        assert!(send_result.is_ok());
383
384        let manager_final_result = rt.handle().block_on(task_manager_handle);
385        assert!(manager_final_result.is_ok(), "TaskManager task should not panic");
386        assert_eq!(manager_final_result.unwrap(), Ok(()));
387
388        let task_join_result = rt.handle().block_on(spawned_task_handle);
389        assert!(task_join_result.is_ok());
390
391        assert!(task_did_shutdown_flag.load(Ordering::Relaxed));
392    }
393}
reth_tasks/lib.rs

reth_tasks/
lib.rs
