You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
93 lines
2.0 KiB
93 lines
2.0 KiB
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();
|
|
if matches!(*state, State::Completed) {
|
|
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!(start.elapsed().as_millis() < 100),
|
|
complete => break,
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|