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 rayon::{ThreadPool, ThreadPoolBuilder};
9use reifydb_value::reifydb_assertions;
10use tokio::{
11	runtime::{self, Handle, Runtime},
12	task::JoinHandle,
13};
14
15use super::PoolConfig;
16use crate::sync::mutex::Mutex;
17
18struct PoolsInner {
19	system: Arc<ThreadPool>,
20	query: Arc<ThreadPool>,
21	commit: Arc<ThreadPool>,
22	background: Arc<ThreadPool>,
23	tokio_handle: Option<Handle>,
24	tokio: Mutex<Option<Runtime>>,
25}
26
27impl PoolsInner {
28	fn take_tokio(&self) -> Option<Runtime> {
29		self.tokio.lock().take()
30	}
31}
32
33impl Drop for PoolsInner {
34	fn drop(&mut self) {
35		if let Some(rt) = self.take_tokio() {
36			if runtime::Handle::try_current().is_err() {
37				rt.shutdown_timeout(Duration::from_secs(5));
38			} else {
39				rt.shutdown_background();
40			}
41		}
42	}
43}
44
45#[derive(Clone)]
46pub struct Pools {
47	inner: Arc<PoolsInner>,
48}
49
50impl Default for Pools {
51	fn default() -> Self {
52		Self::new(PoolConfig::default())
53	}
54}
55
56impl Pools {
57	pub fn new(config: PoolConfig) -> Self {
58		let system = Self::build_pool(config.system_threads, "system-pool");
59		let query = Self::build_pool(config.query_threads, "query-pool");
60		let commit = Self::build_pool(config.commit_threads, "commit-pool");
61		let background = Self::build_pool(config.background_threads, "background-pool");
62		let (tokio_handle, tokio) = Self::build_async_runtime(config.async_threads);
63
64		Self {
65			inner: Arc::new(PoolsInner {
66				system,
67				query,
68				commit,
69				background,
70				tokio_handle,
71				tokio,
72			}),
73		}
74	}
75
76	fn build_pool(threads: usize, name_prefix: &'static str) -> Arc<ThreadPool> {
77		Arc::new(
78			ThreadPoolBuilder::new()
79				.num_threads(threads)
80				.thread_name(move |i| format!("{name_prefix}-{i}"))
81				.build()
82				.unwrap_or_else(|_| panic!("failed to build {name_prefix} thread pool")),
83		)
84	}
85
86	#[inline]
87	fn build_async_runtime(threads: usize) -> (Option<Handle>, Mutex<Option<Runtime>>) {
88		let (tokio_handle, tokio) = if threads > 0 {
89			let rt = runtime::Builder::new_multi_thread()
90				.worker_threads(threads)
91				.thread_name("async")
92				.enable_all()
93				.build()
94				.expect("failed to build tokio runtime");
95			let handle = rt.handle().clone();
96			(Some(handle), Mutex::new(Some(rt)))
97		} else {
98			(None, Mutex::new(None))
99		};
100
101		reifydb_assertions! {
102			let handle_present = tokio_handle.is_some();
103			let runtime_present = tokio.lock().is_some();
104			assert!(
105				handle_present == runtime_present,
106				"async handle/runtime presence must agree (handle_present={handle_present}, runtime_present={runtime_present}); \
107				 a handle without its runtime makes tokio_handle() dispatch onto a runtime that Drop never shuts down (thread leak), \
108				 and a runtime without a handle makes spawn()/block_on() panic via the expect in tokio_handle()"
109			);
110		}
111
112		(tokio_handle, tokio)
113	}
114
115	pub fn shutdown(&self) {
116		if let Some(rt) = self.inner.take_tokio() {
117			if runtime::Handle::try_current().is_err() {
118				rt.shutdown_timeout(Duration::from_secs(5));
119			} else {
120				rt.shutdown_background();
121			}
122		}
123	}
124
125	pub fn system_pool(&self) -> &Arc<ThreadPool> {
126		&self.inner.system
127	}
128
129	pub fn system_thread_count(&self) -> usize {
130		self.inner.system.current_num_threads()
131	}
132
133	pub fn query_pool(&self) -> &Arc<ThreadPool> {
134		&self.inner.query
135	}
136
137	pub fn query_thread_count(&self) -> usize {
138		self.inner.query.current_num_threads()
139	}
140
141	pub fn commit_pool(&self) -> &Arc<ThreadPool> {
142		&self.inner.commit
143	}
144
145	pub fn commit_thread_count(&self) -> usize {
146		self.inner.commit.current_num_threads()
147	}
148
149	pub fn background_pool(&self) -> &Arc<ThreadPool> {
150		&self.inner.background
151	}
152
153	pub fn background_thread_count(&self) -> usize {
154		self.inner.background.current_num_threads()
155	}
156
157	fn tokio_handle(&self) -> Handle {
158		self.inner.tokio_handle.clone().expect("no tokio runtime configured (async_threads = 0)")
159	}
160
161	pub fn handle(&self) -> Handle {
162		self.tokio_handle()
163	}
164
165	pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
166	where
167		F: Future + Send + 'static,
168		F::Output: Send + 'static,
169	{
170		self.tokio_handle().spawn(future)
171	}
172
173	pub fn block_on<F>(&self, future: F) -> F::Output
174	where
175		F: Future,
176	{
177		self.tokio_handle().block_on(future)
178	}
179}
reifydb_runtime/pool/native.rs

reifydb_runtime/pool/
native.rs
