async_ringbuf/traits/
consumer.rs

1use core::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll, Waker},
5};
6use futures_util::future::FusedFuture;
7use ringbuf::traits::Consumer;
8#[cfg(feature = "std")]
9use std::io;
10
11pub trait AsyncConsumer: Consumer {
12    fn register_waker(&self, waker: &Waker);
13
14    fn close(&mut self);
15    /// Whether the corresponding producer was closed.
16    fn is_closed(&self) -> bool {
17        !self.write_is_held()
18    }
19
20    /// Pop item from the ring buffer waiting asynchronously if the buffer is empty.
21    ///
22    /// Future returns:
23    /// + `Some(item)` - an item is taken.
24    /// + `None` - the buffer is empty and the corresponding producer was dropped.
25    ///
26    /// # Cancel safety
27    ///
28    /// If future is cancelled then no item removed from the ring buffer.
29    fn pop(&mut self) -> PopFuture<'_, Self> {
30        PopFuture { owner: self, done: false }
31    }
32
33    /// Wait for the buffer to contain at least `count` items or to close.
34    ///
35    /// In debug mode panics if `count` is greater than buffer capacity.
36    ///
37    /// The method takes `&mut self` because only single [`WaitOccupiedFuture`] is allowed at a time.
38    ///
39    /// # Cancel safety
40    ///
41    /// The future can be safely cancelled.
42    fn wait_occupied(&mut self, count: usize) -> WaitOccupiedFuture<'_, Self> {
43        debug_assert!(count <= self.capacity().get());
44        WaitOccupiedFuture {
45            owner: self,
46            count,
47            done: false,
48        }
49    }
50
51    /// Fill slice with items from the ring buffer waiting asynchronously until slice filled or corresponding producer closed.
52    ///
53    /// Future returns:
54    /// + `Ok` - the whole slice is filled with the items from the buffer.
55    /// + `Err(count)` - the buffer is empty and the corresponding producer was dropped, number of items copied to slice is returned.
56    ///
57    /// # Cancel safety
58    ///
59    /// If future is cancelled then slice can be partially filled.
60    /// The number of items already copied can be examined by [`PopSliceFuture::count`].
61    fn pop_exact<'a: 'b, 'b>(&'a mut self, slice: &'b mut [Self::Item]) -> PopSliceFuture<'a, 'b, Self>
62    where
63        Self::Item: Copy,
64    {
65        PopSliceFuture {
66            owner: self,
67            slice: Some(slice),
68            count: 0,
69        }
70    }
71
72    /// Fill `vec` with items from the ring buffer waiting asynchronously until corresponding producer closed.
73    ///
74    /// # Cancel safety
75    ///
76    /// If future is cancelled then `vec` contains items taken from RB before cancellation.
77    #[cfg(feature = "alloc")]
78    fn pop_until_end<'a: 'b, 'b>(&'a mut self, vec: &'b mut alloc::vec::Vec<Self::Item>) -> PopVecFuture<'a, 'b, Self> {
79        PopVecFuture {
80            owner: self,
81            vec: Some(vec),
82        }
83    }
84
85    /// Poll for the next item in the ring buffer.
86    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>
87    where
88        Self: Unpin,
89    {
90        let mut waker_registered = false;
91        loop {
92            let closed = self.is_closed();
93            if let Some(item) = self.try_pop() {
94                break Poll::Ready(Some(item));
95            }
96            if closed {
97                break Poll::Ready(None);
98            }
99            if waker_registered {
100                break Poll::Pending;
101            }
102            self.register_waker(cx.waker());
103            waker_registered = true;
104        }
105    }
106
107    /// Poll reading bytes from byte buffer.
108    #[cfg(feature = "std")]
109    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>>
110    where
111        Self: AsyncConsumer<Item = u8> + Unpin,
112    {
113        if buf.is_empty() {
114            return Poll::Ready(Ok(0));
115        }
116        let mut waker_registered = false;
117        loop {
118            let closed = self.is_closed();
119            let len = self.pop_slice(buf);
120            if len != 0 || closed {
121                break Poll::Ready(Ok(len));
122            }
123            if waker_registered {
124                break Poll::Pending;
125            }
126            self.register_waker(cx.waker());
127            waker_registered = true;
128        }
129    }
130}
131
132/// # Cancel safety
133///
134/// If future is cancelled then no item removed from the ring buffer.
135#[must_use = "futures do nothing unless you `.await` or poll them"]
136pub struct PopFuture<'a, A: AsyncConsumer + ?Sized> {
137    owner: &'a mut A,
138    done: bool,
139}
140impl<A: AsyncConsumer> Unpin for PopFuture<'_, A> {}
141impl<A: AsyncConsumer> FusedFuture for PopFuture<'_, A> {
142    fn is_terminated(&self) -> bool {
143        self.done
144    }
145}
146impl<A: AsyncConsumer> Future for PopFuture<'_, A> {
147    type Output = Option<A::Item>;
148
149    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
150        let mut waker_registered = false;
151        loop {
152            assert!(!self.done);
153            let closed = self.owner.is_closed();
154            if let Some(item) = self.owner.try_pop() {
155                self.done = true;
156                break Poll::Ready(Some(item));
157            }
158            if closed {
159                self.done = true;
160                break Poll::Ready(None);
161            }
162            if waker_registered {
163                break Poll::Pending;
164            }
165            self.owner.register_waker(cx.waker());
166            waker_registered = true;
167        }
168    }
169}
170
171/// # Cancel safety
172///
173/// If future is cancelled then slice can be partially filled.
174#[must_use = "futures do nothing unless you `.await` or poll them"]
175pub struct PopSliceFuture<'a, 'b, A: AsyncConsumer + ?Sized>
176where
177    A::Item: Copy,
178{
179    owner: &'a mut A,
180    slice: Option<&'b mut [A::Item]>,
181    count: usize,
182}
183impl<A: AsyncConsumer> Unpin for PopSliceFuture<'_, '_, A> where A::Item: Copy {}
184impl<A: AsyncConsumer> FusedFuture for PopSliceFuture<'_, '_, A>
185where
186    A::Item: Copy,
187{
188    fn is_terminated(&self) -> bool {
189        self.slice.is_none()
190    }
191}
192impl<A: AsyncConsumer> Future for PopSliceFuture<'_, '_, A>
193where
194    A::Item: Copy,
195{
196    type Output = Result<(), usize>;
197
198    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
199        let mut waker_registered = false;
200        loop {
201            let closed = self.owner.is_closed();
202            let mut slice = self.slice.take().unwrap();
203            let len = self.owner.pop_slice(slice);
204            slice = &mut slice[len..];
205            self.count += len;
206            if slice.is_empty() {
207                break Poll::Ready(Ok(()));
208            }
209            if closed {
210                break Poll::Ready(Err(self.count));
211            }
212            self.slice.replace(slice);
213            if waker_registered {
214                break Poll::Pending;
215            }
216            self.owner.register_waker(cx.waker());
217            waker_registered = true;
218        }
219    }
220}
221impl<A: AsyncConsumer> PopSliceFuture<'_, '_, A>
222where
223    A::Item: Copy,
224{
225    /// Number of items already copied from the ring buufer to the slice provided.
226    pub fn count(&self) -> usize {
227        self.count
228    }
229}
230
231/// # Cancel safety
232///
233/// If future is cancelled then `vec` contains items taken from RB before cancellation.
234#[cfg(feature = "alloc")]
235#[must_use = "futures do nothing unless you `.await` or poll them"]
236pub struct PopVecFuture<'a, 'b, A: AsyncConsumer + ?Sized> {
237    owner: &'a mut A,
238    vec: Option<&'b mut alloc::vec::Vec<A::Item>>,
239}
240#[cfg(feature = "alloc")]
241impl<A: AsyncConsumer> Unpin for PopVecFuture<'_, '_, A> {}
242#[cfg(feature = "alloc")]
243impl<A: AsyncConsumer> FusedFuture for PopVecFuture<'_, '_, A> {
244    fn is_terminated(&self) -> bool {
245        self.vec.is_none()
246    }
247}
248#[cfg(feature = "alloc")]
249impl<A: AsyncConsumer> Future for PopVecFuture<'_, '_, A> {
250    type Output = ();
251
252    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
253        let mut waker_registered = false;
254        loop {
255            let closed = self.owner.is_closed();
256            let vec = self.vec.take().unwrap();
257
258            loop {
259                if vec.len() == vec.capacity() {
260                    vec.reserve(vec.capacity().max(16));
261                }
262                let n = self.owner.pop_slice_uninit(vec.spare_capacity_mut());
263                if n == 0 {
264                    break;
265                }
266                unsafe { vec.set_len(vec.len() + n) };
267            }
268
269            if closed {
270                break Poll::Ready(());
271            }
272            self.vec.replace(vec);
273            if waker_registered {
274                break Poll::Pending;
275            }
276            self.owner.register_waker(cx.waker());
277            waker_registered = true;
278        }
279    }
280}
281
282/// # Cancel safety
283///
284/// The future can be safely cancelled.
285#[must_use = "futures do nothing unless you `.await` or poll them"]
286pub struct WaitOccupiedFuture<'a, A: AsyncConsumer + ?Sized> {
287    owner: &'a A,
288    count: usize,
289    done: bool,
290}
291impl<A: AsyncConsumer> Unpin for WaitOccupiedFuture<'_, A> {}
292impl<A: AsyncConsumer> FusedFuture for WaitOccupiedFuture<'_, A> {
293    fn is_terminated(&self) -> bool {
294        self.done
295    }
296}
297impl<A: AsyncConsumer> Future for WaitOccupiedFuture<'_, A> {
298    type Output = ();
299
300    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
301        let mut waker_registered = false;
302        loop {
303            assert!(!self.done);
304            let closed = self.owner.is_closed();
305            if self.count <= self.owner.occupied_len() || closed {
306                self.done = true;
307                break Poll::Ready(());
308            }
309            if waker_registered {
310                break Poll::Pending;
311            }
312            self.owner.register_waker(cx.waker());
313            waker_registered = true;
314        }
315    }
316}
async_ringbuf/traits/consumer.rs

async_ringbuf/traits/
consumer.rs
