use pin_project_lite::pin_project;
use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use tokio::runtime::{Handle, Runtime};
pin_project! {
pub struct TokioContext<F> {
#[pin]
inner: F,
handle: Handle,
}
}
impl<F> TokioContext<F> {
pub fn new(future: F, handle: Handle) -> TokioContext<F> {
TokioContext {
inner: future,
handle,
}
}
pub fn handle(&self) -> &Handle {
&self.handle
}
pub fn into_inner(self) -> F {
self.inner
}
}
impl<F: Future> Future for TokioContext<F> {
type Output = F::Output;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.project();
let handle = me.handle;
let fut = me.inner;
let _enter = handle.enter();
fut.poll(cx)
}
}
pub trait RuntimeExt {
fn wrap<F: Future>(&self, fut: F) -> TokioContext<F>;
}
impl RuntimeExt for Runtime {
fn wrap<F: Future>(&self, fut: F) -> TokioContext<F> {
TokioContext {
inner: fut,
handle: self.handle().clone(),
}
}
}