use std::future::Future;
use std::pin::Pin;
use std::sync::OnceLock;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub trait Runtime: Send + Sync + 'static {
fn spawn(&self, fut: BoxFuture<'static, ()>) -> SpawnHandle;
fn sleep(&self, dur: Duration) -> BoxFuture<'static, ()>;
fn now(&self) -> Duration;
fn rng_u64(&self) -> u64;
fn derive_seed(&self, label: &str) -> u64;
fn can_spawn(&self) -> bool {
true
}
}
pub struct SpawnHandle {
abort: Box<dyn Fn() + Send + Sync>,
}
impl SpawnHandle {
pub fn from_abort(abort: impl Fn() + Send + Sync + 'static) -> Self {
Self {
abort: Box::new(abort),
}
}
pub fn noop() -> Self {
Self::from_abort(|| {})
}
pub fn abort(&self) {
(self.abort)()
}
}
#[derive(Default, Clone, Copy)]
pub struct TokioRuntime;
fn tokio_base() -> Instant {
static BASE: OnceLock<Instant> = OnceLock::new();
*BASE.get_or_init(Instant::now)
}
impl Runtime for TokioRuntime {
#[inline]
fn spawn(&self, fut: BoxFuture<'static, ()>) -> SpawnHandle {
let handle = tokio::spawn(fut);
let abort = handle.abort_handle();
SpawnHandle::from_abort(move || abort.abort())
}
#[inline]
fn can_spawn(&self) -> bool {
tokio::runtime::Handle::try_current().is_ok()
}
#[inline]
fn sleep(&self, dur: Duration) -> BoxFuture<'static, ()> {
Box::pin(tokio::time::sleep(dur))
}
#[inline]
fn now(&self) -> Duration {
tokio_base().elapsed()
}
#[inline]
fn rng_u64(&self) -> u64 {
rand::random()
}
#[inline]
fn derive_seed(&self, _label: &str) -> u64 {
rand::random()
}
}
pub fn yield_now() -> YieldNow {
YieldNow { yielded: false }
}
pub struct YieldNow {
yielded: bool,
}
impl Future for YieldNow {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.yielded {
Poll::Ready(())
} else {
self.yielded = true;
cx.waker().wake_by_ref();
Poll::Pending
}
}
}