reifydb_runtime/pool/
native.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4#![allow(clippy::disallowed_types)]
5
6use std::{future::Future, sync::Arc, time::Duration};
7
8use reifydb_value::reifydb_assertions;
9use tokio::{
10	runtime::{self, Handle, Runtime},
11	task::JoinHandle,
12};
13
14use super::PoolConfig;
15use crate::{
16	pool::{
17		actor_pool::{ActorPool, Schedule},
18		compute::ComputePool,
19		task::TaskPool,
20	},
21	sync::mutex::Mutex,
22};
23
24struct PoolsInner {
25	actors: ActorPool,
26	task: TaskPool,
27	compute: ComputePool,
28	tokio_handle: Option<Handle>,
29	tokio: Mutex<Option<Runtime>>,
30}
31
32impl PoolsInner {
33	fn take_tokio(&self) -> Option<Runtime> {
34		self.tokio.lock().take()
35	}
36}
37
38impl Drop for PoolsInner {
39	fn drop(&mut self) {
40		if let Some(rt) = self.take_tokio() {
41			if runtime::Handle::try_current().is_err() {
42				rt.shutdown_timeout(Duration::from_secs(5));
43			} else {
44				rt.shutdown_background();
45			}
46		}
47		self.actors.shutdown();
48		self.task.shutdown();
49	}
50}
51
52#[derive(Clone)]
53pub struct Pools {
54	inner: Arc<PoolsInner>,
55}
56
57impl Default for Pools {
58	fn default() -> Self {
59		Self::new(PoolConfig::default())
60	}
61}
62
63impl Pools {
64	pub fn new(config: PoolConfig) -> Self {
65		let actors = ActorPool::new(config.coordination_threads, config.flow_threads);
66		let task = TaskPool::new(config.task_threads, "task");
67		let compute = ComputePool::new(config.compute_threads, "compute");
68		let (tokio_handle, tokio) = Self::build_async_runtime(config.async_threads);
69
70		Self {
71			inner: Arc::new(PoolsInner {
72				actors,
73				task,
74				compute,
75				tokio_handle,
76				tokio,
77			}),
78		}
79	}
80
81	#[inline]
82	fn build_async_runtime(threads: usize) -> (Option<Handle>, Mutex<Option<Runtime>>) {
83		let (tokio_handle, tokio) = if threads > 0 {
84			let rt = runtime::Builder::new_multi_thread()
85				.worker_threads(threads)
86				.thread_name("async")
87				.enable_all()
88				.build()
89				.expect("failed to build tokio runtime");
90			let handle = rt.handle().clone();
91			(Some(handle), Mutex::new(Some(rt)))
92		} else {
93			(None, Mutex::new(None))
94		};
95
96		reifydb_assertions! {
97			let handle_present = tokio_handle.is_some();
98			let runtime_present = tokio.lock().is_some();
99			assert!(
100				handle_present == runtime_present,
101				"async handle/runtime presence must agree (handle_present={handle_present}, runtime_present={runtime_present}); \
102				 a handle without its runtime makes tokio_handle() dispatch onto a runtime that Drop never shuts down (thread leak), \
103				 and a runtime without a handle makes spawn()/block_on() panic via the expect in tokio_handle()"
104			);
105		}
106
107		(tokio_handle, tokio)
108	}
109
110	pub fn shutdown(&self) {
111		if let Some(rt) = self.inner.take_tokio() {
112			if runtime::Handle::try_current().is_err() {
113				rt.shutdown_timeout(Duration::from_secs(5));
114			} else {
115				rt.shutdown_background();
116			}
117		}
118		self.inner.actors.shutdown();
119		self.inner.task.shutdown();
120	}
121
122	pub fn spawn_task(&self, job: impl FnOnce() + Send + 'static) {
123		self.inner.task.spawn(job);
124	}
125
126	pub fn task_thread_count(&self) -> usize {
127		self.inner.task.thread_count()
128	}
129
130	pub fn compute(&self) -> &ComputePool {
131		&self.inner.compute
132	}
133
134	pub fn compute_thread_count(&self) -> usize {
135		self.inner.compute.thread_count()
136	}
137
138	pub fn coordination_thread_count(&self) -> usize {
139		self.inner.actors.coordination().thread_count()
140	}
141
142	pub fn flow_thread_count(&self) -> usize {
143		self.inner.actors.flow().thread_count()
144	}
145
146	pub(crate) fn actor_pool(&self) -> &ActorPool {
147		&self.inner.actors
148	}
149
150	pub(crate) fn task_injector(&self) -> Schedule {
151		Schedule::Injector(self.inner.task.injector())
152	}
153
154	fn tokio_handle(&self) -> Handle {
155		self.inner.tokio_handle.clone().expect("no tokio runtime configured (async_threads = 0)")
156	}
157
158	pub fn handle(&self) -> Handle {
159		self.tokio_handle()
160	}
161
162	pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
163	where
164		F: Future + Send + 'static,
165		F::Output: Send + 'static,
166	{
167		self.tokio_handle().spawn(future)
168	}
169
170	pub fn block_on<F>(&self, future: F) -> F::Output
171	where
172		F: Future,
173	{
174		self.tokio_handle().block_on(future)
175	}
176}
reifydb_runtime/pool/native.rs

reifydb_runtime/pool/
native.rs
