1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
use crate::consumer::Status;
use crate::prelude::*;
/// A wrapper for a [`Consumer`] which tracks the [`Status`] of the wrapped consumer, and the number of [`items`](Consumer::Item) consumed since the wrapper was created.
///
/// Use the [`into_inner`](Stats::into_inner) method to consume the wrapper and access the wrapped consumer.
///
/// Created via the [`ConsumerExt::to_stats`] method.
///
/// <br/>Counterpart: the [`producer::Stats`] type.
pub struct Stats<C> {
inner: C,
count: usize,
status: Status,
}
impl<C> Stats<C> {
pub(crate) fn new(inner: C) -> Self {
Stats {
inner,
count: 0,
status: Status::Processing,
}
}
/// Returns the number of [`items`](Consumer::Item) which have been consumed since since this wrapper was created.
pub fn count(&self) -> usize {
self.count
}
/// Returns the [`Status`] of the [`Consumer`].
pub fn status(&self) -> Status {
self.status
}
/// Consumes this wrapper to return the inner [`Consumer`].
pub fn into_inner(self) -> C {
self.inner
}
}
impl<C> Consumer for Stats<C>
where
C: Consumer,
{
type Item = C::Item;
type Final = C::Final;
type Error = C::Error;
async fn consume(&mut self, val: Either<Self::Item, Self::Final>) -> Result<(), Self::Error> {
let finalised = val.is_right();
let result = self.inner.consume(val).await;
match result {
Ok(()) => {
if finalised {
self.status = Status::Finalised;
} else {
self.count = self
.count
.checked_add(1)
.expect("consumer stats cannot count more than usize::MAX items");
}
}
Err(_) => {
self.status = Status::Errored;
}
}
result
}
async fn flush(&mut self) -> Result<(), Self::Error> {
let result = self.inner.flush().await;
if result.is_err() {
self.status = Status::Errored;
}
result
}
}
impl<C> BulkConsumer for Stats<C>
where
C: BulkConsumer,
{
async fn expose_slots_gracefully<F, R>(&mut self, f: F) -> Result<R, (F, Self::Error)>
where
F: AsyncFnOnce(&mut [Self::Item]) -> (usize, R),
{
let mut f = Some(f);
let result = match self
.inner
.expose_slots_gracefully(async |items| {
let f = f.take().expect("constructed as a Some variant");
let (consumed, result) = f(items).await;
self.count = self
.count
.checked_add(consumed)
.expect("consumer stats cannot count more than usize::MAX items");
(consumed, result)
})
.await
{
Ok(result) => Ok(result),
Err((_, err)) => {
self.status = Status::Errored;
Err((
f.take()
.expect("provided closure must not be called when an error occurs"),
err,
))
}
};
result
}
}