use std::{future::Future, pin::Pin};
use futures::{FutureExt, future};
use tari_shutdown::Shutdown;
use tokio::{runtime, runtime::Runtime, task, task::JoinError};
pub fn create_runtime() -> Runtime {
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.expect("Could not create runtime")
}
pub fn spawn_until_shutdown<F>(fut: F) -> Shutdown
where F: Future<Output = ()> + Send + 'static {
let shutdown = Shutdown::new();
let signal = shutdown.to_signal();
task::spawn(async move {
futures::pin_mut!(fut);
future::select(signal, fut).await;
});
shutdown
}
type BoxedJoinFuture = Pin<Box<dyn Future<Output = Result<(), JoinError>>>>;
pub struct TestRuntime {
inner: Runtime,
handles: Vec<BoxedJoinFuture>,
}
impl TestRuntime {
pub fn block_on<F: Future>(&mut self, future: F) -> F::Output {
self.inner.block_on(future)
}
pub fn spawn<F>(&mut self, future: F)
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let handle = self.inner.spawn(future);
self.handles.push(
handle
.map(|result| match result {
Ok(_) => Ok(()),
Err(err) => Err(err),
})
.boxed(),
);
}
pub fn spawn_unchecked<F>(&mut self, future: F) -> task::JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
self.inner.spawn(future)
}
pub fn handle(&self) -> &runtime::Handle {
self.inner.handle()
}
}
impl From<Runtime> for TestRuntime {
fn from(rt: Runtime) -> Self {
Self {
inner: rt,
handles: Vec::new(),
}
}
}