use std::time::Duration;
use tokio::runtime::Handle;
use tokio::task::JoinHandle;
use weida_core::Error;
#[derive(Clone)]
pub struct Exec {
handle: Handle,
}
impl Exec {
pub fn from_handle(handle: Handle) -> Exec {
Exec { handle }
}
pub fn current() -> Result<Exec, Error> {
Handle::try_current()
.map(Exec::from_handle)
.map_err(|_| Error::Runtime("no ambient tokio runtime to run on".into()))
}
pub fn owned(worker_threads: usize, thread_name: &str) -> Result<(Exec, OwnedReactor), Error> {
if worker_threads == 0 {
return Err(Error::Runtime("worker_threads must be at least 1".into()));
}
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.worker_threads(worker_threads)
.thread_name(thread_name)
.build()
.map_err(Error::Io)?;
let exec = Exec::from_handle(runtime.handle().clone());
Ok((exec, OwnedReactor(Some(runtime))))
}
pub fn enter(&self) -> tokio::runtime::EnterGuard<'_> {
self.handle.enter()
}
pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
self.handle.spawn(future)
}
pub fn sleep(&self, duration: Duration) -> tokio::time::Sleep {
let _guard = self.handle.enter();
tokio::time::sleep(duration)
}
pub async fn within<F: Future>(&self, limit: Duration, future: F) -> Option<F::Output> {
let deadline = self.sleep(limit);
tokio::select! {
output = future => Some(output),
() = deadline => None,
}
}
pub async fn resolve(
&self,
host: &str,
port: u16,
max_addresses: usize,
) -> Result<Vec<std::net::SocketAddr>, Error> {
use crate::resolve::Resolver;
crate::resolve::SystemResolver
.resolve(self, host, Some(port), max_addresses)
.await
}
}
impl std::fmt::Debug for Exec {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Exec").finish_non_exhaustive()
}
}
pub struct OwnedReactor(Option<tokio::runtime::Runtime>);
impl Drop for OwnedReactor {
fn drop(&mut self) {
if let Some(runtime) = self.0.take() {
runtime.shutdown_background();
}
}
}
impl std::fmt::Debug for OwnedReactor {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OwnedReactor").finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn an_exec_without_an_ambient_reactor_fails() {
let err = Exec::current().unwrap_err();
assert!(matches!(err, Error::Runtime(_)), "{err:?}");
}
#[test]
fn zero_worker_threads_is_rejected() {
let err = Exec::owned(0, "test").unwrap_err();
assert!(matches!(err, Error::Runtime(_)), "{err:?}");
}
#[test]
fn an_owned_reactor_needs_no_ambient_one() {
assert!(Handle::try_current().is_err());
let (exec, reactor) = Exec::owned(1, "test").expect("owned reactor");
let joined = exec.spawn(async { 7u8 });
assert_eq!(futures::executor::block_on(joined).expect("task"), 7);
drop(reactor);
}
#[tokio::test]
async fn within_gives_up_on_a_future_that_never_finishes() {
let exec = Exec::current().expect("ambient runtime");
assert!(
exec.within(Duration::from_millis(1), std::future::pending::<()>())
.await
.is_none()
);
assert_eq!(
exec.within(Duration::from_secs(30), async { 3 }).await,
Some(3)
);
}
}