ufotofu/consumer/
limit.rs

1use crate::{prelude::*, ConsumeLimitError};
2
3/// A `Consumer` adaptor that limits the number of regular items accepted by the Consumer.
4///
5/// Use the [`into_inner`](`Limit::into_inner`) method or [`AsRef<C>`](AsRef) impl to access the wrapped consumer.
6///
7/// Created via the [`ConsumerExt::to_limit`] method.
8///
9/// <br/>Counterpart: the [`producer::Limit`] type.
10#[derive(Copy, Clone, Hash, Ord, Eq, PartialEq, PartialOrd, Debug)]
11pub struct Limit<C> {
12    inner: C,
13    remaining: usize,
14}
15
16impl<C> Limit<C> {
17    pub(crate) fn new(inner: C, limit: usize) -> Self {
18        Limit {
19            inner,
20            remaining: limit,
21        }
22    }
23
24    /// Consumes `self` and returns the wrapped consumer.
25    pub fn into_inner(self) -> C {
26        self.inner
27    }
28
29    /// Returns the number of items that can still be consumed before returning an error.
30    pub fn remaining(&self) -> usize {
31        self.remaining
32    }
33}
34
35impl<C> AsRef<C> for Limit<C> {
36    fn as_ref(&self) -> &C {
37        &self.inner
38    }
39}
40
41impl<C: Consumer> Consumer for Limit<C> {
42    type Item = C::Item;
43    type Final = C::Final;
44    type Error = ConsumeLimitError<C::Error>;
45
46    async fn consume(&mut self, val: Either<Self::Item, Self::Final>) -> Result<(), Self::Error> {
47        match val {
48            Left(item) => match self.remaining.checked_sub(1) {
49                None => Result::Err(ConsumeLimitError::LimitReached),
50                Some(decremented) => {
51                    self.remaining = decremented;
52                    self.inner
53                        .consume_item(item)
54                        .await
55                        .map_err(ConsumeLimitError::Inner)
56                }
57            },
58            Right(fin) => self
59                .inner
60                .consume_final(fin)
61                .await
62                .map_err(ConsumeLimitError::Inner),
63        }
64    }
65
66    async fn flush(&mut self) -> Result<(), Self::Error> {
67        self.inner.flush().await.map_err(ConsumeLimitError::Inner)
68    }
69}
70
71impl<C: BulkConsumer> BulkConsumer for Limit<C> {
72    async fn expose_slots_gracefully<F, R>(&mut self, f: F) -> Result<R, (F, Self::Error)>
73    where
74        F: AsyncFnOnce(&mut [Self::Item]) -> (usize, R),
75    {
76        if self.remaining == 0 {
77            return Err((f, ConsumeLimitError::LimitReached));
78        }
79
80        let mut f = Some(f);
81        let mut remaining = self.remaining;
82
83        let ret = self
84            .inner
85            .expose_slots(async |slots| {
86                let len = remaining.min(slots.len());
87                let inner_f = f.take().expect("This only runs once, so f is safe to take");
88                let (consumed, returned) = inner_f(&mut slots[..len]).await;
89                remaining -= consumed;
90                (consumed, returned)
91            })
92            .await;
93        self.remaining = remaining;
94
95        match ret {
96            Ok(r) => Ok(r),
97            Err(error) => Err((
98                f.take()
99                    .expect("Wrapper closure wasn't called, so f is safe to take"),
100                ConsumeLimitError::Inner(error),
101            )),
102        }
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use crate::prelude::*;
109    use crate::ConsumeLimitError;
110
111    #[test]
112    fn inner_error() {
113        let mut c = consumer::error_immediately("internal_error").to_limit(20);
114        pollster::block_on(async {
115            assert_eq!(
116                c.flush().await,
117                Err(ConsumeLimitError::Inner("internal_error"))
118            );
119        });
120    }
121}
ufotofu/consumer/limit.rs

ufotofu/consumer/
limit.rs
