use crate::runtime::blocking::BlockingPool;
use crate::runtime::scheduler::CurrentThread;
use crate::runtime::{context, EnterGuard, Handle};
use crate::task::JoinHandle;
use std::future::Future;
use std::time::Duration;
cfg_rt_multi_thread! {
use crate::runtime::Builder;
use crate::runtime::scheduler::MultiThread;
}
#[derive(Debug)]
pub struct Runtime {
scheduler: Scheduler,
handle: Handle,
blocking_pool: BlockingPool,
}
#[derive(Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum RuntimeFlavor {
CurrentThread,
MultiThread,
}
#[derive(Debug)]
pub(super) enum Scheduler {
CurrentThread(CurrentThread),
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
MultiThread(MultiThread),
}
impl Runtime {
pub(super) fn from_parts(
scheduler: Scheduler,
handle: Handle,
blocking_pool: BlockingPool,
) -> Runtime {
Runtime {
scheduler,
handle,
blocking_pool,
}
}
cfg_not_wasi! {
#[cfg(feature = "rt-multi-thread")]
#[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))]
pub fn new() -> std::io::Result<Runtime> {
Builder::new_multi_thread().enable_all().build()
}
}
pub fn handle(&self) -> &Handle {
&self.handle
}
#[track_caller]
pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
self.handle.spawn(future)
}
#[track_caller]
pub fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R>
where
F: FnOnce() -> R + Send + 'static,
R: Send + 'static,
{
self.handle.spawn_blocking(func)
}
#[track_caller]
pub fn block_on<F: Future>(&self, future: F) -> F::Output {
#[cfg(all(tokio_unstable, feature = "tracing"))]
let future = crate::util::trace::task(
future,
"block_on",
None,
crate::runtime::task::Id::next().as_u64(),
);
let _enter = self.enter();
match &self.scheduler {
Scheduler::CurrentThread(exec) => exec.block_on(&self.handle.inner, future),
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
Scheduler::MultiThread(exec) => exec.block_on(&self.handle.inner, future),
}
}
pub fn enter(&self) -> EnterGuard<'_> {
self.handle.enter()
}
pub fn shutdown_timeout(mut self, duration: Duration) {
self.handle.inner.shutdown();
self.blocking_pool.shutdown(Some(duration));
}
pub fn shutdown_background(self) {
self.shutdown_timeout(Duration::from_nanos(0))
}
}
#[allow(clippy::single_match)] impl Drop for Runtime {
fn drop(&mut self) {
match &mut self.scheduler {
Scheduler::CurrentThread(current_thread) => {
let _guard = context::try_set_current(&self.handle.inner);
current_thread.shutdown(&self.handle.inner);
}
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
Scheduler::MultiThread(multi_thread) => {
multi_thread.shutdown(&self.handle.inner);
}
}
}
}
cfg_metrics! {
impl Runtime {
pub fn metrics(&self) -> crate::runtime::RuntimeMetrics {
self.handle.metrics()
}
}
}