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()
    .expect("Error setting up async runtime")
    .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();
    match *state {
      State::Completed => return Poll::Ready(()),
      State::NotStarted(duration) => {
        let thread_state = self.state.clone();
        thread::spawn(move || {
          thread::sleep(duration);
          let mut state = thread_state.lock().unwrap();
          if let State::Running(waker) = std::mem::replace(&mut *state, State::Completed) {
            waker.wake();
          };
        });
      }
      State::Running(_) => {}
    };

    *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!(start.elapsed().as_millis() < 100),
          complete => break,
        }
      }
    });
  }
}