use std::{
fmt::Debug,
future::Future,
mem::ManuallyDrop,
ops::{Deref, DerefMut},
};
use tokio::{
runtime::{Handle, Runtime},
task::JoinHandle,
};
pub struct BackgroundShutdownRuntime(ManuallyDrop<Runtime>);
impl Debug for BackgroundShutdownRuntime {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_tuple("BackgroundShutdownRuntime").finish()
}
}
impl Drop for BackgroundShutdownRuntime {
fn drop(&mut self) {
let runtime = unsafe { ManuallyDrop::take(&mut self.0) };
#[cfg(madsim)]
drop(runtime);
#[cfg(not(madsim))]
runtime.shutdown_background();
}
}
impl Deref for BackgroundShutdownRuntime {
type Target = Runtime;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl DerefMut for BackgroundShutdownRuntime {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl From<Runtime> for BackgroundShutdownRuntime {
fn from(runtime: Runtime) -> Self {
Self(ManuallyDrop::new(runtime))
}
}
#[derive(Debug)]
pub struct SingletonHandle(Handle);
impl From<Handle> for SingletonHandle {
fn from(handle: Handle) -> Self {
Self(handle)
}
}
impl SingletonHandle {
pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
self.0.spawn(future)
}
pub fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R>
where
F: FnOnce() -> R + Send + 'static,
R: Send + 'static,
{
self.0.spawn_blocking(func)
}
#[cfg(not(madsim))]
pub fn block_on<F: Future>(&self, future: F) -> F::Output {
self.0.block_on(future)
}
#[cfg(madsim)]
pub fn block_on<F: Future>(&self, _: F) -> F::Output {
unimplemented!("`block_on()` is not supported with madsim")
}
}