use std::{future::Future, io, sync::Arc};
use tokio::task::{JoinHandle, LocalSet};
#[derive(Debug)]
enum RuntimeInner {
Owned(tokio::runtime::Runtime),
Shared(Arc<tokio::runtime::Runtime>),
Static(&'static tokio::runtime::Runtime),
}
#[derive(Debug)]
pub struct Runtime {
local: LocalSet,
rt: RuntimeInner,
}
pub(crate) fn default_tokio_runtime() -> io::Result<tokio::runtime::Runtime> {
tokio::runtime::Builder::new_current_thread()
.enable_io()
.enable_time()
.build()
}
impl Runtime {
#[allow(clippy::new_ret_no_self)]
pub fn new() -> io::Result<Self> {
let rt = default_tokio_runtime()?;
Ok(Runtime {
rt: RuntimeInner::Owned(rt),
local: LocalSet::new(),
})
}
fn tokio_runtime_ref(&self) -> &tokio::runtime::Runtime {
match &self.rt {
RuntimeInner::Owned(rt) => rt,
RuntimeInner::Shared(rt) => rt,
RuntimeInner::Static(rt) => rt,
}
}
#[track_caller]
pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future + 'static,
{
self.local.spawn_local(future)
}
pub fn tokio_runtime(&self) -> &tokio::runtime::Runtime {
self.tokio_runtime_ref()
}
#[track_caller]
pub fn block_on<F>(&self, f: F) -> F::Output
where
F: Future,
{
self.local.block_on(self.tokio_runtime_ref(), f)
}
}
impl From<tokio::runtime::Runtime> for Runtime {
fn from(rt: tokio::runtime::Runtime) -> Self {
Self {
local: LocalSet::new(),
rt: RuntimeInner::Owned(rt),
}
}
}
impl From<Arc<tokio::runtime::Runtime>> for Runtime {
fn from(rt: Arc<tokio::runtime::Runtime>) -> Self {
Self {
local: LocalSet::new(),
rt: RuntimeInner::Shared(rt),
}
}
}
impl From<&'static tokio::runtime::Runtime> for Runtime {
fn from(rt: &'static tokio::runtime::Runtime) -> Self {
Self {
local: LocalSet::new(),
rt: RuntimeInner::Static(rt),
}
}
}