reifydb_runtime/pool/
host.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 =
66			ActorPool::new(config.coordination_threads, config.flow_threads, config.maintenance_threads);
67		let task = TaskPool::new(config.task_threads, "task");
68		let compute = ComputePool::new(config.compute_threads, "compute");
69		let (tokio_handle, tokio) = Self::build_async_runtime(config.async_threads);
70
71		Self {
72			inner: Arc::new(PoolsInner {
73				actors,
74				task,
75				compute,
76				tokio_handle,
77				tokio,
78			}),
79		}
80	}
81
82	#[inline]
83	fn build_async_runtime(threads: usize) -> (Option<Handle>, Mutex<Option<Runtime>>) {
84		let (tokio_handle, tokio) = if threads > 0 {
85			let rt = runtime::Builder::new_multi_thread()
86				.worker_threads(threads)
87				.thread_name("async")
88				.enable_all()
89				.build()
90				.expect("failed to build tokio runtime");
91			let handle = rt.handle().clone();
92			(Some(handle), Mutex::new(Some(rt)))
93		} else {
94			(None, Mutex::new(None))
95		};
96
97		reifydb_assertions! {
98			let handle_present = tokio_handle.is_some();
99			let runtime_present = tokio.lock().is_some();
100			assert!(
101				handle_present == runtime_present,
102				"async handle/runtime presence must agree (handle_present={handle_present}, runtime_present={runtime_present}); \
103				 a handle without its runtime makes tokio_handle() dispatch onto a runtime that Drop never shuts down (thread leak), \
104				 and a runtime without a handle makes spawn()/block_on() panic via the expect in tokio_handle()"
105			);
106		}
107
108		(tokio_handle, tokio)
109	}
110
111	pub fn shutdown(&self) {
112		if let Some(rt) = self.inner.take_tokio() {
113			if runtime::Handle::try_current().is_err() {
114				rt.shutdown_timeout(Duration::from_secs(5));
115			} else {
116				rt.shutdown_background();
117			}
118		}
119		self.inner.actors.shutdown();
120		self.inner.task.shutdown();
121	}
122
123	pub fn spawn_task(&self, job: impl FnOnce() + Send + 'static) {
124		self.inner.task.spawn(job);
125	}
126
127	pub fn task_thread_count(&self) -> usize {
128		self.inner.task.thread_count()
129	}
130
131	pub fn compute(&self) -> &ComputePool {
132		&self.inner.compute
133	}
134
135	pub fn compute_thread_count(&self) -> usize {
136		self.inner.compute.thread_count()
137	}
138
139	pub fn coordination_thread_count(&self) -> usize {
140		self.inner.actors.coordination().thread_count()
141	}
142
143	pub fn flow_thread_count(&self) -> usize {
144		self.inner.actors.flow().thread_count()
145	}
146
147	pub fn maintenance_thread_count(&self) -> usize {
148		self.inner.actors.maintenance().thread_count()
149	}
150
151	pub(crate) fn actor_pool(&self) -> &ActorPool {
152		&self.inner.actors
153	}
154
155	pub(crate) fn task_injector(&self) -> Schedule {
156		Schedule::Injector(self.inner.task.injector())
157	}
158
159	fn tokio_handle(&self) -> Handle {
160		self.inner.tokio_handle.clone().expect("no tokio runtime configured (async_threads = 0)")
161	}
162
163	pub fn handle(&self) -> Handle {
164		self.tokio_handle()
165	}
166
167	pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
168	where
169		F: Future + Send + 'static,
170		F::Output: Send + 'static,
171	{
172		self.tokio_handle().spawn(future)
173	}
174
175	pub fn block_on<F>(&self, future: F) -> F::Output
176	where
177		F: Future,
178	{
179		self.tokio_handle().block_on(future)
180	}
181}
reifydb_runtime/pool/host.rs

reifydb_runtime/pool/
host.rs
