diff --git a/timely/src/dataflow/channels/pushers/counter.rs b/timely/src/dataflow/channels/pushers/counter.rs index 60064af22..aa8a416d5 100644 --- a/timely/src/dataflow/channels/pushers/counter.rs +++ b/timely/src/dataflow/channels/pushers/counter.rs @@ -13,6 +13,8 @@ use crate::Accountable; pub struct Counter { pushee: P, produced: Rc>>, + /// Set on message push, unset on flush. + dirty: bool, } impl Push> for Counter where P: Push> { @@ -23,10 +25,11 @@ impl Push> for Counter wher for time in message.stamp.iter() { produced.update(time.clone(), message.data.record_count()); } + self.dirty = true; } - // only propagate `None` if dirty (indicates flush) - if message.is_some() || !self.produced.borrow_mut().is_empty() { + // only propagate `None` if a message has been pushed. + if message.is_some() || std::mem::take(&mut self.dirty) { self.pushee.push(message); } } @@ -38,6 +41,7 @@ impl Counter where T : Ord+Clone+'static { Counter { pushee, produced: Rc::new(RefCell::new(ChangeBatch::new())), + dirty: false, } } /// A references to shared changes in counts, for cloning or draining. diff --git a/timely/tests/gh_824.rs b/timely/tests/gh_824.rs new file mode 100644 index 000000000..6bc04a221 --- /dev/null +++ b/timely/tests/gh_824.rs @@ -0,0 +1,38 @@ +use timely::dataflow::operators::{Exchange, Input, Probe}; +use timely::dataflow::InputHandle; +use timely::Config; + +/// A record flushed into an exchange, followed by a step that drains the input's +/// produced counts, must still be delivered when the input later advances. +#[test] +fn gh_824() { + timely::execute(Config::process(2), |worker| { + let mut input = InputHandle::new(); + let probe = worker.dataflow::(|scope| { + scope + .input_from(&mut input) + .container::>() + .exchange(|x: &u64| *x) + .probe() + .0 + }); + + if worker.index() == 0 { + input.advance_to(1); + input.send(0); + input.flush(); + worker.step(); + } + input.advance_to(2); + + while probe.less_than(&2) { + worker.step(); + } + + input.close(); + while !probe.done() { + worker.step(); + } + }) + .unwrap(); +}