1#[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
37pub const DEFAULT_THREAD_KEEP_ALIVE: Duration = Duration::from_secs(15);
39
40pub const DEFAULT_RESERVED_CPU_CORES: usize = 2;
42
43pub const DEFAULT_STORAGE_POOL_THREADS: usize = 16;
45
46#[cfg(feature = "rayon")]
48pub const DEFAULT_STATE_TRIE_OVERLAY_WORKER_THREADS: usize = 4;
49
50pub const DEFAULT_MAX_BLOCKING_TASKS: usize = 512;
52
53#[derive(Debug, Clone)]
55pub enum TokioConfig {
56 Owned {
58 worker_threads: Option<usize>,
60 thread_keep_alive: Duration,
62 thread_name: &'static str,
64 },
65 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 pub const fn existing_handle(handle: Handle) -> Self {
82 Self::ExistingHandle(handle)
83 }
84
85 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#[derive(Debug, Clone)]
97#[cfg(feature = "rayon")]
98pub struct RayonConfig {
99 pub cpu_threads: Option<usize>,
102 pub reserved_cpu_cores: usize,
104 pub rpc_threads: Option<usize>,
107 pub storage_threads: Option<usize>,
110 pub max_blocking_tasks: usize,
112 pub proof_storage_worker_threads: Option<usize>,
115 pub proof_account_worker_threads: Option<usize>,
118 pub prewarming_threads: Option<usize>,
121 pub bal_streaming_threads: Option<usize>,
124 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 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 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 pub const fn with_rpc_threads(mut self, rpc_threads: usize) -> Self {
163 self.rpc_threads = Some(rpc_threads);
164 self
165 }
166
167 pub const fn with_storage_threads(mut self, storage_threads: usize) -> Self {
169 self.storage_threads = Some(storage_threads);
170 self
171 }
172
173 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 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 pub const fn with_prewarming_threads(mut self, prewarming_threads: usize) -> Self {
193 self.prewarming_threads = Some(prewarming_threads);
194 self
195 }
196
197 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 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 fn default_thread_count(&self) -> usize {
214 let _ = self.reserved_cpu_cores;
217 self.cpu_threads.unwrap_or_else(|| available_parallelism().map_or(1, NonZeroUsize::get))
218 }
219
220 fn proof_storage_worker_threads_override(&self) -> Option<usize> {
223 self.proof_storage_worker_threads.filter(|&threads| threads > 0)
224 }
225
226 fn proof_account_worker_threads_override(&self) -> Option<usize> {
229 self.proof_account_worker_threads.filter(|&threads| threads > 0)
230 }
231}
232
233#[derive(Debug, Clone, Default)]
235pub struct RuntimeConfig {
236 pub tokio: TokioConfig,
238 #[cfg(feature = "rayon")]
240 pub rayon: RayonConfig,
241}
242
243impl RuntimeConfig {
244 pub fn with_tokio(mut self, tokio: TokioConfig) -> Self {
246 self.tokio = tokio;
247 self
248 }
249
250 #[cfg(feature = "rayon")]
252 pub const fn with_rayon(mut self, rayon: RayonConfig) -> Self {
253 self.rayon = rayon;
254 self
255 }
256
257 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#[derive(Debug, thiserror::Error)]
274pub enum RuntimeBuildError {
275 #[error("Failed to build tokio runtime: {0}")]
277 TokioBuild(#[from] std::io::Error),
278 #[cfg(feature = "rayon")]
280 #[error("Failed to build rayon thread pool: {0}")]
281 RayonBuild(#[from] rayon::ThreadPoolBuildError),
282}
283
284struct RuntimeInner {
287 tokio_runtime: Mutex<Option<TokioRuntime>>,
289 handle: Handle,
291 on_shutdown: Shutdown,
293 task_events_tx: UnboundedSender<TaskEvent>,
295 metrics: TaskExecutorMetrics,
297 graceful_tasks: Arc<AtomicUsize>,
299 #[cfg(feature = "rayon")]
301 cpu_pool: rayon::ThreadPool,
302 #[cfg(feature = "rayon")]
304 rpc_pool: BlockingTaskPool,
305 #[cfg(feature = "rayon")]
307 storage_pool: rayon::ThreadPool,
308 #[cfg(feature = "rayon")]
310 blocking_guard: BlockingTaskGuard,
311 #[cfg(feature = "rayon")]
313 proof_storage_worker_pool: WorkerPool,
314 #[cfg(feature = "rayon")]
316 proof_account_worker_pool: WorkerPool,
317 #[cfg(feature = "rayon")]
319 proof_storage_worker_threads_override: Option<usize>,
320 #[cfg(feature = "rayon")]
322 proof_account_worker_threads_override: Option<usize>,
323 #[cfg(feature = "rayon")]
325 prewarming_pool: WorkerPool,
326 #[cfg(feature = "rayon")]
328 bal_streaming_pool: WorkerPool,
329 #[cfg(feature = "rayon")]
331 state_trie_overlay_worker_pool: Arc<WorkerPool>,
332 worker_map: WorkerMap,
335 task_manager_handle: Mutex<Option<JoinHandle<Result<(), PanickedTaskError>>>>,
339}
340
341impl Drop for RuntimeInner {
342 fn drop(&mut self) {
343 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#[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
370impl Runtime {
373 pub fn take_task_manager_handle(&self) -> Option<JoinHandle<Result<(), PanickedTaskError>>> {
379 self.0.task_manager_handle.lock().unwrap().take()
380 }
381
382 pub fn handle(&self) -> &Handle {
384 &self.0.handle
385 }
386
387 #[cfg(feature = "rayon")]
389 pub fn cpu_pool(&self) -> &rayon::ThreadPool {
390 &self.0.cpu_pool
391 }
392
393 #[cfg(feature = "rayon")]
395 pub fn rpc_pool(&self) -> &BlockingTaskPool {
396 &self.0.rpc_pool
397 }
398
399 #[cfg(feature = "rayon")]
401 pub fn storage_pool(&self) -> &rayon::ThreadPool {
402 &self.0.storage_pool
403 }
404
405 #[cfg(feature = "rayon")]
407 pub fn blocking_guard(&self) -> BlockingTaskGuard {
408 self.0.blocking_guard.clone()
409 }
410
411 #[cfg(feature = "rayon")]
413 pub fn proof_storage_worker_pool(&self) -> &WorkerPool {
414 &self.0.proof_storage_worker_pool
415 }
416
417 #[cfg(feature = "rayon")]
419 pub fn proof_account_worker_pool(&self) -> &WorkerPool {
420 &self.0.proof_account_worker_pool
421 }
422
423 #[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 #[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 #[cfg(feature = "rayon")]
443 pub fn prewarming_pool(&self) -> &WorkerPool {
444 &self.0.prewarming_pool
445 }
446
447 #[cfg(feature = "rayon")]
449 pub fn bal_streaming_pool(&self) -> &WorkerPool {
450 &self.0.bal_streaming_pool
451 }
452
453 #[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
460impl Runtime {
463 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
499enum TaskKind {
503 Default,
505 Blocking,
507}
508
509impl Runtime {
510 pub fn on_shutdown_signal(&self) -> &Shutdown {
512 &self.0.on_shutdown
513 }
514
515 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 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 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 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 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 pub fn spawn_drop<T: Send + 'static>(&self, value: T) {
595 self.spawn_blocking_named("drop", move || drop(value));
596 }
597
598 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 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 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 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 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 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 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 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 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 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 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 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 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 pub fn graceful_shutdown(&self) {
879 let _ = self.do_graceful_shutdown(None);
880 }
881
882 pub fn graceful_shutdown_with_timeout(&self, timeout: Duration) -> bool {
887 self.do_graceful_shutdown(Some(timeout))
888 }
889
890 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#[derive(Debug, Clone)]
920pub struct RuntimeBuilder {
921 config: RuntimeConfig,
922}
923
924impl RuntimeBuilder {
925 pub const fn new(config: RuntimeConfig) -> Self {
927 Self { config }
928 }
929
930 #[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 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 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 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}