timely/dataflow/channels/pushers/
counter.rs

1//! A wrapper which counts the number of records pushed past and updates a shared count map.
2
3use std::rc::Rc;
4use std::cell::RefCell;
5
6use crate::progress::ChangeBatch;
7use crate::dataflow::channels::Message;
8use crate::communication::Push;
9use crate::Accountable;
10
11/// A wrapper which updates shared `produced` based on the number of records pushed.
12#[derive(Debug)]
13pub struct Counter<T, P> {
14    pushee: P,
15    produced: Rc<RefCell<ChangeBatch<T>>>,
16    /// Set on message push, unset on flush.
17    dirty: bool,
18}
19
20impl<T: Clone+Ord, C: Accountable, P> Push<Message<T, C>> for Counter<T, P> where P: Push<Message<T, C>> {
21    #[inline]
22    fn push(&mut self, message: &mut Option<Message<T, C>>) {
23        if let Some(message) = message {
24            let mut produced = self.produced.borrow_mut();
25            for time in message.stamp.iter() {
26                produced.update(time.clone(), message.data.record_count());
27            }
28            self.dirty = true;
29        }
30
31        // only propagate `None` if a message has been pushed.
32        if message.is_some() || std::mem::take(&mut self.dirty) {
33            self.pushee.push(message);
34        }
35    }
36}
37
38impl<T, P> Counter<T, P> where T : Ord+Clone+'static {
39    /// Allocates a new `Counter` from a pushee and shared counts.
40    pub fn new(pushee: P) -> Counter<T, P> {
41        Counter {
42            pushee,
43            produced: Rc::new(RefCell::new(ChangeBatch::new())),
44            dirty: false,
45        }
46    }
47    /// A references to shared changes in counts, for cloning or draining.
48    #[inline]
49    pub fn produced(&self) -> &Rc<RefCell<ChangeBatch<T>>> {
50        &self.produced
51    }
52}
timely/dataflow/channels/pushers/counter.rs

timely/dataflow/channels/pushers/
counter.rs
