1mod clock;
 2mod executor;
 3mod test_scheduler;
 4#[cfg(test)]
 5mod tests;
 6
 7pub use clock::*;
 8pub use executor::*;
 9pub use test_scheduler::*;
10
11use async_task::Runnable;
12use futures::{FutureExt as _, channel::oneshot, future::LocalBoxFuture};
13use std::{
14    future::Future,
15    pin::Pin,
16    task::{Context, Poll},
17    time::Duration,
18};
19
20pub trait Scheduler: Send + Sync {
21    fn block(&self, future: LocalBoxFuture<()>, timeout: Option<Duration>);
22    fn schedule_foreground(&self, session_id: SessionId, runnable: Runnable);
23    fn schedule_background(&self, runnable: Runnable);
24    fn timer(&self, timeout: Duration) -> Timer;
25    fn is_main_thread(&self) -> bool;
26}
27
28impl dyn Scheduler {
29    pub fn block_on<Fut: Future>(&self, future: Fut) -> Fut::Output {
30        let mut output = None;
31        self.block(async { output = Some(future.await) }.boxed_local(), None);
32        output.unwrap()
33    }
34
35    pub fn block_with_timeout<Fut: Unpin + Future>(
36        &self,
37        future: &mut Fut,
38        timeout: Duration,
39    ) -> Option<Fut::Output> {
40        let mut output = None;
41        self.block(
42            async { output = Some(future.await) }.boxed_local(),
43            Some(timeout),
44        );
45        output
46    }
47}
48
49#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
50pub struct SessionId(u16);
51
52pub struct Timer(oneshot::Receiver<()>);
53
54impl Future for Timer {
55    type Output = ();
56
57    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
58        match self.0.poll_unpin(cx) {
59            Poll::Ready(_) => Poll::Ready(()),
60            Poll::Pending => Poll::Pending,
61        }
62    }
63}