queue_ext/
queue_stream.rs

1use std::collections::VecDeque;
2use std::fmt;
3use std::marker::PhantomData;
4use std::marker::Unpin;
5use std::ops::{Deref, DerefMut};
6use std::pin::Pin;
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::Arc;
9use std::sync::Mutex;
10use std::task::{Context, Poll};
11
12use futures::task::AtomicWaker;
13use futures::Stream;
14use pin_project_lite::pin_project;
15
16use super::Waker;
17
18pin_project! {
19    #[must_use = "streams do nothing unless polled"]
20    pub struct QueueStream<Q, Item, F> {
21        #[pin]
22        q: Q,
23        #[pin]
24        f: F,
25        recv_task: Arc<AtomicWaker>,
26        parked_queue: Arc<Mutex<VecDeque<std::task::Waker>>>,
27        closed: Arc<AtomicBool>,
28        _item: PhantomData<Item>,
29    }
30}
31
32unsafe impl<Q, Item, F> Sync for QueueStream<Q, Item, F> {}
33
34unsafe impl<Q, Item, F> Send for QueueStream<Q, Item, F> {}
35
36impl<Q, Item, F> Clone for QueueStream<Q, Item, F>
37where
38    Q: Clone,
39    F: Clone,
40{
41    #[inline]
42    fn clone(&self) -> Self {
43        Self {
44            q: self.q.clone(),
45            f: self.f.clone(),
46            recv_task: self.recv_task.clone(),
47            parked_queue: self.parked_queue.clone(),
48            closed: self.closed.clone(),
49            _item: PhantomData,
50        }
51    }
52}
53
54impl<Q, Item, F> fmt::Debug for QueueStream<Q, Item, F>
55where
56    Q: fmt::Debug,
57{
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        f.debug_struct("QueueStream")
60            .field("queue", &self.q)
61            .finish()
62    }
63}
64
65impl<Q: Unpin, Item, F> QueueStream<Q, Item, F> {
66    #[inline]
67    pub(super) fn new(q: Q, f: F) -> Self {
68        Self {
69            q,
70            f,
71            recv_task: Arc::new(AtomicWaker::new()),
72            parked_queue: Arc::new(Mutex::new(VecDeque::default())),
73            closed: Arc::new(AtomicBool::new(false)),
74            _item: PhantomData,
75        }
76    }
77
78    #[inline]
79    pub fn is_closed(&self) -> bool {
80        self.closed.load(Ordering::SeqCst)
81    }
82}
83
84impl<Q, Item, F> Waker for QueueStream<Q, Item, F> {
85    #[inline]
86    fn rx_wake(&self) {
87        self.recv_task.wake()
88    }
89
90    #[inline]
91    fn tx_park(&self, w: std::task::Waker) {
92        self.parked_queue.lock().unwrap().push_back(w);
93    }
94
95    #[inline]
96    fn close_channel(&self) {
97        if !self.closed.load(Ordering::SeqCst) {
98            self.closed.store(true, Ordering::SeqCst);
99            self.rx_wake();
100            if let Some(w) = self.parked_queue.lock().unwrap().pop_front() {
101                w.wake();
102            }
103        }
104    }
105
106    #[inline]
107    fn is_closed(&self) -> bool {
108        self.closed.load(Ordering::SeqCst)
109    }
110}
111
112impl<Q, Item, F> Stream for QueueStream<Q, Item, F>
113where
114    Q: Unpin,
115    F: Fn(Pin<&mut Q>, &mut Context<'_>) -> Poll<Option<Item>>,
116{
117    type Item = Item;
118
119    fn poll_next(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
120        let mut this = self.project();
121        let f = this.f.as_mut();
122        match f(this.q.as_mut(), ctx) {
123            Poll::Ready(msg) => {
124                if let Some(w) = this.parked_queue.lock().unwrap().pop_front() {
125                    w.wake();
126                }
127                Poll::Ready(msg)
128            }
129            Poll::Pending => {
130                if this.closed.load(Ordering::SeqCst) {
131                    Poll::Ready(None)
132                } else {
133                    this.recv_task.register(ctx.waker());
134                    f(this.q.as_mut(), ctx)
135                }
136            }
137        }
138    }
139}
140
141impl<Q, Item, F> Deref for QueueStream<Q, Item, F> {
142    type Target = Q;
143    #[inline]
144    fn deref(&self) -> &Self::Target {
145        &self.q
146    }
147}
148
149impl<Q, Item, F> DerefMut for QueueStream<Q, Item, F> {
150    #[inline]
151    fn deref_mut(&mut self) -> &mut Self::Target {
152        &mut self.q
153    }
154}
queue_ext/queue_stream.rs

queue_ext/
queue_stream.rs
