From 3187b3a2b30aea4465f358b6f0f8d76fd5146bfc Mon Sep 17 00:00:00 2001 From: Kord Boniadi <67992622+kboniadi@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:01:05 -0700 Subject: [PATCH 1/2] Coalesce communication channel notifications Preserve an immediate first-message wake while batching redundant notifications until the pusher's done boundary. Close input batches explicitly and cover same-epoch reactivation and inter-thread tail delivery. --- communication/src/allocator/counters.rs | 75 +++++++------ communication/src/lib.rs | 3 +- communication/tests/counters.rs | 113 ++++++++++++++++++++ timely/src/dataflow/operators/core/input.rs | 22 ++-- timely/tests/coalesced_input.rs | 38 +++++++ 5 files changed, 206 insertions(+), 45 deletions(-) create mode 100644 communication/tests/counters.rs create mode 100644 timely/tests/coalesced_input.rs diff --git a/communication/src/allocator/counters.rs b/communication/src/allocator/counters.rs index 6af216e8f..9861c1816 100644 --- a/communication/src/allocator/counters.rs +++ b/communication/src/allocator/counters.rs @@ -9,7 +9,7 @@ use crate::{Push, Pull}; /// The push half of an intra-thread channel. pub struct Pusher> { index: usize, - // count: usize, + pushed: usize, events: Rc>>, pusher: P, phantom: ::std::marker::PhantomData, @@ -20,7 +20,7 @@ impl> Pusher { pub fn new(pusher: P, index: usize, events: Rc>>) -> Self { Pusher { index, - // count: 0, + pushed: 0, events, pusher, phantom: ::std::marker::PhantomData, @@ -31,31 +31,30 @@ impl> Pusher { impl> Push for Pusher { #[inline] fn push(&mut self, element: &mut Option) { - // 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) + // 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> { index: usize, - // count: usize, + pushed: usize, events: Sender, pusher: P, phantom: ::std::marker::PhantomData, @@ -67,7 +66,7 @@ impl> ArcPusher { pub fn new(pusher: P, index: usize, events: Sender, buzzer: crate::buzzer::Buzzer) -> Self { ArcPusher { index, - // count: 0, + pushed: 0, events, pusher, phantom: ::std::marker::PhantomData, @@ -79,27 +78,27 @@ impl> ArcPusher { impl> Push for ArcPusher { #[inline] fn push(&mut self, element: &mut Option) { - // 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); // 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); + } } } diff --git a/communication/src/lib.rs b/communication/src/lib.rs index 3daac7b43..e88618c58 100644 --- a/communication/src/lib.rs +++ b/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 { /// Pushes `element` with the opportunity to take ownership. fn push(&mut self, element: &mut Option); diff --git a/communication/tests/counters.rs b/communication/tests/counters.rs new file mode 100644 index 000000000..627afd837 --- /dev/null +++ b/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 { + pushed: Rc>>>, +} + +impl Push for RecordingPusher { + fn push(&mut self, element: &mut Option) { + self.pushed.borrow_mut().push(element.take()); + } +} + +struct ChannelPusher(std::sync::mpsc::Sender); + +impl Push for ChannelPusher { + fn push(&mut self, element: &mut Option) { + 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(); +} diff --git a/timely/src/dataflow/operators/core/input.rs b/timely/src/dataflow/operators/core/input.rs index 20a1413aa..5e4f511b4 100644 --- a/timely/src/dataflow/operators/core/input.rs +++ b/timely/src/dataflow/operators/core/input.rs @@ -368,9 +368,14 @@ impl> Handle { } } - /// 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> Handle { // 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> Handle { /// 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> Handle { 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(); } } diff --git a/timely/tests/coalesced_input.rs b/timely/tests/coalesced_input.rs new file mode 100644 index 000000000..6c3a75b04 --- /dev/null +++ b/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::(|scope| { + let (input, stream) = scope.new_input::>(); + 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]); + }); +} From a84793c46981ede10ce0f19a2195d62c0049ef72 Mon Sep 17 00:00:00 2001 From: Kord Boniadi <67992622+kboniadi@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:01:05 -0700 Subject: [PATCH 2/2] Document communication counter events --- communication/src/allocator/counters.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/communication/src/allocator/counters.rs b/communication/src/allocator/counters.rs index 9861c1816..5d6ba1ebf 100644 --- a/communication/src/allocator/counters.rs +++ b/communication/src/allocator/counters.rs @@ -34,6 +34,8 @@ impl> Push for Pusher { 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. // Notify promptly for the first message, and once more at the end // if later messages may have arrived after the receiver drained. if done { @@ -81,6 +83,8 @@ impl> Push for ArcPusher { 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