futures_prio/
lib.rs

1#![cfg_attr(not(test), no_std)]
2#![warn(missing_docs)]
3
4//! Futures combinator to set the priority of a Future.
5
6use core::future::Future;
7use core::pin::Pin;
8use core::task::{Context, Poll};
9use pin_project::pin_project;
10
11/// Future to set priority.
12#[pin_project]
13pub struct Priority<T: Future> {
14    #[pin]
15    inner: T,
16    prio: usize,
17    waited: usize,
18}
19
20impl<T: Future> Future for Priority<T> {
21    type Output = T::Output;
22
23    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
24        let this = self.project();
25
26        if this.waited >= this.prio {
27            *this.waited = 0;
28            this.inner.poll(cx)
29        } else {
30            *this.waited += 1;
31            cx.waker().wake_by_ref();
32            Poll::Pending
33        }
34    }
35}
36
37/// Extension trait to set priority on Futures.
38pub trait PriorityExt: Future + Sized {
39    /// Set the priority of a Future. A priority of 0 is equivalent to not calling this method, and
40    /// higher priorities will cause a Future to be polled less often.
41    fn priority(self, prio: usize) -> Priority<Self>;
42}
43
44impl<T: Future> PriorityExt for T {
45    fn priority(self, prio: usize) -> Priority<Self> {
46        Priority {
47            inner: self,
48            waited: 0,
49            prio,
50        }
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57    use core::future::Future;
58    use core::pin::Pin;
59    use core::task::{Context, Poll};
60    use futures_executor::block_on;
61
62    struct DummyFuture;
63
64    impl Future for DummyFuture {
65        type Output = ();
66
67        fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
68            Poll::Ready(())
69        }
70    }
71
72    #[test]
73    fn sanity_check() {
74        block_on(async {
75            assert_eq!(DummyFuture.await, ());
76            assert_eq!(DummyFuture.priority(1).await, ());
77            assert_eq!(DummyFuture.priority(2).await, ());
78        });
79    }
80}