timeout_spawn.rs

use async_io::Timer;
use async_std::task::{spawn, Context, Poll};
use std::{future::Future, pin::Pin, time::Duration};

/// Awaits a future or spawns it out after a duration of time.
///
/// If timeout is reached, the future is not dropped, but instead is spawned out as a free task,
/// and `None` is returned.
///
/// # Examples
///
/// ```
/// # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
/// #
/// use std::time::Duration;
/// use async_std::future;
/// use libow::timeout_spawn;
///
/// let never = future::pending::<()>();
/// # let never = future::timeout(Duration::from_millis(20), never);
/// let dur = Duration::from_millis(5);
/// assert!(timeout_spawn(dur, never).await.is_none());
/// #
/// # Ok(()) }) }
/// ```
pub async fn timeout_spawn<F, T>(dur: Duration, f: F) -> Option<T>
where
	F: Future<Output = T> + Send + 'static,
	T: Send,
{
	TimeoutSpawn::new(f, dur).await
}

#[pin_project::pin_project]
struct TimeoutSpawn<F> {
	future: Option<Pin<Box<F>>>,
	#[pin]
	delay: Timer,
}

impl<F> TimeoutSpawn<F> {
	fn new(future: F, dur: Duration) -> Self {
		Self {
			future: Some(Box::pin(future)),
			delay: Timer::after(dur),
		}
	}
}

impl<F: Future + Send + 'static> Future for TimeoutSpawn<F>
where
	F::Output: Send,
{
	type Output = Option<F::Output>;

	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
		let this = self.project();
		let fut = this.future.take();
		match fut {
			Some(mut fut) => match fut.as_mut().poll(cx) {
				Poll::Ready(v) => Poll::Ready(Some(v)),
				Poll::Pending => match this.delay.poll(cx) {
					Poll::Ready(_) => {
						spawn(fut);
						Poll::Ready(None)
					}
					Poll::Pending => {
						this.future.replace(fut);
						Poll::Pending
					}
				},
			},
			None => Poll::Ready(None),
		}
	}
}

Public domain (CC0)