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        let mut waker_registered = false;
119        loop {
120            if self.is_closed() {
121                break Poll::Ready(Ok(0));
122            }
123            let count = self.push_slice(buf);
124            if count > 0 {
125                break Poll::Ready(Ok(count));
126            }
127            if waker_registered {
128                break Poll::Pending;
129            }
130            self.register_waker(cx.waker());
131            waker_registered = true;
132        }
133    }
134}
135
136/// # Cancel safety
137///
138/// If future is cancelled no item pushed to the RB.
139#[must_use = "futures do nothing unless you `.await` or poll them"]
140pub struct PushFuture<'a, A: AsyncProducer + ?Sized> {
141    owner: &'a mut A,
142    item: Option<A::Item>,
143}
144impl<A: AsyncProducer> Unpin for PushFuture<'_, A> {}
145impl<A: AsyncProducer> FusedFuture for PushFuture<'_, A> {
146    fn is_terminated(&self) -> bool {
147        self.item.is_none()
148    }
149}
150impl<A: AsyncProducer> Future for PushFuture<'_, A> {
151    type Output = Result<(), A::Item>;
152
153    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
154        let mut waker_registered = false;
155        loop {
156            let item = self.item.take().unwrap();
157            if self.owner.is_closed() {
158                break Poll::Ready(Err(item));
159            }
160            let push_result = self.owner.try_push(item);
161            if push_result.is_ok() {
162                break Poll::Ready(Ok(()));
163            }
164            self.item.replace(push_result.unwrap_err());
165            if waker_registered {
166                break Poll::Pending;
167            }
168            self.owner.register_waker(cx.waker());
169            waker_registered = true;
170        }
171    }
172}
173
174/// # Cancel safety
175///
176/// On cancel the slice can be copied partially.
177#[must_use = "futures do nothing unless you `.await` or poll them"]
178pub struct PushSliceFuture<'a, 'b, A: AsyncProducer + ?Sized>
179where
180    A::Item: Copy,
181{
182    owner: &'a mut A,
183    slice: Option<&'b [A::Item]>,
184    count: usize,
185}
186impl<A: AsyncProducer> Unpin for PushSliceFuture<'_, '_, A> where A::Item: Copy {}
187impl<A: AsyncProducer> FusedFuture for PushSliceFuture<'_, '_, A>
188where
189    A::Item: Copy,
190{
191    fn is_terminated(&self) -> bool {
192        self.slice.is_none()
193    }
194}
195impl<A: AsyncProducer> Future for PushSliceFuture<'_, '_, A>
196where
197    A::Item: Copy,
198{
199    type Output = Result<(), usize>;
200
201    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
202        let mut waker_registered = false;
203        loop {
204            let mut slice = self.slice.take().unwrap();
205            if self.owner.is_closed() {
206                break Poll::Ready(Err(self.count));
207            }
208            let len = self.owner.push_slice(slice);
209            slice = &slice[len..];
210            self.count += len;
211            if slice.is_empty() {
212                break Poll::Ready(Ok(()));
213            }
214            self.slice.replace(slice);
215            if waker_registered {
216                break Poll::Pending;
217            }
218            self.owner.register_waker(cx.waker());
219            waker_registered = true;
220        }
221    }
222}
223impl<A: AsyncProducer> PushSliceFuture<'_, '_, A>
224where
225    A::Item: Copy,
226{
227    pub fn count(&self) -> usize {
228        self.count
229    }
230}
231
232/// # Cancel safety
233///
234/// If future is cancelled then remaining items are left in iterator.
235#[must_use = "futures do nothing unless you `.await` or poll them"]
236pub struct PushIterFuture<'a, A: AsyncProducer + ?Sized, I: Iterator<Item = A::Item>> {
237    owner: &'a mut A,
238    iter: Option<Peekable<I>>,
239}
240impl<A: AsyncProducer, I: Iterator<Item = A::Item>> Unpin for PushIterFuture<'_, A, I> {}
241impl<A: AsyncProducer, I: Iterator<Item = A::Item>> FusedFuture for PushIterFuture<'_, A, I> {
242    fn is_terminated(&self) -> bool {
243        self.iter.is_none() || self.owner.is_closed()
244    }
245}
246impl<A: AsyncProducer, I: Iterator<Item = A::Item>> Future for PushIterFuture<'_, A, I> {
247    type Output = bool;
248
249    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
250        let mut waker_registered = false;
251        loop {
252            let mut iter = self.iter.take().unwrap();
253            if self.owner.is_closed() {
254                break Poll::Ready(false);
255            }
256            self.owner.push_iter(&mut iter);
257            if iter.peek().is_none() {
258                break Poll::Ready(true);
259            }
260            self.iter.replace(iter);
261            if waker_registered {
262                break Poll::Pending;
263            }
264            self.owner.register_waker(cx.waker());
265            waker_registered = true;
266        }
267    }
268}
269impl<A: AsyncProducer, I: Iterator<Item = A::Item>> PushIterFuture<'_, A, I> {
270    pub fn inner(&self) -> &Peekable<I> {
271        self.iter.as_ref().unwrap()
272    }
273    pub fn inner_mut(&mut self) -> &mut Peekable<I> {
274        self.iter.as_mut().unwrap()
275    }
276    pub fn into_inner(self) -> Peekable<I> {
277        self.iter.unwrap()
278    }
279}
280
281/// # Cancel safety
282///
283/// You can safely cancel this future.
284#[must_use = "futures do nothing unless you `.await` or poll them"]
285pub struct WaitVacantFuture<'a, A: AsyncProducer + ?Sized> {
286    owner: &'a A,
287    count: usize,
288    done: bool,
289}
290impl<A: AsyncProducer> Unpin for WaitVacantFuture<'_, A> {}
291impl<A: AsyncProducer> FusedFuture for WaitVacantFuture<'_, A> {
292    fn is_terminated(&self) -> bool {
293        self.done
294    }
295}
296impl<A: AsyncProducer> Future for WaitVacantFuture<'_, A> {
297    type Output = ();
298
299    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
300        let mut waker_registered = false;
301        loop {
302            assert!(!self.done);
303            let closed = self.owner.is_closed();
304            if self.count <= self.owner.vacant_len() || closed {
305                break Poll::Ready(());
306            }
307            if waker_registered {
308                break Poll::Pending;
309            }
310            self.owner.register_waker(cx.waker());
311            waker_registered = true;
312        }
313    }
314}