timely/dataflow/channels/pullers/
counter.rs

1//! A wrapper which accounts records pulled past in a shared count map.
2
3use std::rc::Rc;
4use std::cell::RefCell;
5
6use crate::dataflow::channels::Message;
7use crate::progress::{ChangeBatch, Stamp};
8use crate::communication::Pull;
9use crate::Accountable;
10
11/// A wrapper which accounts records pulled past in a shared count map.
12pub struct Counter<T, C, P> {
13    pullable: P,
14    consumed: Rc<RefCell<ChangeBatch<T>>>,
15    phantom: ::std::marker::PhantomData<C>,
16}
17
18/// A guard type that updates the change batch counts on drop
19pub struct ConsumedGuard<T: Ord + Clone + 'static> {
20    consumed: Rc<RefCell<ChangeBatch<T>>>,
21    stamp: Option<Stamp<T>>,
22    record_count: i64,
23}
24
25impl<T:Ord+Clone+'static> ConsumedGuard<T> {
26    #[inline]
27    pub(crate) fn stamp(&self) -> &Stamp<T> {
28        self.stamp.as_ref().unwrap()
29    }
30}
31
32impl<T:Ord+Clone+'static> Drop for ConsumedGuard<T> {
33    fn drop(&mut self) {
34        // SAFETY: we're in a Drop impl, so this runs at most once
35        let stamp = self.stamp.take().unwrap();
36        let mut consumed = self.consumed.borrow_mut();
37        for time in stamp.iter() {
38            consumed.update(time.clone(), self.record_count);
39        }
40    }
41}
42
43impl<T:Ord+Clone+'static, C: Accountable, P: Pull<Message<T, C>>> Counter<T, C, P> {
44    /// Retrieves the next timestamp and batch of data.
45    #[inline]
46    pub fn next(&mut self) -> Option<&mut Message<T, C>> {
47        self.next_guarded().map(|(_guard, bundle)| bundle)
48    }
49
50    #[inline]
51    pub(crate) fn next_guarded(&mut self) -> Option<(ConsumedGuard<T>, &mut Message<T, C>)> {
52        if let Some(message) = self.pullable.pull() {
53            let guard = ConsumedGuard {
54                consumed: Rc::clone(&self.consumed),
55                stamp: Some(message.stamp.clone()),
56                record_count: message.data.record_count(),
57            };
58            Some((guard, message))
59        }
60        else { None }
61    }
62}
63
64impl<T:Ord+Clone+'static, C, P: Pull<Message<T, C>>> Counter<T, C, P> {
65    /// Allocates a new `Counter` from a boxed puller.
66    pub fn new(pullable: P) -> Self {
67        Counter {
68            phantom: ::std::marker::PhantomData,
69            pullable,
70            consumed: Rc::new(RefCell::new(ChangeBatch::new())),
71        }
72    }
73    /// A references to shared changes in counts, for cloning or draining.
74    pub fn consumed(&self) -> &Rc<RefCell<ChangeBatch<T>>> {
75        &self.consumed
76    }
77}
timely/dataflow/channels/pullers/counter.rs

timely/dataflow/channels/pullers/
counter.rs
