Skip to main content

async_ringbuf/traits/
producer.rs

1use core::{
2    future::Future,
3    iter::Peekable,
4    pin::Pin,
5    task::{Context, Poll, Waker},
6};
7use futures_util::future::FusedFuture;
8use ringbuf::traits::Producer;
9#[cfg(feature = "std")]
10use std::io;
11
12pub trait AsyncProducer: Producer {
13    fn register_waker(&self, waker: &Waker);
14
15    fn close(&mut self);
16    /// Whether the corresponding consumer was closed.
17    fn is_closed(&self) -> bool {
18        !self.read_is_held()
19    }
20
21    /// Push item to the ring buffer waiting asynchronously if the buffer is full.
22    ///
23    /// Future returns:
24    /// + `Ok` - item successfully pushed.
25    /// + `Err(item)` - the corresponding consumer was dropped, item is returned back.
26    ///
27    /// # Cancel safety
28    ///
29    /// If future is cancelled no item pushed to the RB.
30    fn push(&mut self, item: Self::Item) -> PushFuture<'_, Self> {
31        PushFuture {
32            owner: self,
33            item: Some(item),
34        }
35    }
36
37    /// Push items from iterator waiting asynchronously if the buffer is full.
38    ///
39    /// Future returns:
40    /// + `true` - iterator ended.
41    /// + `false` - the corresponding consumer was dropped.
42    ///
43    /// # Cancel safety
44    ///
45    /// If future is cancelled then remaining items are left in iterator.
46    /// You can get the iterator by using [`PushIterFuture::inner`], [`PushIterFuture::inner_mut`] and [`PushIterFuture::into_inner`].
47    /// *Note that the iterator is [`Peekable`].*
48    fn push_iter_all<I: Iterator<Item = Self::Item>>(&mut self, iter: I) -> PushIterFuture<'_, Self, I> {
49        PushIterFuture {
50            owner: self,
51            iter: Some(iter.peekable()),
52        }
53    }
54
55    /// Wait for the buffer to have at least `count` free places for items or to close.
56    ///
57    /// In debug mode panics if `count` is greater than buffer capacity.
58    ///
59    /// The method takes `&mut self` because only single [`WaitVacantFuture`] is allowed at a time.
60    ///
61    /// # Cancel safety
62    ///
63    /// You can safely cancel this future.
64    fn wait_vacant(&mut self, count: usize) -> WaitVacantFuture<'_, Self> {
65        debug_assert!(count <= self.capacity().get());
66        WaitVacantFuture {
67            owner: self,
68            count,
69            done: false,
70        }
71    }
72
73    /// Copy slice contents to the buffer waiting asynchronously if the buffer is full.
74    ///
75    /// Future returns:
76    /// + `Ok` - all slice contents are copied.
77    /// + `Err(count)` - the corresponding consumer was dropped, number of copied items returned.
78    ///
79    /// # Cancel safety
80    ///
81    /// On cancel the slice can be copied partially.
82    /// The number of items already copied can be examined by [`PushSliceFuture::count`].
83    fn push_exact<'a: 'b, 'b>(&'a mut self, slice: &'b [Self::Item]) -> PushSliceFuture<'a, 'b, Self>
84    where
85        Self::Item: Copy,
86    {
87        PushSliceFuture {
88            owner: self,
89            slice: Some(slice),
90            count: 0,
91        }
92    }
93
94    /// Poll the ring buffer has free slot for at least one item and the corresponding consumer is not closed.
95    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<bool> {
96        let mut waker_registered = false;
97        loop {
98            if self.is_closed() {
99                break Poll::Ready(false);
100            }
101            if !self.is_full() {
102                break Poll::Ready(true);
103            }
104            if waker_registered {
105                break Poll::Pending;
106            }
107            self.register_waker(cx.waker());
108            waker_registered = true;
109        }
110    }
111
112    /// Poll writing bytes into byte buffer.
113    #[cfg(feature = "std")]
114    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>>
115    where
116        Self: AsyncProducer<Item = u8> + Unpin,
117    {
118        if buf.is_empty() {
119            return Poll::Ready(Ok(0));
120        }
121        let mut waker_registered = false;
122        loop {
123            if self.is_closed() {
124                break Poll::Ready(Ok(0));
125            }
126            let count = self.push_slice(buf);
127            if count > 0 {
128                break Poll::Ready(Ok(count));
129            }
130            if waker_registered {
131                break Poll::Pending;
132            }
133            self.register_waker(cx.waker());
134            waker_registered = true;
135        }
136    }
137}
138
139/// # Cancel safety
140///
141/// If future is cancelled no item pushed to the RB.
142#[must_use = "futures do nothing unless you `.await` or poll them"]
143pub struct PushFuture<'a, A: AsyncProducer + ?Sized> {
144    owner: &'a mut A,
145    item: Option<A::Item>,
146}
147impl<A: AsyncProducer> Unpin for PushFuture<'_, A> {}
148impl<A: AsyncProducer> FusedFuture for PushFuture<'_, A> {
149    fn is_terminated(&self) -> bool {
150        self.item.is_none()
151    }
152}
153impl<A: AsyncProducer> Future for PushFuture<'_, A> {
154    type Output = Result<(), A::Item>;
155
156    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
157        let mut waker_registered = false;
158        loop {
159            let item = self.item.take().unwrap();
160            if self.owner.is_closed() {
161                break Poll::Ready(Err(item));
162            }
163            let push_result = self.owner.try_push(item);
164            if push_result.is_ok() {
165                break Poll::Ready(Ok(()));
166            }
167            self.item.replace(push_result.unwrap_err());
168            if waker_registered {
169                break Poll::Pending;
170            }
171            self.owner.register_waker(cx.waker());
172            waker_registered = true;
173        }
174    }
175}
176
177/// # Cancel safety
178///
179/// On cancel the slice can be copied partially.
180#[must_use = "futures do nothing unless you `.await` or poll them"]
181pub struct PushSliceFuture<'a, 'b, A: AsyncProducer + ?Sized>
182where
183    A::Item: Copy,
184{
185    owner: &'a mut A,
186    slice: Option<&'b [A::Item]>,
187    count: usize,
188}
189impl<A: AsyncProducer> Unpin for PushSliceFuture<'_, '_, A> where A::Item: Copy {}
190impl<A: AsyncProducer> FusedFuture for PushSliceFuture<'_, '_, A>
191where
192    A::Item: Copy,
193{
194    fn is_terminated(&self) -> bool {
195        self.slice.is_none()
196    }
197}
198impl<A: AsyncProducer> Future for PushSliceFuture<'_, '_, A>
199where
200    A::Item: Copy,
201{
202    type Output = Result<(), usize>;
203
204    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
205        let mut waker_registered = false;
206        loop {
207            let mut slice = self.slice.take().unwrap();
208            if self.owner.is_closed() {
209                break Poll::Ready(Err(self.count));
210            }
211            let len = self.owner.push_slice(slice);
212            slice = &slice[len..];
213            self.count += len;
214            if slice.is_empty() {
215                break Poll::Ready(Ok(()));
216            }
217            self.slice.replace(slice);
218            if waker_registered {
219                break Poll::Pending;
220            }
221            self.owner.register_waker(cx.waker());
222            waker_registered = true;
223        }
224    }
225}
226impl<A: AsyncProducer> PushSliceFuture<'_, '_, A>
227where
228    A::Item: Copy,
229{
230    pub fn count(&self) -> usize {
231        self.count
232    }
233}
234
235/// # Cancel safety
236///
237/// If future is cancelled then remaining items are left in iterator.
238#[must_use = "futures do nothing unless you `.await` or poll them"]
239pub struct PushIterFuture<'a, A: AsyncProducer + ?Sized, I: Iterator<Item = A::Item>> {
240    owner: &'a mut A,
241    iter: Option<Peekable<I>>,
242}
243impl<A: AsyncProducer, I: Iterator<Item = A::Item>> Unpin for PushIterFuture<'_, A, I> {}
244impl<A: AsyncProducer, I: Iterator<Item = A::Item>> FusedFuture for PushIterFuture<'_, A, I> {
245    fn is_terminated(&self) -> bool {
246        self.iter.is_none()
247    }
248}
249impl<A: AsyncProducer, I: Iterator<Item = A::Item>> Future for PushIterFuture<'_, A, I> {
250    type Output = bool;
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 mut iter = self.iter.take().unwrap();
256            if self.owner.is_closed() {
257                break Poll::Ready(false);
258            }
259            self.owner.push_iter(&mut iter);
260            if iter.peek().is_none() {
261                break Poll::Ready(true);
262            }
263            self.iter.replace(iter);
264            if waker_registered {
265                break Poll::Pending;
266            }
267            self.owner.register_waker(cx.waker());
268            waker_registered = true;
269        }
270    }
271}
272impl<A: AsyncProducer, I: Iterator<Item = A::Item>> PushIterFuture<'_, A, I> {
273    pub fn inner(&self) -> &Peekable<I> {
274        self.iter.as_ref().unwrap()
275    }
276    pub fn inner_mut(&mut self) -> &mut Peekable<I> {
277        self.iter.as_mut().unwrap()
278    }
279    pub fn into_inner(self) -> Peekable<I> {
280        self.iter.unwrap()
281    }
282}
283
284/// # Cancel safety
285///
286/// You can safely cancel this future.
287#[must_use = "futures do nothing unless you `.await` or poll them"]
288pub struct WaitVacantFuture<'a, A: AsyncProducer + ?Sized> {
289    owner: &'a A,
290    count: usize,
291    done: bool,
292}
293impl<A: AsyncProducer> Unpin for WaitVacantFuture<'_, A> {}
294impl<A: AsyncProducer> FusedFuture for WaitVacantFuture<'_, A> {
295    fn is_terminated(&self) -> bool {
296        self.done
297    }
298}
299impl<A: AsyncProducer> Future for WaitVacantFuture<'_, A> {
300    type Output = ();
301
302    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
303        let mut waker_registered = false;
304        loop {
305            assert!(!self.done);
306            let closed = self.owner.is_closed();
307            if self.count <= self.owner.vacant_len() || closed {
308                self.done = true;
309                break Poll::Ready(());
310            }
311            if waker_registered {
312                break Poll::Pending;
313            }
314            self.owner.register_waker(cx.waker());
315            waker_registered = true;
316        }
317    }
318}