schematics/src/async_utils.rs

use std::{
  future::Future,
  pin::Pin,
  sync::{Arc, Mutex},
  task::{Context, Poll, Waker},
  thread,
  time::Duration,
};
use tokio::runtime::Builder;

pub use async_trait::async_trait;

pub fn run<F: Future>(future: F) -> F::Output {
  Builder::new_current_thread()
    .enable_io()
    .build()
    .unwrap()
    .block_on(future)
}
pub use tokio::join;
pub use tokio::try_join;

#[derive(Debug)]
pub struct TimerFuture {
  state: Arc<Mutex<State>>,
}

#[derive(Debug)]
enum State {
  NotStarted(Duration),
  Running(Waker),
  Completed,
}

impl Future for TimerFuture {
  type Output = ();
  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    let mut state = self.state.lock().unwrap();
    if let State::Completed = *state {
      return Poll::Ready(());
    }

    if let State::NotStarted(duration) = *state {
      let thread_state = self.state.clone();
      thread::spawn(move || {
        thread::sleep(duration);
        let mut state = thread_state.lock().unwrap();
        let waker = if let State::Running(waker) = &*state {
          Some(waker.clone())
        } else {
          None
        };
        *state = State::Completed;
        if let Some(w) = waker {
          w.wake()
        }
      });
    }

    *state = State::Running(cx.waker().clone());
    Poll::Pending
  }
}

pub fn sleep(duration: Duration) -> impl Future<Output = ()> {
  TimerFuture {
    state: Arc::new(Mutex::new(State::NotStarted(duration))),
  }
}

#[cfg(test)]
mod test {
  use crate::async_utils::{run, sleep};
  use futures_util::future::FutureExt;
  use std::time::{Duration, Instant};

  #[test]
  fn test_sleep() {
    run(async {
      let start = Instant::now();
      let sleep = sleep(Duration::from_millis(100)).fuse();
      let ok = async {}.fuse();
      futures_util::pin_mut!(sleep, ok);
      loop {
        futures_util::select! {
          _ = sleep => {},
          _ = ok => assert!((Instant::now() - start).as_millis() < 100),
          complete => break,
        }
      }
    })
  }
}