use std::sync::{Arc, OnceLock};
use tokio::runtime::{Handle, Runtime};
use crate::RuntimeConfig;
static GLOBAL_RUNTIME: OnceLock<GlobalRuntime> = OnceLock::new();
#[derive(Debug, Clone)]
pub struct GlobalRuntime {
tokio_rt: Option<Arc<Runtime>>,
tokio_handle: Handle,
config: RuntimeConfig,
}
impl RuntimeConfig {
#[inline]
pub fn auto() -> Self {
let snapshot = zenith_capability::detection::EnvironmentDetector::new().detect();
let cpus = snapshot.cpu_cores.max(1);
let stack = if snapshot.numa_nodes >= 2 || cpus >= 16 {
4 * 1024 * 1024
} else {
2 * 1024 * 1024
};
Self {
worker_threads: cpus,
enable_io: true,
enable_time: true,
stack_size: stack,
}
}
}
#[inline]
pub fn init_global(config: RuntimeConfig) -> &'static GlobalRuntime {
GLOBAL_RUNTIME.get_or_init(|| build_global_runtime(config))
}
#[inline]
pub fn global_runtime() -> &'static GlobalRuntime {
if let Some(rt) = GLOBAL_RUNTIME.get() {
return rt;
}
init_global(RuntimeConfig::auto())
}
#[inline]
pub fn handle() -> Handle {
global_runtime().handle().clone()
}
#[inline]
pub fn spawn<F>(future: F) -> tokio::task::JoinHandle<F::Output>
where
F: std::future::Future + Send + 'static,
F::Output: Send + 'static,
{
global_runtime().spawn(future)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockOnError {
ReentrantCurrentThread,
}
impl std::fmt::Display for BlockOnError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BlockOnError::ReentrantCurrentThread => write!(
f,
"cannot block_on from within a current_thread runtime context"
),
}
}
}
impl std::error::Error for BlockOnError {}
#[inline]
pub fn block_on<F: std::future::Future>(
future: F,
) -> Result<F::Output, BlockOnError> {
global_runtime().block_on(future)
}
impl GlobalRuntime {
#[inline]
pub fn handle(&self) -> &Handle {
&self.tokio_handle
}
#[inline]
pub fn config(&self) -> &RuntimeConfig {
&self.config
}
#[inline]
pub fn spawn<F>(&self, future: F) -> tokio::task::JoinHandle<F::Output>
where
F: std::future::Future + Send + 'static,
F::Output: Send + 'static,
{
self.tokio_handle.spawn(future)
}
#[inline]
pub fn block_on<F: std::future::Future>(
&self,
future: F,
) -> Result<F::Output, BlockOnError> {
match Handle::try_current() {
Err(_) => Ok(match &self.tokio_rt {
Some(rt) => rt.block_on(future),
None => self.tokio_handle.block_on(future),
}),
Ok(current) => match current.runtime_flavor() {
tokio::runtime::RuntimeFlavor::MultiThread => {
Ok(tokio::task::block_in_place(|| match &self.tokio_rt {
Some(rt) => rt.block_on(future),
None => self.tokio_handle.block_on(future),
}))
}
_ => Err(BlockOnError::ReentrantCurrentThread),
},
}
}
}
fn build_global_runtime(config: RuntimeConfig) -> GlobalRuntime {
if let Some(rt) = build_tokio_runtime(&config) {
let arc_rt = Arc::new(rt);
let handle = arc_rt.handle().clone();
return GlobalRuntime {
tokio_rt: Some(arc_rt),
tokio_handle: handle,
config,
};
}
if let Ok(h) = Handle::try_current() {
eprintln!("[zenith-warn] Global runtime borrowing external tokio Handle - lifetime risk if external runtime shuts down");
return GlobalRuntime {
tokio_rt: None,
tokio_handle: h,
config,
};
}
let rt = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(rt) => rt,
Err(_e) => {
tracing::error!("Runtime build failed, aborting process");
std::process::abort();
}
};
let handle = rt.handle().clone();
GlobalRuntime {
tokio_rt: Some(Arc::new(rt)),
tokio_handle: handle,
config,
}
}
fn build_tokio_runtime(config: &RuntimeConfig) -> Option<Runtime> {
let mut builder = tokio::runtime::Builder::new_multi_thread();
builder.worker_threads(config.worker_threads.max(1));
builder.thread_stack_size(config.stack_size.max(4096));
if config.enable_io || config.enable_time {
if config.enable_io {
builder.enable_io();
}
if config.enable_time {
builder.enable_time();
}
}
builder.build().ok()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_config_auto_cpu_sane() {
let cfg = RuntimeConfig::auto();
assert!(cfg.worker_threads >= 1);
assert!(cfg.enable_io);
assert!(cfg.enable_time);
assert!(cfg.stack_size >= 4096);
}
#[test]
fn test_config_auto_vs_new_equal_same_input() {
let a = RuntimeConfig::auto();
let b = RuntimeConfig::auto();
assert_eq!(a.worker_threads, b.worker_threads);
assert_eq!(a.stack_size, b.stack_size);
}
#[test]
fn test_global_runtime_init_idempotent() {
let cfg = RuntimeConfig::new().with_worker_threads(2);
let r1 = init_global(cfg);
let r2 = init_global(RuntimeConfig::new().with_worker_threads(32));
assert!(std::ptr::eq(r1, r2));
assert_eq!(r1.config().worker_threads, r2.config().worker_threads);
}
#[test]
fn test_global_spawn_works() {
let result = block_on(async {
let h = spawn(async { 42_u32 });
h.await.unwrap()
})
.unwrap();
assert_eq!(result, 42);
}
#[test]
fn test_handle_clone_works() {
let h1 = handle();
let h2 = handle();
let r = block_on(async move {
let j1 = h1.spawn(async { 1 });
let j2 = h2.spawn(async { 2 });
j1.await.unwrap() + j2.await.unwrap()
})
.unwrap();
assert_eq!(r, 3);
}
#[test]
fn test_block_on_reentrant_current_thread_returns_err() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let result = block_on(async { 1u32 });
assert_eq!(result, Err(BlockOnError::ReentrantCurrentThread));
});
}
#[test]
fn test_block_on_reentrant_multi_thread_ok() {
let rt = tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let result = block_on(async { 7u32 }).unwrap();
assert_eq!(result, 7);
});
}
#[test]
fn test_build_with_zero_threads_no_panic() {
let cfg = RuntimeConfig::new().with_worker_threads(0).with_stack_size(0);
let gr = build_global_runtime(cfg);
let _ = gr.block_on(async { 1 });
}
#[test]
fn test_global_runtime_config_snapshot_kept() {
let cfg = RuntimeConfig::new().with_worker_threads(1).with_stack_size(8 * 1024);
let gr = build_global_runtime(cfg);
assert!(gr.config().worker_threads >= 1);
assert!(gr.config().stack_size >= 8 * 1024);
}
const _: fn() = || {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<GlobalRuntime>();
};
}