1#![cfg_attr(not(test), no_std)]
2#![warn(missing_docs)]
3
4use core::future::Future;
7use core::pin::Pin;
8use core::task::{Context, Poll};
9use pin_project::pin_project;
10
11#[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
37pub trait PriorityExt: Future + Sized {
39 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}