Coalesce communication channel notifications

communication/src/allocator/counters.rs

@@ -9,7 +9,7 @@ use crate::{Push, Pull};
 /// The push half of an intra-thread channel.
 pub struct Pusher<T, P: Push<T>> {
     index: usize,
-    // count: usize,
+    pushed: usize,
     events: Rc<RefCell<Vec<usize>>>,
     pusher: P,
     phantom: ::std::marker::PhantomData<T>,
@@ -20,7 +20,7 @@ impl<T, P: Push<T>>  Pusher<T, P> {
     pub fn new(pusher: P, index: usize, events: Rc<RefCell<Vec<usize>>>) -> Self {
         Pusher {
             index,
-            // count: 0,
+            pushed: 0,
             events,
             pusher,
             phantom: ::std::marker::PhantomData,
@@ -31,31 +31,32 @@ impl<T, P: Push<T>>  Pusher<T, P> {
 impl<T, P: Push<T>> Push<T> for Pusher<T, P> {
     #[inline]
     fn push(&mut self, element: &mut Option<T>) {
-        // if element.is_none() {
-        //     if self.count != 0 {
-        //         self.events
-        //             .borrow_mut()
-        //             .push_back(self.index);
-        //         self.count = 0;
-        //     }
-        // }
-        // else {
-        //     self.count += 1;
-        // }
-        // TODO: Version above is less chatty, but can be a bit late in
-        //       moving information along. Better, but needs cooperation.
-        self.events
-            .borrow_mut()
-            .push(self.index);
+        let done = element.is_none();
+        self.pusher.push(element);
 
-        self.pusher.push(element)
+        // An empty batch emits no event; a single message emits one prompt
+        // event; multi-message batches emit a second event when the batch ends.
+        // Notify promptly for the first message, and once more at the end
+        // if later messages may have arrived after the receiver drained.
+        if done {
+            if self.pushed > 1 {
+                self.events.borrow_mut().push(self.index);
+            }
+            self.pushed = 0;
+        }
+        else {
+            if self.pushed == 0 {
+                self.events.borrow_mut().push(self.index);
+            }
+            self.pushed = self.pushed.saturating_add(1);
+        }
     }
 }
 
-/// The push half of an intra-thread channel.
+/// The push half of an inter-thread channel.
 pub struct ArcPusher<T, P: Push<T>> {
     index: usize,
-    // count: usize,
+    pushed: usize,
     events: Sender<usize>,
     pusher: P,
     phantom: ::std::marker::PhantomData<T>,
@@ -67,7 +68,7 @@ impl<T, P: Push<T>>  ArcPusher<T, P> {
     pub fn new(pusher: P, index: usize, events: Sender<usize>, buzzer: crate::buzzer::Buzzer) -> Self {
         ArcPusher {
             index,
-            // count: 0,
+            pushed: 0,
             events,
             pusher,
             phantom: ::std::marker::PhantomData,
@@ -79,27 +80,29 @@ impl<T, P: Push<T>>  ArcPusher<T, P> {
 impl<T, P: Push<T>> Push<T> for ArcPusher<T, P> {
     #[inline]
     fn push(&mut self, element: &mut Option<T>) {
-        // if element.is_none() {
-        //     if self.count != 0 {
-        //         self.events
-        //             .send((self.index, Event::Pushed(self.count)))
-        //             .expect("Failed to send message count");
-        //         self.count = 0;
-        //     }
-        // }
-        // else {
-        //     self.count += 1;
-        // }
+        let done = element.is_none();
+        self.pusher.push(element);
 
+        // An empty batch emits no event; a single message emits one prompt
+        // event; multi-message batches emit a second event when the batch ends.
         // These three calls should happen in this order, to ensure that
         // we first enqueue data, second enqueue interest in the channel,
         // and finally awaken the thread. Other orders are defective when
         // multiple threads are involved.
-        self.pusher.push(element);
-        let _ = self.events.send(self.index);
-            // TODO : Perhaps this shouldn't be a fatal error (e.g. in shutdown).
-            // .expect("Failed to send message count");
-        self.buzzer.buzz();
+        if done {
+            if self.pushed > 1 {
+                let _ = self.events.send(self.index);
+                self.buzzer.buzz();
+            }
+            self.pushed = 0;
+        }
+        else {
+            if self.pushed == 0 {
+                let _ = self.events.send(self.index);
+                self.buzzer.buzz();
+            }
+            self.pushed = self.pushed.saturating_add(1);
+        }
     }
 }
 

communication/src/lib.rs

@@ -129,7 +129,8 @@ pub trait Bytesable {
 ///
 /// Conventionally, a sequence of calls to `push()` should conclude with
 /// a call of `push(&mut None)` or `done()` to signal to implementors that
-/// another call to `push()` may not be coming.
+/// another call to `push()` may not be coming. Implementors may coalesce
+/// notifications for subsequent messages until they observe this boundary.
 pub trait Push<T> {
     /// Pushes `element` with the opportunity to take ownership.
     fn push(&mut self, element: &mut Option<T>);

communication/tests/counters.rs

@@ -0,0 +1,113 @@
+use std::cell::RefCell;
+use std::rc::Rc;
+use std::sync::mpsc::{channel, TryRecvError};
+use std::time::Duration;
+
+use timely_communication::allocator::counters::{ArcPusher, Pusher};
+use timely_communication::Push;
+
+struct RecordingPusher<T> {
+    pushed: Rc<RefCell<Vec<Option<T>>>>,
+}
+
+impl<T> Push<T> for RecordingPusher<T> {
+    fn push(&mut self, element: &mut Option<T>) {
+        self.pushed.borrow_mut().push(element.take());
+    }
+}
+
+struct ChannelPusher<T>(std::sync::mpsc::Sender<T>);
+
+impl<T> Push<T> for ChannelPusher<T> {
+    fn push(&mut self, element: &mut Option<T>) {
+        if let Some(element) = element.take() {
+            let _ = self.0.send(element);
+        }
+    }
+}
+
+#[test]
+fn pusher_coalesces_non_empty_batches() {
+    let pushed = Rc::new(RefCell::new(Vec::new()));
+    let events = Rc::new(RefCell::new(Vec::new()));
+    let mut pusher = Pusher::new(
+        RecordingPusher {
+            pushed: Rc::clone(&pushed),
+        },
+        7,
+        Rc::clone(&events),
+    );
+
+    pusher.send(1);
+    pusher.send(2);
+    assert_eq!(&*events.borrow(), &[7]);
+
+    pusher.done();
+    assert_eq!(&*events.borrow(), &[7, 7]);
+    assert_eq!(&*pushed.borrow(), &[Some(1), Some(2), None]);
+
+    pusher.done();
+    assert_eq!(&*events.borrow(), &[7, 7]);
+
+    pusher.send(3);
+    assert_eq!(&*events.borrow(), &[7, 7, 7]);
+    pusher.done();
+    assert_eq!(&*events.borrow(), &[7, 7, 7]);
+}
+
+#[test]
+fn arc_pusher_coalesces_non_empty_batches() {
+    let pushed = Rc::new(RefCell::new(Vec::new()));
+    let (events_tx, events_rx) = channel();
+    let mut pusher = ArcPusher::new(
+        RecordingPusher {
+            pushed: Rc::clone(&pushed),
+        },
+        11,
+        events_tx,
+        timely_communication::buzzer::Buzzer::default(),
+    );
+
+    pusher.send(1);
+    assert_eq!(events_rx.recv().unwrap(), 11);
+    pusher.send(2);
+    assert_eq!(events_rx.try_recv(), Err(TryRecvError::Empty));
+
+    pusher.done();
+    assert_eq!(events_rx.recv().unwrap(), 11);
+    assert_eq!(events_rx.try_recv(), Err(TryRecvError::Empty));
+    assert_eq!(&*pushed.borrow(), &[Some(1), Some(2), None]);
+
+    pusher.done();
+    assert_eq!(events_rx.try_recv(), Err(TryRecvError::Empty));
+
+    pusher.send(3);
+    assert_eq!(events_rx.recv().unwrap(), 11);
+    pusher.done();
+    assert_eq!(events_rx.try_recv(), Err(TryRecvError::Empty));
+}
+
+#[test]
+fn arc_pusher_notifies_messages_arriving_after_the_first_wake() {
+    let timeout = Duration::from_secs(5);
+    let (data_tx, data_rx) = channel();
+    let (events_tx, events_rx) = channel();
+    let (continue_tx, continue_rx) = channel();
+    let buzzer = timely_communication::buzzer::Buzzer::default();
+
+    let sender = std::thread::spawn(move || {
+        let mut pusher = ArcPusher::new(ChannelPusher(data_tx), 13, events_tx, buzzer);
+        pusher.send(1);
+        continue_rx.recv_timeout(timeout).unwrap();
+        pusher.send(2);
+        pusher.done();
+    });
+
+    assert_eq!(events_rx.recv_timeout(timeout).unwrap(), 13);
+    assert_eq!(data_rx.recv_timeout(timeout).unwrap(), 1);
+    continue_tx.send(()).unwrap();
+    assert_eq!(events_rx.recv_timeout(timeout).unwrap(), 13);
+    assert_eq!(data_rx.recv_timeout(timeout).unwrap(), 2);
+
+    sender.join().unwrap();
+}

timely/src/dataflow/operators/core/input.rs

@@ -368,9 +368,14 @@ impl<T: Timestamp, CB: ContainerBuilder<Container: Clone>> Handle<T, CB> {
         }
     }
 
-    /// Flush all contents and distribute to downstream operators.
+    /// Flush all contents, distribute them downstream, and close the current batch.
     #[inline]
     pub fn flush(&mut self) {
+        self.flush_builder();
+        self.flush_pushers();
+    }
+
+    fn flush_builder(&mut self) {
         while let Some(container) = self.builder.finish() {
             Self::send_container(container, &mut self.buffer, &mut self.pushers, &self.now_at);
         }
@@ -402,9 +407,6 @@ impl<T: Timestamp, CB: ContainerBuilder<Container: Clone>> Handle<T, CB> {
     // TODO: Find a better name for this function.
     fn close_epoch(&mut self) {
         self.flush();
-        for pusher in self.pushers.iter_mut() {
-            pusher.done();
-        }
         for progress in self.progress.iter() {
             progress.borrow_mut().update(self.now_at.clone(), -1);
         }
@@ -416,7 +418,8 @@ impl<T: Timestamp, CB: ContainerBuilder<Container: Clone>> Handle<T, CB> {
 
     /// Sends a batch of records into the corresponding timely dataflow [Stream], at the current epoch.
     ///
-    /// This method flushes single elements previously sent with `send`, to keep the insertion order.
+    /// This method flushes single elements previously sent with `send`, to keep the insertion order,
+    /// and closes the batch after sending it.
     ///
     /// # Examples
     /// ```
@@ -445,8 +448,15 @@ impl<T: Timestamp, CB: ContainerBuilder<Container: Clone>> Handle<T, CB> {
     pub fn send_batch(&mut self, buffer: &mut CB::Container) {
         if !buffer.is_empty() {
             // flush buffered elements to ensure local fifo.
-            self.flush();
+            self.flush_builder();
             Self::send_container(buffer, &mut self.buffer, &mut self.pushers, &self.now_at);
+            self.flush_pushers();
+        }
+    }
+
+    fn flush_pushers(&mut self) {
+        for pusher in self.pushers.iter_mut() {
+            pusher.done();
         }
     }
 

timely/tests/coalesced_input.rs

@@ -0,0 +1,38 @@
+use std::cell::RefCell;
+use std::rc::Rc;
+
+use timely::dataflow::operators::{Input, Inspect};
+
+#[test]
+fn input_can_send_again_at_the_same_epoch_after_idle() {
+    timely::execute_directly(|worker| {
+        let seen = Rc::new(RefCell::new(Vec::new()));
+        let output = Rc::clone(&seen);
+        let (mut input, ()) = worker.dataflow::<u64, _, _>(|scope| {
+            let (input, stream) = scope.new_input::<Vec<u64>>();
+            stream.inspect(move |item| output.borrow_mut().push(*item));
+            (input, ())
+        });
+
+        input.send_batch(&mut vec![1]);
+        worker.step();
+        worker.step();
+        assert_eq!(&*seen.borrow(), &[1]);
+
+        input.send_batch(&mut vec![2]);
+        worker.step();
+        assert_eq!(&*seen.borrow(), &[1, 2]);
+
+        worker.step();
+        input.send(3);
+        input.flush();
+        worker.step();
+        assert_eq!(&*seen.borrow(), &[1, 2, 3]);
+
+        worker.step();
+        input.send(4);
+        input.flush();
+        worker.step();
+        assert_eq!(&*seen.borrow(), &[1, 2, 3, 4]);
+    });
+}