Skip to main content

reth_tasks/
runtime.rs

1//! Centralized management of async and parallel execution.
2//!
3//! This module provides [`Runtime`], a cheaply cloneable handle that manages:
4//! - Tokio runtime (either owned or attached)
5//! - Task spawning with shutdown awareness and panic monitoring
6//! - Dedicated rayon thread pools for different workloads (with `rayon` feature)
7//! - [`BlockingTaskGuard`] for rate-limiting expensive operations (with `rayon` feature)
8
9#[cfg(feature = "rayon")]
10use crate::pool::{build_pool_with_panic_handler, BlockingTaskGuard, BlockingTaskPool, WorkerPool};
11use crate::{
12    metrics::{IncCounterOnDrop, TaskExecutorMetrics},
13    shutdown::{GracefulShutdown, GracefulShutdownGuard, Shutdown},
14    worker_map::WorkerMap,
15    PanickedTaskError, TaskEvent, TaskManager,
16};
17use futures_util::{future::select, Future, FutureExt, TryFutureExt};
18use std::{
19    num::NonZeroUsize,
20    pin::pin,
21    sync::{
22        atomic::{AtomicUsize, Ordering},
23        Arc, Mutex,
24    },
25    thread,
26    time::{Duration, Instant},
27};
28use tokio::{runtime::Handle, sync::mpsc::UnboundedSender, task::JoinHandle};
29use tracing::{debug, error};
30use tracing_futures::Instrument;
31
32use tokio::runtime::Runtime as TokioRuntime;
33
34#[cfg(feature = "rayon")]
35use std::thread::available_parallelism;
36
37/// Default thread keep-alive duration for the tokio runtime.
38pub const DEFAULT_THREAD_KEEP_ALIVE: Duration = Duration::from_secs(15);
39
40/// Default reserved CPU cores for OS and other processes.
41pub const DEFAULT_RESERVED_CPU_CORES: usize = 2;
42
43/// Default number of threads for the storage I/O pool.
44pub const DEFAULT_STORAGE_POOL_THREADS: usize = 16;
45
46/// Default number of threads for the state trie overlay worker pool.
47#[cfg(feature = "rayon")]
48pub const DEFAULT_STATE_TRIE_OVERLAY_WORKER_THREADS: usize = 4;
49
50/// Default maximum number of concurrent blocking tasks (for RPC tracing guard).
51pub const DEFAULT_MAX_BLOCKING_TASKS: usize = 512;
52
53/// Configuration for the tokio runtime.
54#[derive(Debug, Clone)]
55pub enum TokioConfig {
56    /// Build and own a new tokio runtime.
57    Owned {
58        /// Number of worker threads. If `None`, uses tokio's default (number of CPU cores).
59        worker_threads: Option<usize>,
60        /// How long to keep worker threads alive when idle.
61        thread_keep_alive: Duration,
62        /// Thread name prefix.
63        thread_name: &'static str,
64    },
65    /// Attach to an existing tokio runtime handle.
66    ExistingHandle(Handle),
67}
68
69impl Default for TokioConfig {
70    fn default() -> Self {
71        Self::Owned {
72            worker_threads: None,
73            thread_keep_alive: DEFAULT_THREAD_KEEP_ALIVE,
74            thread_name: "tokio-rt",
75        }
76    }
77}
78
79impl TokioConfig {
80    /// Create a config that attaches to an existing runtime handle.
81    pub const fn existing_handle(handle: Handle) -> Self {
82        Self::ExistingHandle(handle)
83    }
84
85    /// Create a config for an owned runtime with the specified number of worker threads.
86    pub const fn with_worker_threads(worker_threads: usize) -> Self {
87        Self::Owned {
88            worker_threads: Some(worker_threads),
89            thread_keep_alive: DEFAULT_THREAD_KEEP_ALIVE,
90            thread_name: "tokio-rt",
91        }
92    }
93}
94
95/// Configuration for the rayon thread pools.
96#[derive(Debug, Clone)]
97#[cfg(feature = "rayon")]
98pub struct RayonConfig {
99    /// Number of threads for the general CPU pool.
100    /// If `None`, derived from available parallelism minus reserved cores.
101    pub cpu_threads: Option<usize>,
102    /// Number of CPU cores to reserve for OS and other processes.
103    pub reserved_cpu_cores: usize,
104    /// Number of threads for the RPC blocking pool (trace calls, `eth_getProof`, etc.).
105    /// If `None`, uses the same as `cpu_threads`.
106    pub rpc_threads: Option<usize>,
107    /// Number of threads for the storage I/O pool (static file, `RocksDB` writes in
108    /// `save_blocks`). If `None`, uses [`DEFAULT_STORAGE_POOL_THREADS`].
109    pub storage_threads: Option<usize>,
110    /// Maximum number of concurrent blocking tasks for the RPC guard semaphore.
111    pub max_blocking_tasks: usize,
112    /// Number of threads for the proof storage worker pool (trie storage proof workers).
113    /// If `None` or zero, uses four times the resolved CPU pool thread count.
114    pub proof_storage_worker_threads: Option<usize>,
115    /// Number of threads for the proof account worker pool (trie account proof workers).
116    /// If `None` or zero, uses four times the resolved CPU pool thread count.
117    pub proof_account_worker_threads: Option<usize>,
118    /// Number of threads for the prewarming pool (execution prewarming workers).
119    /// If `None`, derived from available parallelism.
120    pub prewarming_threads: Option<usize>,
121    /// Number of threads for the BAL streaming pool (BAL hashed state streaming).
122    /// If `None`, derived from available parallelism.
123    pub bal_streaming_threads: Option<usize>,
124    /// Number of threads for the state trie overlay worker pool.
125    /// If `None`, uses [`DEFAULT_STATE_TRIE_OVERLAY_WORKER_THREADS`].
126    pub state_trie_overlay_worker_threads: Option<usize>,
127}
128
129#[cfg(feature = "rayon")]
130impl Default for RayonConfig {
131    fn default() -> Self {
132        Self {
133            cpu_threads: None,
134            reserved_cpu_cores: DEFAULT_RESERVED_CPU_CORES,
135            rpc_threads: None,
136            storage_threads: None,
137            max_blocking_tasks: DEFAULT_MAX_BLOCKING_TASKS,
138            proof_storage_worker_threads: None,
139            proof_account_worker_threads: None,
140            prewarming_threads: None,
141            bal_streaming_threads: None,
142            state_trie_overlay_worker_threads: None,
143        }
144    }
145}
146
147#[cfg(feature = "rayon")]
148impl RayonConfig {
149    /// Set the number of reserved CPU cores.
150    pub const fn with_reserved_cpu_cores(mut self, reserved_cpu_cores: usize) -> Self {
151        self.reserved_cpu_cores = reserved_cpu_cores;
152        self
153    }
154
155    /// Set the maximum number of concurrent blocking tasks.
156    pub const fn with_max_blocking_tasks(mut self, max_blocking_tasks: usize) -> Self {
157        self.max_blocking_tasks = max_blocking_tasks;
158        self
159    }
160
161    /// Set the number of threads for the RPC blocking pool.
162    pub const fn with_rpc_threads(mut self, rpc_threads: usize) -> Self {
163        self.rpc_threads = Some(rpc_threads);
164        self
165    }
166
167    /// Set the number of threads for the storage I/O pool.
168    pub const fn with_storage_threads(mut self, storage_threads: usize) -> Self {
169        self.storage_threads = Some(storage_threads);
170        self
171    }
172
173    /// Set the number of threads for the proof storage worker pool.
174    pub const fn with_proof_storage_worker_threads(
175        mut self,
176        proof_storage_worker_threads: usize,
177    ) -> Self {
178        self.proof_storage_worker_threads = Some(proof_storage_worker_threads);
179        self
180    }
181
182    /// Set the number of threads for the proof account worker pool.
183    pub const fn with_proof_account_worker_threads(
184        mut self,
185        proof_account_worker_threads: usize,
186    ) -> Self {
187        self.proof_account_worker_threads = Some(proof_account_worker_threads);
188        self
189    }
190
191    /// Set the number of threads for the prewarming pool.
192    pub const fn with_prewarming_threads(mut self, prewarming_threads: usize) -> Self {
193        self.prewarming_threads = Some(prewarming_threads);
194        self
195    }
196
197    /// Set the number of threads for the BAL streaming pool.
198    pub const fn with_bal_streaming_threads(mut self, bal_streaming_threads: usize) -> Self {
199        self.bal_streaming_threads = Some(bal_streaming_threads);
200        self
201    }
202
203    /// Set the number of threads for the state trie overlay worker pool.
204    pub const fn with_state_trie_overlay_worker_threads(
205        mut self,
206        state_trie_overlay_worker_threads: usize,
207    ) -> Self {
208        self.state_trie_overlay_worker_threads = Some(state_trie_overlay_worker_threads);
209        self
210    }
211
212    /// Compute the default number of threads based on available parallelism.
213    fn default_thread_count(&self) -> usize {
214        // TODO: reserved_cpu_cores is currently ignored because subtracting from thread pool
215        // sizes doesn't actually reserve CPU cores for other processes.
216        let _ = self.reserved_cpu_cores;
217        self.cpu_threads.unwrap_or_else(|| available_parallelism().map_or(1, NonZeroUsize::get))
218    }
219
220    /// Returns the explicitly configured proof storage worker thread count, treating zero as
221    /// unset.
222    fn proof_storage_worker_threads_override(&self) -> Option<usize> {
223        self.proof_storage_worker_threads.filter(|&threads| threads > 0)
224    }
225
226    /// Returns the explicitly configured proof account worker thread count, treating zero as
227    /// unset.
228    fn proof_account_worker_threads_override(&self) -> Option<usize> {
229        self.proof_account_worker_threads.filter(|&threads| threads > 0)
230    }
231}
232
233/// Configuration for building a [`Runtime`].
234#[derive(Debug, Clone, Default)]
235pub struct RuntimeConfig {
236    /// Tokio runtime configuration.
237    pub tokio: TokioConfig,
238    /// Rayon thread pool configuration.
239    #[cfg(feature = "rayon")]
240    pub rayon: RayonConfig,
241}
242
243impl RuntimeConfig {
244    /// Set the tokio configuration.
245    pub fn with_tokio(mut self, tokio: TokioConfig) -> Self {
246        self.tokio = tokio;
247        self
248    }
249
250    /// Set the rayon configuration.
251    #[cfg(feature = "rayon")]
252    pub const fn with_rayon(mut self, rayon: RayonConfig) -> Self {
253        self.rayon = rayon;
254        self
255    }
256
257    /// Override the available CPU count used to size thread pools.
258    ///
259    /// Explicit thread counts and externally supplied Tokio runtimes are preserved.
260    /// Fixed-size pools, such as the storage I/O pool, are not affected.
261    pub fn with_cpu_cores(mut self, cpu_cores: NonZeroUsize) -> Self {
262        let cpu_cores = cpu_cores.get();
263        if let TokioConfig::Owned { worker_threads, .. } = &mut self.tokio {
264            worker_threads.get_or_insert(cpu_cores);
265        }
266        #[cfg(feature = "rayon")]
267        self.rayon.cpu_threads.get_or_insert(cpu_cores);
268        self
269    }
270}
271
272/// Error returned when [`RuntimeBuilder::build`] fails.
273#[derive(Debug, thiserror::Error)]
274pub enum RuntimeBuildError {
275    /// Failed to build the tokio runtime.
276    #[error("Failed to build tokio runtime: {0}")]
277    TokioBuild(#[from] std::io::Error),
278    /// Failed to build a rayon thread pool.
279    #[cfg(feature = "rayon")]
280    #[error("Failed to build rayon thread pool: {0}")]
281    RayonBuild(#[from] rayon::ThreadPoolBuildError),
282}
283
284// ── RuntimeInner ──────────────────────────────────────────────────────
285
286struct RuntimeInner {
287    /// Owned tokio runtime, taken during shutdown even while other handles remain alive.
288    tokio_runtime: Mutex<Option<TokioRuntime>>,
289    /// Handle to the tokio runtime.
290    handle: Handle,
291    /// Receiver of the shutdown signal.
292    on_shutdown: Shutdown,
293    /// Sender half for sending task events to the [`TaskManager`].
294    task_events_tx: UnboundedSender<TaskEvent>,
295    /// Task executor metrics.
296    metrics: TaskExecutorMetrics,
297    /// How many [`GracefulShutdown`] tasks are currently active.
298    graceful_tasks: Arc<AtomicUsize>,
299    /// General-purpose rayon CPU pool.
300    #[cfg(feature = "rayon")]
301    cpu_pool: rayon::ThreadPool,
302    /// RPC blocking pool.
303    #[cfg(feature = "rayon")]
304    rpc_pool: BlockingTaskPool,
305    /// Storage I/O pool.
306    #[cfg(feature = "rayon")]
307    storage_pool: rayon::ThreadPool,
308    /// Rate limiter for expensive RPC operations.
309    #[cfg(feature = "rayon")]
310    blocking_guard: BlockingTaskGuard,
311    /// Proof storage worker pool (trie storage proof computation).
312    #[cfg(feature = "rayon")]
313    proof_storage_worker_pool: WorkerPool,
314    /// Proof account worker pool (trie account proof computation).
315    #[cfg(feature = "rayon")]
316    proof_account_worker_pool: WorkerPool,
317    /// Proof storage worker thread count the operator configured explicitly, if any.
318    #[cfg(feature = "rayon")]
319    proof_storage_worker_threads_override: Option<usize>,
320    /// Proof account worker thread count the operator configured explicitly, if any.
321    #[cfg(feature = "rayon")]
322    proof_account_worker_threads_override: Option<usize>,
323    /// Prewarming pool (execution prewarming workers).
324    #[cfg(feature = "rayon")]
325    prewarming_pool: WorkerPool,
326    /// BAL streaming pool (BAL hashed state streaming).
327    #[cfg(feature = "rayon")]
328    bal_streaming_pool: WorkerPool,
329    /// State trie overlay worker pool.
330    #[cfg(feature = "rayon")]
331    state_trie_overlay_worker_pool: Arc<WorkerPool>,
332    /// Named single-thread worker map. Each unique name gets a dedicated OS thread
333    /// that is reused across all tasks submitted under that name.
334    worker_map: WorkerMap,
335    /// Handle to the spawned [`TaskManager`] background task.
336    /// The task monitors critical tasks for panics and fires the shutdown signal.
337    /// Can be taken via [`Runtime::take_task_manager_handle`] to poll for panic errors.
338    task_manager_handle: Mutex<Option<JoinHandle<Result<(), PanickedTaskError>>>>,
339}
340
341impl Drop for RuntimeInner {
342    fn drop(&mut self) {
343        // The last handle can belong to an async task, including one running on this runtime.
344        // Without explicit shutdown, let blocking tasks finish in the background in that case.
345        if Handle::try_current().is_ok() &&
346            let Some(runtime) = self.tokio_runtime.get_mut().unwrap().take()
347        {
348            runtime.shutdown_background();
349        }
350    }
351}
352
353// ── Runtime ───────────────────────────────────────────────────────────
354
355/// A cheaply cloneable handle to the runtime resources.
356///
357/// Wraps an `Arc<RuntimeInner>` and provides access to:
358/// - The tokio [`Handle`]
359/// - Task spawning with shutdown awareness and panic monitoring
360/// - Rayon thread pools (with `rayon` feature)
361#[derive(Clone)]
362pub struct Runtime(Arc<RuntimeInner>);
363
364impl std::fmt::Debug for Runtime {
365    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
366        f.debug_struct("Runtime").field("handle", &self.0.handle).finish()
367    }
368}
369
370// ── Pool accessors ────────────────────────────────────────────────────
371
372impl Runtime {
373    /// Takes the [`TaskManager`] handle out of this runtime, if one is stored.
374    ///
375    /// The handle resolves with `Err(PanickedTaskError)` if a critical task panicked,
376    /// or `Ok(())` if shutdown was requested. If not taken, the background task still
377    /// runs and logs panics at `debug!` level.
378    pub fn take_task_manager_handle(&self) -> Option<JoinHandle<Result<(), PanickedTaskError>>> {
379        self.0.task_manager_handle.lock().unwrap().take()
380    }
381
382    /// Returns the tokio runtime [`Handle`].
383    pub fn handle(&self) -> &Handle {
384        &self.0.handle
385    }
386
387    /// Get the general-purpose rayon CPU thread pool.
388    #[cfg(feature = "rayon")]
389    pub fn cpu_pool(&self) -> &rayon::ThreadPool {
390        &self.0.cpu_pool
391    }
392
393    /// Get the RPC blocking task pool.
394    #[cfg(feature = "rayon")]
395    pub fn rpc_pool(&self) -> &BlockingTaskPool {
396        &self.0.rpc_pool
397    }
398
399    /// Get the storage I/O pool.
400    #[cfg(feature = "rayon")]
401    pub fn storage_pool(&self) -> &rayon::ThreadPool {
402        &self.0.storage_pool
403    }
404
405    /// Get a clone of the [`BlockingTaskGuard`].
406    #[cfg(feature = "rayon")]
407    pub fn blocking_guard(&self) -> BlockingTaskGuard {
408        self.0.blocking_guard.clone()
409    }
410
411    /// Get the proof storage worker pool.
412    #[cfg(feature = "rayon")]
413    pub fn proof_storage_worker_pool(&self) -> &WorkerPool {
414        &self.0.proof_storage_worker_pool
415    }
416
417    /// Get the proof account worker pool.
418    #[cfg(feature = "rayon")]
419    pub fn proof_account_worker_pool(&self) -> &WorkerPool {
420        &self.0.proof_account_worker_pool
421    }
422
423    /// Returns the proof storage worker thread count the operator configured explicitly, if any.
424    /// A configured count of zero is treated as unset.
425    ///
426    /// Callers that size the worker count per block must use this verbatim instead of scaling it.
427    #[cfg(feature = "rayon")]
428    pub fn proof_storage_worker_threads_override(&self) -> Option<usize> {
429        self.0.proof_storage_worker_threads_override
430    }
431
432    /// Returns the proof account worker thread count the operator configured explicitly, if any.
433    /// A configured count of zero is treated as unset.
434    ///
435    /// Callers that size the worker count per block must use this verbatim instead of scaling it.
436    #[cfg(feature = "rayon")]
437    pub fn proof_account_worker_threads_override(&self) -> Option<usize> {
438        self.0.proof_account_worker_threads_override
439    }
440
441    /// Get the prewarming pool.
442    #[cfg(feature = "rayon")]
443    pub fn prewarming_pool(&self) -> &WorkerPool {
444        &self.0.prewarming_pool
445    }
446
447    /// Get the BAL streaming pool.
448    #[cfg(feature = "rayon")]
449    pub fn bal_streaming_pool(&self) -> &WorkerPool {
450        &self.0.bal_streaming_pool
451    }
452
453    /// Get the state trie overlay worker pool.
454    #[cfg(feature = "rayon")]
455    pub fn state_trie_overlay_worker_pool(&self) -> Arc<WorkerPool> {
456        Arc::clone(&self.0.state_trie_overlay_worker_pool)
457    }
458}
459
460// ── Test helpers ──────────────────────────────────────────────────────
461
462impl Runtime {
463    /// Creates a lightweight [`Runtime`] for tests with minimal thread pools.
464    ///
465    /// If called from within a tokio runtime (e.g. `#[tokio::test]`), attaches to the existing
466    /// handle to avoid shutdown panics when the test runtime is dropped.
467    pub fn test() -> Self {
468        let config = match Handle::try_current() {
469            Ok(handle) => Self::test_config().with_tokio(TokioConfig::existing_handle(handle)),
470            Err(_) => Self::test_config(),
471        };
472        RuntimeBuilder::new(config).build().expect("failed to build test Runtime")
473    }
474
475    const fn test_config() -> RuntimeConfig {
476        RuntimeConfig {
477            tokio: TokioConfig::Owned {
478                worker_threads: Some(2),
479                thread_keep_alive: DEFAULT_THREAD_KEEP_ALIVE,
480                thread_name: "tokio-test",
481            },
482            #[cfg(feature = "rayon")]
483            rayon: RayonConfig {
484                cpu_threads: Some(2),
485                reserved_cpu_cores: 0,
486                rpc_threads: Some(2),
487                storage_threads: Some(2),
488                max_blocking_tasks: 16,
489                proof_storage_worker_threads: Some(2),
490                proof_account_worker_threads: Some(2),
491                prewarming_threads: Some(2),
492                bal_streaming_threads: Some(2),
493                state_trie_overlay_worker_threads: Some(2),
494            },
495        }
496    }
497}
498
499// ── Spawn methods ─────────────────────────────────────────────────────
500
501/// Determines how a task is spawned.
502enum TaskKind {
503    /// Spawn the task to the default executor [`Handle::spawn`].
504    Default,
505    /// Spawn the task to the blocking executor [`Handle::spawn_blocking`].
506    Blocking,
507}
508
509impl Runtime {
510    /// Returns the receiver of the shutdown signal.
511    pub fn on_shutdown_signal(&self) -> &Shutdown {
512        &self.0.on_shutdown
513    }
514
515    /// Spawns a future on the tokio runtime depending on the [`TaskKind`].
516    fn spawn_on_rt<F>(&self, fut: F, task_kind: TaskKind) -> JoinHandle<()>
517    where
518        F: Future<Output = ()> + Send + 'static,
519    {
520        match task_kind {
521            TaskKind::Default => self.0.handle.spawn(fut),
522            TaskKind::Blocking => {
523                let handle = self.0.handle.clone();
524                self.0.handle.spawn_blocking(move || handle.block_on(fut))
525            }
526        }
527    }
528
529    /// Spawns a regular task depending on the given [`TaskKind`].
530    fn spawn_task_as<F>(&self, fut: F, task_kind: TaskKind) -> JoinHandle<()>
531    where
532        F: Future<Output = ()> + Send + 'static,
533    {
534        match task_kind {
535            TaskKind::Default => self.0.metrics.inc_regular_tasks(),
536            TaskKind::Blocking => self.0.metrics.inc_regular_blocking_tasks(),
537        }
538        let on_shutdown = self.0.on_shutdown.clone();
539
540        let finished_counter = match task_kind {
541            TaskKind::Default => self.0.metrics.finished_regular_tasks_total.clone(),
542            TaskKind::Blocking => self.0.metrics.finished_regular_blocking_tasks_total.clone(),
543        };
544
545        let task = {
546            async move {
547                let _inc_counter_on_drop = IncCounterOnDrop::new(finished_counter);
548                let fut = pin!(fut);
549                let _ = select(on_shutdown, fut).await;
550            }
551        }
552        .in_current_span();
553
554        self.spawn_on_rt(task, task_kind)
555    }
556
557    /// Spawns the task onto the runtime.
558    /// The given future resolves as soon as the [Shutdown] signal is received.
559    ///
560    /// See also [`Handle::spawn`].
561    pub fn spawn_task<F>(&self, fut: F) -> JoinHandle<()>
562    where
563        F: Future<Output = ()> + Send + 'static,
564    {
565        self.spawn_task_as(fut, TaskKind::Default)
566    }
567
568    /// Spawns a blocking task onto the runtime.
569    /// The given future resolves as soon as the [Shutdown] signal is received.
570    ///
571    /// See also [`Handle::spawn_blocking`].
572    pub fn spawn_blocking_task<F>(&self, fut: F) -> JoinHandle<()>
573    where
574        F: Future<Output = ()> + Send + 'static,
575    {
576        self.spawn_task_as(fut, TaskKind::Blocking)
577    }
578
579    /// Spawns a blocking closure directly on the tokio runtime, bypassing shutdown
580    /// awareness. Useful for raw CPU-bound work.
581    pub fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R>
582    where
583        F: FnOnce() -> R + Send + 'static,
584        R: Send + 'static,
585    {
586        self.0.handle.spawn_blocking(func)
587    }
588
589    /// Moves the given value to a dedicated background thread for deallocation.
590    ///
591    /// This is useful when dropping a value is expensive (e.g. large nested collections)
592    /// and should not block the current task. Uses a persistent named thread (`"drop"`)
593    /// to avoid thread creation overhead on hot paths.
594    pub fn spawn_drop<T: Send + 'static>(&self, value: T) {
595        self.spawn_blocking_named("drop", move || drop(value));
596    }
597
598    /// Spawns a blocking closure on a dedicated, named OS thread.
599    ///
600    /// Unlike [`spawn_blocking`](Self::spawn_blocking) which uses tokio's blocking thread pool,
601    /// this reuses the same OS thread for all tasks submitted under the same `name`. The thread
602    /// is created lazily on first use and its OS thread name is set to `name`.
603    ///
604    /// This is useful for tasks that benefit from running on a stable thread, e.g. for
605    /// thread-local state reuse or to avoid thread creation overhead on hot paths.
606    ///
607    /// Returns a [`LazyHandle`](crate::LazyHandle) handle that resolves on first access and caches
608    /// the result.
609    pub fn spawn_blocking_named<F, R>(&self, name: &'static str, func: F) -> crate::LazyHandle<R>
610    where
611        F: FnOnce() -> R + Send + 'static,
612        R: Send + 'static,
613    {
614        crate::LazyHandle::new(self.0.worker_map.spawn_on(name, func))
615    }
616
617    /// Attempts to spawn a blocking closure on a dedicated, named OS thread.
618    ///
619    /// Returns `None` if the named worker already has a task running or queued, allowing the caller
620    /// to fall back to another executor instead of serializing behind the named worker.
621    pub fn try_spawn_blocking_named<F, R>(
622        &self,
623        name: &'static str,
624        func: F,
625    ) -> Option<crate::LazyHandle<R>>
626    where
627        F: FnOnce() -> R + Send + 'static,
628        R: Send + 'static,
629    {
630        self.0.worker_map.try_spawn_on(name, func).map(crate::LazyHandle::new)
631    }
632
633    /// Spawns a blocking closure on a named OS thread if it is idle, otherwise falls back to
634    /// tokio's unnamed blocking thread pool.
635    ///
636    /// Returns `true` if the closure was spawned on the named thread.
637    pub fn spawn_blocking_named_or_tokio<F>(&self, name: &'static str, func: F) -> bool
638    where
639        F: FnOnce() + Send + 'static,
640    {
641        let func = Arc::new(parking_lot::Mutex::new(Some(func)));
642        let named_func = func.clone();
643
644        if self
645            .try_spawn_blocking_named(name, move || {
646                if let Some(func) = named_func.lock().take() {
647                    func();
648                }
649            })
650            .is_some()
651        {
652            return true
653        }
654
655        if let Some(func) = func.lock().take() {
656            self.spawn_blocking(func);
657        }
658        false
659    }
660
661    /// Spawns the task onto the runtime.
662    /// The given future resolves as soon as the [Shutdown] signal is received.
663    ///
664    /// See also [`Handle::spawn`].
665    pub fn spawn_with_signal<F>(&self, f: impl FnOnce(Shutdown) -> F) -> JoinHandle<()>
666    where
667        F: Future<Output = ()> + Send + 'static,
668    {
669        let on_shutdown = self.0.on_shutdown.clone();
670        let fut = f(on_shutdown);
671        let task = fut.in_current_span();
672        self.0.handle.spawn(task)
673    }
674
675    /// Spawns a critical task depending on the given [`TaskKind`].
676    fn spawn_critical_as<F>(
677        &self,
678        name: &'static str,
679        fut: F,
680        task_kind: TaskKind,
681    ) -> JoinHandle<()>
682    where
683        F: Future<Output = ()> + Send + 'static,
684    {
685        self.0.metrics.inc_critical_tasks();
686        let panicked_tasks_tx = self.0.task_events_tx.clone();
687        let on_shutdown = self.0.on_shutdown.clone();
688
689        // wrap the task in catch unwind
690        let task = std::panic::AssertUnwindSafe(fut)
691            .catch_unwind()
692            .map_err(move |error| {
693                let task_error = PanickedTaskError::new(name, error);
694                error!("{task_error}");
695                let _ = panicked_tasks_tx.send(TaskEvent::Panic(task_error));
696            })
697            .in_current_span();
698
699        let finished_critical_tasks_total_metrics =
700            self.0.metrics.finished_critical_tasks_total.clone();
701        let task = async move {
702            let _inc_counter_on_drop = IncCounterOnDrop::new(finished_critical_tasks_total_metrics);
703            let task = pin!(task);
704            let _ = select(on_shutdown, task).await;
705        };
706
707        self.spawn_on_rt(task, task_kind)
708    }
709
710    /// This spawns a critical task onto the runtime.
711    /// The given future resolves as soon as the [Shutdown] signal is received.
712    ///
713    /// If this task panics, the [`TaskManager`] is notified.
714    pub fn spawn_critical_task<F>(&self, name: &'static str, fut: F) -> JoinHandle<()>
715    where
716        F: Future<Output = ()> + Send + 'static,
717    {
718        self.spawn_critical_as(name, fut, TaskKind::Default)
719    }
720
721    /// This spawns a critical blocking task onto the runtime.
722    /// The given future resolves as soon as the [Shutdown] signal is received.
723    ///
724    /// If this task panics, the [`TaskManager`] is notified.
725    pub fn spawn_critical_blocking_task<F>(&self, name: &'static str, fut: F) -> JoinHandle<()>
726    where
727        F: Future<Output = ()> + Send + 'static,
728    {
729        self.spawn_critical_as(name, fut, TaskKind::Blocking)
730    }
731
732    /// This spawns a critical task onto a dedicated named OS thread.
733    /// The given future resolves as soon as the [`Shutdown`] signal is received.
734    ///
735    /// If this task panics, the [`TaskManager`] is notified.
736    pub fn spawn_critical_os_thread<F>(
737        &self,
738        thread_name: &'static str,
739        task_name: &'static str,
740        fut: F,
741    ) -> thread::JoinHandle<()>
742    where
743        F: Future<Output = ()> + Send + 'static,
744    {
745        self.0.metrics.inc_critical_tasks();
746        let handle = self.0.handle.clone();
747        let panicked_tasks_tx = self.0.task_events_tx.clone();
748        let on_shutdown = self.0.on_shutdown.clone();
749
750        let task = std::panic::AssertUnwindSafe(fut)
751            .catch_unwind()
752            .map_err(move |error| {
753                let task_error = PanickedTaskError::new(task_name, error);
754                error!("{task_error}");
755                let _ = panicked_tasks_tx.send(TaskEvent::Panic(task_error));
756            })
757            .in_current_span();
758
759        let finished_critical_tasks_total_metrics =
760            self.0.metrics.finished_critical_tasks_total.clone();
761        let task = async move {
762            let _inc_counter_on_drop = IncCounterOnDrop::new(finished_critical_tasks_total_metrics);
763            let task = pin!(task);
764            let _ = select(on_shutdown, task).await;
765        };
766
767        thread::Builder::new()
768            .name(thread_name.to_string())
769            .spawn(move || {
770                let _guard = handle.enter();
771                handle.block_on(task);
772            })
773            .unwrap_or_else(|e| panic!("failed to spawn critical OS thread {thread_name:?}: {e}"))
774    }
775
776    /// This spawns a critical task onto the runtime.
777    ///
778    /// If this task panics, the [`TaskManager`] is notified.
779    /// The [`TaskManager`] will wait until the given future has completed before shutting down.
780    ///
781    /// # Example
782    ///
783    /// ```no_run
784    /// # async fn t(executor: reth_tasks::TaskExecutor) {
785    ///
786    /// executor.spawn_critical_with_graceful_shutdown_signal("grace", async move |shutdown| {
787    ///     // await the shutdown signal
788    ///     let guard = shutdown.await;
789    ///     // do work before exiting the program
790    ///     tokio::time::sleep(std::time::Duration::from_secs(1)).await;
791    ///     // allow graceful shutdown
792    ///     drop(guard);
793    /// });
794    /// # }
795    /// ```
796    pub fn spawn_critical_with_graceful_shutdown_signal<F>(
797        &self,
798        name: &'static str,
799        f: impl FnOnce(GracefulShutdown) -> F,
800    ) -> JoinHandle<()>
801    where
802        F: Future<Output = ()> + Send + 'static,
803    {
804        let panicked_tasks_tx = self.0.task_events_tx.clone();
805        let on_shutdown = GracefulShutdown::new(
806            self.0.on_shutdown.clone(),
807            GracefulShutdownGuard::new(Arc::clone(&self.0.graceful_tasks)),
808        );
809        let fut = f(on_shutdown);
810
811        // wrap the task in catch unwind
812        let task = std::panic::AssertUnwindSafe(fut)
813            .catch_unwind()
814            .map_err(move |error| {
815                let task_error = PanickedTaskError::new(name, error);
816                error!("{task_error}");
817                let _ = panicked_tasks_tx.send(TaskEvent::Panic(task_error));
818            })
819            .map(drop)
820            .in_current_span();
821
822        self.0.handle.spawn(task)
823    }
824
825    /// This spawns a regular task onto the runtime.
826    ///
827    /// The [`TaskManager`] will wait until the given future has completed before shutting down.
828    ///
829    /// # Example
830    ///
831    /// ```no_run
832    /// # async fn t(executor: reth_tasks::TaskExecutor) {
833    ///
834    /// executor.spawn_with_graceful_shutdown_signal(async move |shutdown| {
835    ///     // await the shutdown signal
836    ///     let guard = shutdown.await;
837    ///     // do work before exiting the program
838    ///     tokio::time::sleep(std::time::Duration::from_secs(1)).await;
839    ///     // allow graceful shutdown
840    ///     drop(guard);
841    /// });
842    /// # }
843    /// ```
844    pub fn spawn_with_graceful_shutdown_signal<F>(
845        &self,
846        f: impl FnOnce(GracefulShutdown) -> F,
847    ) -> JoinHandle<()>
848    where
849        F: Future<Output = ()> + Send + 'static,
850    {
851        let on_shutdown = GracefulShutdown::new(
852            self.0.on_shutdown.clone(),
853            GracefulShutdownGuard::new(Arc::clone(&self.0.graceful_tasks)),
854        );
855        let fut = f(on_shutdown);
856
857        self.0.handle.spawn(fut)
858    }
859
860    /// Sends a request to the `TaskManager` to initiate a graceful shutdown.
861    ///
862    /// Caution: This will terminate the entire program.
863    pub fn initiate_graceful_shutdown(
864        &self,
865    ) -> Result<GracefulShutdown, tokio::sync::mpsc::error::SendError<()>> {
866        self.0
867            .task_events_tx
868            .send(TaskEvent::GracefulShutdown)
869            .map_err(|_send_error_with_task_event| tokio::sync::mpsc::error::SendError(()))?;
870
871        Ok(GracefulShutdown::new(
872            self.0.on_shutdown.clone(),
873            GracefulShutdownGuard::new(Arc::clone(&self.0.graceful_tasks)),
874        ))
875    }
876
877    /// Fires the shutdown signal and waits until all graceful tasks complete.
878    pub fn graceful_shutdown(&self) {
879        let _ = self.do_graceful_shutdown(None);
880    }
881
882    /// Fires the shutdown signal and waits until all graceful tasks complete or the timeout
883    /// elapses.
884    ///
885    /// Returns `true` if all tasks completed before the timeout.
886    pub fn graceful_shutdown_with_timeout(&self, timeout: Duration) -> bool {
887        self.do_graceful_shutdown(Some(timeout))
888    }
889
890    /// Shuts down the owned tokio runtime for all clones, waiting at most `timeout` for blocking
891    /// tasks. Externally supplied runtimes are left running.
892    ///
893    /// Panics if called from an async context while an owned runtime remains.
894    pub fn shutdown_timeout(self, timeout: Duration) {
895        let runtime = self.0.tokio_runtime.lock().unwrap().take();
896        if let Some(runtime) = runtime {
897            runtime.shutdown_timeout(timeout);
898        }
899    }
900
901    fn do_graceful_shutdown(&self, timeout: Option<Duration>) -> bool {
902        let _ = self.0.task_events_tx.send(TaskEvent::GracefulShutdown);
903        let deadline = timeout.map(|t| Instant::now() + t);
904        while self.0.graceful_tasks.load(Ordering::SeqCst) > 0 {
905            if deadline.is_some_and(|d| Instant::now() > d) {
906                debug!("graceful shutdown timed out");
907                return false;
908            }
909            std::thread::yield_now();
910        }
911        debug!("gracefully shut down");
912        true
913    }
914}
915
916// ── RuntimeBuilder ────────────────────────────────────────────────────
917
918/// Builder for constructing a [`Runtime`].
919#[derive(Debug, Clone)]
920pub struct RuntimeBuilder {
921    config: RuntimeConfig,
922}
923
924impl RuntimeBuilder {
925    /// Create a new builder with the given configuration.
926    pub const fn new(config: RuntimeConfig) -> Self {
927        Self { config }
928    }
929
930    /// Build the [`Runtime`].
931    ///
932    /// The [`TaskManager`] is automatically spawned as a background task that monitors
933    /// critical tasks for panics. Use [`Runtime::take_task_manager_handle`] to extract
934    /// the join handle if you need to poll for panic errors.
935    #[tracing::instrument(name = "RuntimeBuilder::build", level = "debug", skip_all)]
936    pub fn build(self) -> Result<Runtime, RuntimeBuildError> {
937        debug!(?self.config, "Building runtime");
938        let config = self.config;
939
940        let (owned_runtime, handle) = match &config.tokio {
941            TokioConfig::Owned { worker_threads, thread_keep_alive, thread_name } => {
942                let mut builder = tokio::runtime::Builder::new_multi_thread();
943                builder
944                    .enable_all()
945                    .thread_keep_alive(*thread_keep_alive)
946                    .thread_name(*thread_name);
947
948                if let Some(threads) = worker_threads {
949                    builder.worker_threads(*threads);
950                }
951
952                let runtime = builder.build()?;
953                let h = runtime.handle().clone();
954                (Some(runtime), h)
955            }
956            TokioConfig::ExistingHandle(h) => (None, h.clone()),
957        };
958
959        let (task_manager, on_shutdown, task_events_tx, graceful_tasks) =
960            TaskManager::new_parts(handle.clone());
961
962        #[cfg(feature = "rayon")]
963        let (
964            cpu_pool,
965            rpc_pool,
966            storage_pool,
967            blocking_guard,
968            proof_storage_worker_pool,
969            proof_account_worker_pool,
970            prewarming_pool,
971            bal_streaming_pool,
972            state_trie_overlay_worker_pool,
973        ) = {
974            let default_threads = config.rayon.default_thread_count();
975            let rpc_threads = config.rayon.rpc_threads.unwrap_or(default_threads);
976
977            let cpu_pool = build_pool_with_panic_handler(
978                rayon::ThreadPoolBuilder::new()
979                    .num_threads(default_threads)
980                    .thread_name(|i| format!("cpu-{i:02}")),
981            )?;
982
983            let rpc_raw = build_pool_with_panic_handler(
984                rayon::ThreadPoolBuilder::new()
985                    .num_threads(rpc_threads)
986                    .thread_name(|i| format!("rpc-{i:02}")),
987            )?;
988            let rpc_pool = BlockingTaskPool::new(rpc_raw);
989
990            let storage_threads =
991                config.rayon.storage_threads.unwrap_or(DEFAULT_STORAGE_POOL_THREADS);
992            let storage_pool = build_pool_with_panic_handler(
993                rayon::ThreadPoolBuilder::new()
994                    .num_threads(storage_threads)
995                    .thread_name(|i| format!("storage-{i:02}")),
996            )?;
997
998            let blocking_guard = BlockingTaskGuard::new(config.rayon.max_blocking_tasks);
999
1000            // `cpu_threads` may be zero, which rayon resolves to the automatic count when the
1001            // cpu pool is built; size the proof pools from that resolved count so a zero never
1002            // reaches the per-block worker budget.
1003            let default_proof_worker_threads = cpu_pool.current_num_threads() * 4;
1004
1005            let proof_storage_worker_threads = config
1006                .rayon
1007                .proof_storage_worker_threads_override()
1008                .unwrap_or(default_proof_worker_threads);
1009            let proof_storage_worker_pool =
1010                WorkerPool::new(proof_storage_worker_threads, "proof-strg");
1011
1012            let proof_account_worker_threads = config
1013                .rayon
1014                .proof_account_worker_threads_override()
1015                .unwrap_or(default_proof_worker_threads);
1016            let proof_account_worker_pool =
1017                WorkerPool::new(proof_account_worker_threads, "proof-acct");
1018
1019            let prewarming_threads = config.rayon.prewarming_threads.unwrap_or(default_threads);
1020            let prewarming_pool = WorkerPool::new(prewarming_threads, "prewarm");
1021
1022            let bal_streaming_threads =
1023                config.rayon.bal_streaming_threads.unwrap_or(default_threads);
1024            let bal_streaming_pool = WorkerPool::new(bal_streaming_threads, "bal-stream");
1025
1026            let state_trie_overlay_worker_threads = config
1027                .rayon
1028                .state_trie_overlay_worker_threads
1029                .unwrap_or(DEFAULT_STATE_TRIE_OVERLAY_WORKER_THREADS);
1030            let state_trie_overlay_worker_pool =
1031                Arc::new(WorkerPool::new(state_trie_overlay_worker_threads, "state-ovly"));
1032
1033            debug!(
1034                default_threads,
1035                rpc_threads,
1036                storage_threads,
1037                proof_storage_worker_threads,
1038                proof_account_worker_threads,
1039                prewarming_threads,
1040                bal_streaming_threads,
1041                state_trie_overlay_worker_threads,
1042                max_blocking_tasks = config.rayon.max_blocking_tasks,
1043                "Configured lazy rayon worker pools"
1044            );
1045
1046            (
1047                cpu_pool,
1048                rpc_pool,
1049                storage_pool,
1050                blocking_guard,
1051                proof_storage_worker_pool,
1052                proof_account_worker_pool,
1053                prewarming_pool,
1054                bal_streaming_pool,
1055                state_trie_overlay_worker_pool,
1056            )
1057        };
1058
1059        let task_manager_handle = handle.spawn(async move {
1060            let result = task_manager.await;
1061            if let Err(ref err) = result {
1062                debug!("{err}");
1063            }
1064            result
1065        });
1066
1067        let inner = RuntimeInner {
1068            tokio_runtime: Mutex::new(owned_runtime),
1069            handle,
1070            on_shutdown,
1071            task_events_tx,
1072            metrics: Default::default(),
1073            graceful_tasks,
1074            #[cfg(feature = "rayon")]
1075            cpu_pool,
1076            #[cfg(feature = "rayon")]
1077            rpc_pool,
1078            #[cfg(feature = "rayon")]
1079            storage_pool,
1080            #[cfg(feature = "rayon")]
1081            blocking_guard,
1082            #[cfg(feature = "rayon")]
1083            proof_storage_worker_pool,
1084            #[cfg(feature = "rayon")]
1085            proof_account_worker_pool,
1086            #[cfg(feature = "rayon")]
1087            proof_storage_worker_threads_override: config
1088                .rayon
1089                .proof_storage_worker_threads_override(),
1090            #[cfg(feature = "rayon")]
1091            proof_account_worker_threads_override: config
1092                .rayon
1093                .proof_account_worker_threads_override(),
1094            #[cfg(feature = "rayon")]
1095            prewarming_pool,
1096            #[cfg(feature = "rayon")]
1097            bal_streaming_pool,
1098            #[cfg(feature = "rayon")]
1099            state_trie_overlay_worker_pool,
1100            worker_map: WorkerMap::new(),
1101            task_manager_handle: Mutex::new(Some(task_manager_handle)),
1102        };
1103
1104        Ok(Runtime(Arc::new(inner)))
1105    }
1106}
1107
1108#[cfg(test)]
1109mod tests {
1110    use super::*;
1111
1112    #[test]
1113    fn test_runtime_config_default() {
1114        let config = RuntimeConfig::default();
1115        assert!(matches!(config.tokio, TokioConfig::Owned { .. }));
1116    }
1117
1118    #[test]
1119    fn cpu_cores_override_sizes_tokio_workers() {
1120        let config = RuntimeConfig::default().with_cpu_cores(NonZeroUsize::new(8).unwrap());
1121        assert!(matches!(config.tokio, TokioConfig::Owned { worker_threads: Some(8), .. }));
1122    }
1123
1124    #[test]
1125    fn cpu_cores_override_preserves_explicit_tokio_workers() {
1126        let config = RuntimeConfig::default()
1127            .with_tokio(TokioConfig::with_worker_threads(2))
1128            .with_cpu_cores(NonZeroUsize::new(8).unwrap());
1129        assert!(matches!(config.tokio, TokioConfig::Owned { worker_threads: Some(2), .. }));
1130    }
1131
1132    #[test]
1133    fn cpu_cores_override_preserves_existing_tokio_runtime() {
1134        let runtime = TokioRuntime::new().unwrap();
1135        let config = RuntimeConfig::default()
1136            .with_tokio(TokioConfig::existing_handle(runtime.handle().clone()))
1137            .with_cpu_cores(NonZeroUsize::new(8).unwrap());
1138        assert!(matches!(config.tokio, TokioConfig::ExistingHandle(_)));
1139    }
1140
1141    #[cfg(feature = "rayon")]
1142    #[test]
1143    fn cpu_cores_override_sizes_cpu_pools() {
1144        let runtime = RuntimeBuilder::new(
1145            RuntimeConfig::default().with_cpu_cores(NonZeroUsize::new(5).unwrap()),
1146        )
1147        .build()
1148        .unwrap();
1149        assert_eq!(runtime.handle().metrics().num_workers(), 5);
1150        assert_eq!(runtime.cpu_pool().current_num_threads(), 5);
1151        assert_eq!(
1152            runtime
1153                .handle()
1154                .block_on(runtime.rpc_pool().spawn(rayon::current_num_threads))
1155                .unwrap(),
1156            5
1157        );
1158        assert_eq!(runtime.prewarming_pool().current_num_threads(), 5);
1159        assert_eq!(runtime.bal_streaming_pool().current_num_threads(), 5);
1160        assert_eq!(runtime.proof_storage_worker_pool().num_threads(), 20);
1161        assert_eq!(runtime.proof_account_worker_pool().num_threads(), 20);
1162        assert_eq!(runtime.storage_pool().current_num_threads(), DEFAULT_STORAGE_POOL_THREADS);
1163        assert_eq!(
1164            runtime.state_trie_overlay_worker_pool().num_threads(),
1165            DEFAULT_STATE_TRIE_OVERLAY_WORKER_THREADS
1166        );
1167    }
1168
1169    #[cfg(feature = "rayon")]
1170    #[test]
1171    fn cpu_cores_override_preserves_explicit_cpu_threads() {
1172        let config = RuntimeConfig::default()
1173            .with_rayon(RayonConfig { cpu_threads: Some(2), ..Default::default() })
1174            .with_cpu_cores(NonZeroUsize::new(8).unwrap());
1175        assert_eq!(config.rayon.default_thread_count(), 2);
1176    }
1177
1178    #[test]
1179    fn test_runtime_config_existing_handle() {
1180        let rt = TokioRuntime::new().unwrap();
1181        let config =
1182            Runtime::test_config().with_tokio(TokioConfig::existing_handle(rt.handle().clone()));
1183        assert!(matches!(config.tokio, TokioConfig::ExistingHandle(_)));
1184    }
1185
1186    #[cfg(feature = "rayon")]
1187    #[test]
1188    fn test_rayon_config_thread_count() {
1189        let config = RayonConfig::default();
1190        let count = config.default_thread_count();
1191        assert!(count >= 1);
1192    }
1193
1194    #[cfg(feature = "rayon")]
1195    #[test]
1196    fn zero_cpu_threads_still_size_the_proof_pools() {
1197        let rt = TokioRuntime::new().unwrap();
1198        let mut config =
1199            Runtime::test_config().with_tokio(TokioConfig::existing_handle(rt.handle().clone()));
1200        config.rayon.cpu_threads = Some(0);
1201        config.rayon.proof_storage_worker_threads = None;
1202        config.rayon.proof_account_worker_threads = None;
1203        let runtime = RuntimeBuilder::new(config).build().unwrap();
1204
1205        let resolved = runtime.cpu_pool().current_num_threads();
1206        assert!(resolved >= 1);
1207        assert_eq!(runtime.proof_storage_worker_pool().current_num_threads(), resolved * 4);
1208        assert_eq!(runtime.proof_account_worker_pool().current_num_threads(), resolved * 4);
1209    }
1210
1211    #[cfg(feature = "rayon")]
1212    #[test]
1213    fn proof_worker_threads_keep_pinned_counts() {
1214        let rt = TokioRuntime::new().unwrap();
1215        for configured in [1, 8, 16, 17, 64] {
1216            let mut config = Runtime::test_config()
1217                .with_tokio(TokioConfig::existing_handle(rt.handle().clone()));
1218            config.rayon.proof_storage_worker_threads = Some(configured);
1219            config.rayon.proof_account_worker_threads = Some(configured);
1220            let runtime = RuntimeBuilder::new(config).build().unwrap();
1221            assert_eq!(runtime.proof_storage_worker_pool().current_num_threads(), configured);
1222            assert_eq!(runtime.proof_account_worker_pool().current_num_threads(), configured);
1223            assert_eq!(runtime.proof_storage_worker_threads_override(), Some(configured));
1224            assert_eq!(runtime.proof_account_worker_threads_override(), Some(configured));
1225        }
1226    }
1227
1228    #[cfg(feature = "rayon")]
1229    #[test]
1230    fn zero_proof_worker_threads_use_the_default() {
1231        let rt = TokioRuntime::new().unwrap();
1232        let mut config =
1233            Runtime::test_config().with_tokio(TokioConfig::existing_handle(rt.handle().clone()));
1234        config.rayon.proof_storage_worker_threads = Some(0);
1235        config.rayon.proof_account_worker_threads = Some(0);
1236        let runtime = RuntimeBuilder::new(config).build().unwrap();
1237
1238        let default = runtime.cpu_pool().current_num_threads() * 4;
1239        assert_eq!(runtime.proof_storage_worker_pool().current_num_threads(), default);
1240        assert_eq!(runtime.proof_account_worker_pool().current_num_threads(), default);
1241        assert_eq!(runtime.proof_storage_worker_threads_override(), None);
1242        assert_eq!(runtime.proof_account_worker_threads_override(), None);
1243    }
1244
1245    #[test]
1246    fn test_runtime_builder() {
1247        let rt = TokioRuntime::new().unwrap();
1248        let config =
1249            Runtime::test_config().with_tokio(TokioConfig::existing_handle(rt.handle().clone()));
1250        let runtime = RuntimeBuilder::new(config).build().unwrap();
1251        let _ = runtime.handle();
1252    }
1253
1254    #[tokio::test]
1255    async fn last_runtime_clone_can_drop_on_its_worker() {
1256        let runtime = RuntimeBuilder::new(Runtime::test_config()).build().unwrap();
1257        let task_runtime = runtime.clone();
1258        let (release, wait) = tokio::sync::oneshot::channel();
1259        let task = runtime.handle().spawn(async move {
1260            wait.await.unwrap();
1261            drop(task_runtime);
1262        });
1263
1264        thread::spawn(move || drop(runtime)).join().unwrap();
1265        release.send(()).unwrap();
1266        task.await.unwrap();
1267    }
1268
1269    #[tokio::test]
1270    async fn shutdown_with_live_runtime_clones() {
1271        let runtime = RuntimeBuilder::new(Runtime::test_config()).build().unwrap();
1272        let retained = runtime.clone();
1273        let task_runtime = runtime.clone();
1274        let (started, ready) = tokio::sync::oneshot::channel();
1275        let task = runtime.handle().spawn(async move {
1276            let _runtime = task_runtime;
1277            started.send(()).unwrap();
1278            std::future::pending::<()>().await;
1279        });
1280        ready.await.unwrap();
1281
1282        thread::spawn(move || runtime.shutdown_timeout(Duration::from_secs(5))).join().unwrap();
1283
1284        assert!(task.await.unwrap_err().is_cancelled());
1285        assert!(retained.handle().spawn(async {}).await.unwrap_err().is_cancelled());
1286        retained.shutdown_timeout(Duration::from_secs(5));
1287    }
1288
1289    #[tokio::test]
1290    async fn shutdown_preserves_external_runtime() {
1291        let runtime = RuntimeBuilder::new(
1292            Runtime::test_config().with_tokio(TokioConfig::existing_handle(Handle::current())),
1293        )
1294        .build()
1295        .unwrap();
1296        let retained = runtime.clone();
1297
1298        runtime.shutdown_timeout(Duration::from_secs(5));
1299
1300        assert_eq!(retained.handle().spawn(async { 42 }).await.unwrap(), 42);
1301    }
1302
1303    #[test]
1304    fn critical_os_thread_uses_requested_name() {
1305        let runtime = Runtime::test();
1306        let (tx, rx) = std::sync::mpsc::channel();
1307
1308        let handle = runtime.spawn_critical_os_thread(
1309            "critical-os-test",
1310            "critical os thread test",
1311            async move {
1312                let name = thread::current().name().unwrap().to_string();
1313                tx.send(name).unwrap();
1314            },
1315        );
1316
1317        let name = rx.recv_timeout(Duration::from_secs(5)).unwrap();
1318        assert_eq!(name, "critical-os-test");
1319        handle.join().unwrap();
1320    }
1321
1322    #[test]
1323    fn critical_os_thread_panic_is_reported() {
1324        let runtime = Runtime::test();
1325        let manager_handle = runtime.take_task_manager_handle().unwrap();
1326
1327        let handle = runtime.spawn_critical_os_thread(
1328            "critical-os-panic",
1329            "critical os thread panic test",
1330            async { panic!("critical os thread panic") },
1331        );
1332
1333        let err =
1334            runtime.handle().block_on(async move { manager_handle.await.unwrap().unwrap_err() });
1335        assert_eq!(err.task_name, "critical os thread panic test");
1336        assert_eq!(err.error, Some("critical os thread panic".to_string()));
1337        handle.join().unwrap();
1338    }
1339
1340    #[cfg(feature = "rayon")]
1341    #[test]
1342    fn test_worker_pools_are_lazy() {
1343        let runtime = Runtime::test();
1344
1345        // Worker pools are lazy — not initialized until first access.
1346        assert!(!runtime.0.bal_streaming_pool.is_initialized());
1347        assert!(!runtime.0.proof_storage_worker_pool.is_initialized());
1348        assert!(!runtime.0.state_trie_overlay_worker_pool.is_initialized());
1349
1350        // Accessing them triggers initialization and returns the configured thread count.
1351        assert_eq!(runtime.bal_streaming_pool().current_num_threads(), 2);
1352        assert!(runtime.0.bal_streaming_pool.is_initialized());
1353
1354        assert_eq!(runtime.proof_storage_worker_pool().current_num_threads(), 2);
1355        assert_eq!(runtime.proof_account_worker_pool().current_num_threads(), 2);
1356        assert_eq!(runtime.prewarming_pool().current_num_threads(), 2);
1357        assert_eq!(runtime.state_trie_overlay_worker_pool().current_num_threads(), 2);
1358    }
1359}