From 54906da143daf40604f5a22a623117ff73c89b72 Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@feldera.com>
Date: Tue, 9 Jun 2026 09:29:43 -0700
Subject: [PATCH 1/8] [dbsp] Simplify dyn_shard_accumulate().

We don't need so many special cases because `dyn_shard()` is cached, for
the case where there is a sharded version already, and a no-op, for the
case where there is only one worker.

Signed-off-by: Ben Pfaff <blp@feldera.com>
---
 .../operator/dynamic/sharded_accumulator.rs   | 69 ++++++++-----------
 1 file changed, 30 insertions(+), 39 deletions(-)

diff --git a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
index af793f00e2e..493b49a443e 100644
--- a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
+++ b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
@@ -57,50 +57,41 @@ where
     where
         B: Batch<Time = ()>,
     {
-        if let Some(sharded) = self.get_sharded_version() {
-            sharded.dyn_accumulate(factories)
-        } else if Runtime::num_workers() == 1 {
-            self.dyn_accumulate(factories)
-        } else if Runtime::with_dev_tweaks(|d| !d.streaming_exchange()) {
-            self.dyn_shard(factories).dyn_accumulate(factories)
-        } else {
+        if !self.has_sharded_version()
+            && Runtime::with_dev_tweaks(|d| d.streaming_exchange())
+            && let Some(runtime) = Runtime::runtime()
+            && runtime.layout().n_workers() > 1
+            && runtime.get_step_size() == StepSize::Microsteps
+        {
             self.circuit()
                 .cache_get_or_insert_with(ShardedAccumulatorId::new(self.stream_id()), || {
-                    let runtime = Runtime::runtime().unwrap();
-                    match runtime.get_step_size() {
-                        StepSize::Microsteps => {
-                            let exchange_id: ExchangeId =
-                                runtime.sequence_next().try_into().unwrap();
-                            let exchange = ShardedAccumulator::<B>::with_runtime(
-                                &runtime,
-                                Runtime::worker_index(),
-                                exchange_id,
-                                factories,
-                            );
-                            let enable_count = exchange.enable_count.clone();
-                            let stream = self
-                                .circuit()
-                                .add_exchange(
-                                    ShardedAccumulatorSender::new(
-                                        Some(Location::caller()),
-                                        exchange.clone(),
-                                    ),
-                                    ShardedAccumulatorReceiver::new(
-                                        Some(Location::caller()),
-                                        exchange,
-                                    ),
-                                    self,
-                                )
-                                .mark_sharded();
-                            Accumulation {
-                                stream,
-                                enable_count,
-                            }
-                        }
-                        StepSize::FullSteps => self.dyn_shard(factories).dyn_accumulate(factories),
+                    let exchange_id: ExchangeId = runtime.sequence_next().try_into().unwrap();
+                    let exchange = ShardedAccumulator::<B>::with_runtime(
+                        &runtime,
+                        Runtime::worker_index(),
+                        exchange_id,
+                        factories,
+                    );
+                    let enable_count = exchange.enable_count.clone();
+                    let stream = self
+                        .circuit()
+                        .add_exchange(
+                            ShardedAccumulatorSender::new(
+                                Some(Location::caller()),
+                                exchange.clone(),
+                            ),
+                            ShardedAccumulatorReceiver::new(Some(Location::caller()), exchange),
+                            self,
+                        )
+                        .mark_sharded();
+                    Accumulation {
+                        stream,
+                        enable_count,
                     }
                 })
                 .clone()
+        } else {
+            self.dyn_shard(factories).dyn_accumulate(factories)
         }
     }
 

From a8ee7a6c6facb31a3fa6acfe7a8c5395d3a7bc86 Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@feldera.com>
Date: Tue, 9 Jun 2026 10:55:31 -0700
Subject: [PATCH 2/8] [dbsp] Implement streaming exchange to a subset of
 workers.

This implements `dyn_shard_workers_accumulate` in a streaming manner by
sharding only to a subset of workers.

Signed-off-by: Ben Pfaff <blp@feldera.com>
---
 .../operator/dynamic/communication/shard.rs   |  18 +-
 .../operator/dynamic/sharded_accumulator.rs   | 167 ++++++++++--------
 crates/dbsp/src/trace.rs                      |  26 ++-
 3 files changed, 129 insertions(+), 82 deletions(-)

diff --git a/crates/dbsp/src/operator/dynamic/communication/shard.rs b/crates/dbsp/src/operator/dynamic/communication/shard.rs
index 44f43d59c54..f5cd70af397 100644
--- a/crates/dbsp/src/operator/dynamic/communication/shard.rs
+++ b/crates/dbsp/src/operator/dynamic/communication/shard.rs
@@ -535,17 +535,25 @@ where
     /// across workers, otherwise this will cause the dataflow to yield
     /// incorrect results
     pub fn mark_sharded(&self) -> Self {
-        self.circuit().cache_insert(
-            ShardId::new((self.stream_id(), all_workers())),
-            self.clone(),
-        );
+        self.mark_sharded_workers(all_workers())
+    }
+
+    pub fn mark_sharded_workers(&self, workers: Range<usize>) -> Self {
+        self.circuit()
+            .cache_insert(ShardId::new((self.stream_id(), workers)), self.clone());
         self.clone()
     }
 
     /// Returns `true` if a sharded version of the current stream exists
     pub fn has_sharded_version(&self) -> bool {
+        self.has_workers_sharded_version(all_workers())
+    }
+
+    /// Returns `true` if a version of the current stream sharded to `workers`
+    /// exists.
+    pub fn has_workers_sharded_version(&self, workers: Range<usize>) -> bool {
         self.circuit()
-            .cache_contains(&ShardId::<C, T>::new((self.stream_id(), all_workers())))
+            .cache_contains(&ShardId::<C, T>::new((self.stream_id(), workers)))
     }
 
     pub fn get_sharded_version(&self) -> Option<Self> {
diff --git a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
index 493b49a443e..5f6f1ee62e1 100644
--- a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
+++ b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
@@ -40,7 +40,7 @@ use crate::{
 
 circuit_cache_key!(local StreamingExchangeCacheId<B: Batch>(ExchangeId => Arc<ShardedAccumulator<B>>));
 
-circuit_cache_key!(ShardedAccumulatorId<C, B: Batch>(StreamId => Accumulation<Stream<C, Option<Spine<B>>>>));
+circuit_cache_key!(ShardedAccumulatorId<C, B: Batch>((StreamId, Range<usize>) => Accumulation<Stream<C, Option<Spine<B>>>>));
 
 impl<C, B> Stream<C, B>
 where
@@ -57,42 +57,7 @@ where
     where
         B: Batch<Time = ()>,
     {
-        if !self.has_sharded_version()
-            && Runtime::with_dev_tweaks(|d| d.streaming_exchange())
-            && let Some(runtime) = Runtime::runtime()
-            && runtime.layout().n_workers() > 1
-            && runtime.get_step_size() == StepSize::Microsteps
-        {
-            self.circuit()
-                .cache_get_or_insert_with(ShardedAccumulatorId::new(self.stream_id()), || {
-                    let exchange_id: ExchangeId = runtime.sequence_next().try_into().unwrap();
-                    let exchange = ShardedAccumulator::<B>::with_runtime(
-                        &runtime,
-                        Runtime::worker_index(),
-                        exchange_id,
-                        factories,
-                    );
-                    let enable_count = exchange.enable_count.clone();
-                    let stream = self
-                        .circuit()
-                        .add_exchange(
-                            ShardedAccumulatorSender::new(
-                                Some(Location::caller()),
-                                exchange.clone(),
-                            ),
-                            ShardedAccumulatorReceiver::new(Some(Location::caller()), exchange),
-                            self,
-                        )
-                        .mark_sharded();
-                    Accumulation {
-                        stream,
-                        enable_count,
-                    }
-                })
-                .clone()
-        } else {
-            self.dyn_shard(factories).dyn_accumulate(factories)
-        }
+        self.dyn_shard_workers_accumulate(factories, 0..Runtime::num_workers())
     }
 
     #[track_caller]
@@ -104,13 +69,44 @@ where
     where
         B: Batch<Time = ()>,
     {
-        if Runtime::num_workers() == 1 {
-            self.dyn_accumulate(factories)
-        } else if Runtime::with_dev_tweaks(|d| !d.streaming_exchange()) {
-            self.dyn_shard_workers(workers, factories)
-                .dyn_accumulate(factories)
+        if !self.has_workers_sharded_version(workers.clone())
+            && Runtime::with_dev_tweaks(|d| d.streaming_exchange())
+            && let Some(runtime) = Runtime::runtime()
+            && runtime.layout().n_workers() > 1
+            && runtime.get_step_size() == StepSize::Microsteps
+        {
+            self.circuit()
+                .cache_get_or_insert_with(
+                    ShardedAccumulatorId::new((self.stream_id(), workers.clone())),
+                    || {
+                        let exchange_id: ExchangeId = runtime.sequence_next().try_into().unwrap();
+                        let exchange = ShardedAccumulator::<B>::with_runtime(
+                            &runtime,
+                            Runtime::worker_index(),
+                            workers.clone(),
+                            exchange_id,
+                            factories,
+                        );
+                        let enable_count = exchange.enable_count.clone();
+                        let stream = self
+                            .circuit()
+                            .add_exchange(
+                                ShardedAccumulatorSender::new(
+                                    Some(Location::caller()),
+                                    exchange.clone(),
+                                ),
+                                ShardedAccumulatorReceiver::new(Some(Location::caller()), exchange),
+                                self,
+                            )
+                            .mark_sharded_workers(workers.clone());
+                        Accumulation {
+                            stream,
+                            enable_count,
+                        }
+                    },
+                )
+                .clone()
         } else {
-            // TODO: implement using streaming exchange.
             self.dyn_shard_workers(workers, factories)
                 .dyn_accumulate(factories)
         }
@@ -124,8 +120,11 @@ where
     name: OperatorName,
     exchange_id: ExchangeId,
 
-    /// The number of communicating peers.
-    npeers: usize,
+    /// The destination workers.
+    ///
+    /// Usually this is all workers, but output operators shard to the workers
+    /// on one particular host.
+    workers: Range<usize>,
 
     factories: B::Factories,
 
@@ -147,6 +146,7 @@ where
     fn with_runtime(
         runtime: &Runtime,
         worker_index: usize,
+        workers: Range<usize>,
         exchange_id: ExchangeId,
         factories: &B::Factories,
     ) -> Arc<Self> {
@@ -163,6 +163,7 @@ where
                 ShardedAccumulator::new(
                     runtime,
                     worker_index,
+                    workers,
                     clients,
                     exchange_id,
                     &directory,
@@ -173,10 +174,12 @@ where
             .clone()
     }
 
-    /// Create a new streaming exchange operator for `npeers` communicating threads.
+    /// Create a new streaming exchange operator to shard and accumulate to the
+    /// given `workers` (which is often all the workers in the runtime).
     fn new(
         runtime: &Runtime,
         worker_index: usize,
+        workers: Range<usize>,
         clients: Arc<ExchangeClients>,
         exchange_id: ExchangeId,
         directory: &ExchangeDirectory,
@@ -188,7 +191,7 @@ where
         let name = OperatorName::new("ShardedAccumulatorReceiver");
         let exchange = Arc::new(Self {
             exchange_id,
-            npeers,
+            workers,
             local_workers: layout.local_workers(),
             factories: factories.clone(),
             clients,
@@ -249,7 +252,7 @@ where
         let mut batches = Vec::with_capacity(layout.n_workers());
         shard_batch(
             batch,
-            &(0..self.npeers),
+            &self.workers,
             &mut builders,
             &mut batches,
             &self.factories,
@@ -283,29 +286,40 @@ where
                     }
                 }
                 WorkerLocation::Remote => {
-                    let items = receivers
-                        .clone()
-                        .map(|_| {
-                            let mut fbuf = data
-                                .next()
-                                .expect("data should include one item per peer")
-                                .into_tx()
-                                .expect("remote mailboxes should always be serialized");
-                            fbuf.push(flush as u8);
-                            fbuf
-                        })
-                        .inspect(|serialized| {
-                            serialized_bytes += serialized.len();
-                        })
-                        .collect_vec();
-                    let this = self.clone();
-                    if let Some(waiter) = this.clients.connect(receivers.start).await.send(
-                        name.clone(),
-                        this.exchange_id,
-                        sender,
-                        items,
-                    ) {
-                        remote_waiters.push(waiter);
+                    let mut empty = true;
+                    let mut items = Vec::with_capacity(receivers.len());
+                    for _ in receivers.clone() {
+                        let fbuf = data
+                            .next()
+                            .expect("data should include one item per peer")
+                            .into_tx()
+                            .expect("remote mailboxes should always be serialized");
+                        if !fbuf.is_empty() {
+                            serialized_bytes += fbuf.len();
+                            empty = false;
+                        }
+                        items.push(fbuf);
+                    }
+
+                    // Skip sending to the remote host if there's no data to
+                    // send (unless we're flushing).
+                    //
+                    // The common case for no data to send is when we're
+                    // sharding to workers on only one host for an output
+                    // connector.
+                    if !empty || flush {
+                        for item in &mut items {
+                            item.push(flush as u8);
+                        }
+                        let this = self.clone();
+                        if let Some(waiter) = this.clients.connect(receivers.start).await.send(
+                            name.clone(),
+                            this.exchange_id,
+                            sender,
+                            items,
+                        ) {
+                            remote_waiters.push(waiter);
+                        }
                     }
                 }
             }
@@ -391,7 +405,11 @@ where
     /// Our worker index within the runtime.
     worker_index: usize,
 
-    /// Total number of worker threads in this circuit.
+    /// Total number of worker threads (senders) in this circuit.
+    ///
+    /// This is important because each sender must send, in addition to any
+    /// data, a flush notification to indicate that all data for a step has been
+    /// sent.
     npeers: usize,
 
     /// A deque of spines under construction for delivery to the circuit.  The
@@ -711,6 +729,11 @@ where
         //   then this could cause livelock, spinning uselessly with 100% CPU,
         //   if there's nothing for the circuit to do while data transmits.
         //
+        // `ShardedAccumulatorReceiver` operators in workers that are outside
+        // the set of workers being sharded to seem like a further special case,
+        // since we know in advance that they won't receive any data, but I
+        // don't see a way to do better with them.
+        //
         // I don't know the right solution.
         true
     }
diff --git a/crates/dbsp/src/trace.rs b/crates/dbsp/src/trace.rs
index 1a9e1818644..b97546eff6b 100644
--- a/crates/dbsp/src/trace.rs
+++ b/crates/dbsp/src/trace.rs
@@ -1544,14 +1544,25 @@ impl IndexedWSetSerializer {
         });
     }
 
+    fn is_empty(&self) -> bool {
+        self.n_keys == 0 && self.n_values == 0
+    }
+
+    /// Returns the serialized form as an [FBuf].
+    ///
+    /// If no data was pushed into the serializer, then the [FBuf] will be
+    /// empty.
     pub fn done(mut self, serializer_inner: &mut SerializerInner) -> FBuf {
         #[cfg(debug_assertions)]
         debug_assert_eq!(self.state, State::Key);
-        self.offsets[0] = self.n_keys;
-        self.offsets[1] = self.n_values;
-        serializer_inner.with(FBufSerializer::new(&mut self.fbuf), |s| {
-            s.serialize_value(&self.offsets).unwrap()
-        });
+
+        if !self.is_empty() {
+            self.offsets[0] = self.n_keys;
+            self.offsets[1] = self.n_values;
+            serializer_inner.with(FBufSerializer::new(&mut self.fbuf), |s| {
+                s.serialize_value(&self.offsets).unwrap()
+            });
+        }
         self.fbuf
     }
 }
@@ -1585,6 +1596,11 @@ where
     V: DataTrait + ?Sized,
     R: WeightTrait + ?Sized,
 {
+    if data.is_empty() {
+        // An empty batch is serialized as an empty FBuf.
+        return B::dyn_empty(factories);
+    }
+
     let offsets = unsafe { archived_root::<Vec<usize>>(data) };
     let n_keys = offsets[0] as usize;
     let n_values = offsets[1] as usize;

From 53e856a93700f594f1de4cb160f45b1bad7fb4f7 Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@feldera.com>
Date: Wed, 10 Jun 2026 13:13:36 -0700
Subject: [PATCH 3/8] [dbsp] Turn `ExchangeReceiver` into a "synchronous"
 operator.

DBSP has the notion of an asynchronous operator, which is one that it does
not evaluate until it indicates that it is ready.  Since this notion was
introduced, DBSP was changed to use Rust async for its evaluation
functions.  Rust async is more flexible and easier to use than DBSP
asynchronous operators, so this commit changes ExchangeReceiver to use it
instead.

This is preparation for updating ExchangeReceiver::eval() to work in a
way that doesn't fit as well with the DBSP asynchronous operator concept.

After this commit, there is only one DBSP asynchronous operator left,
which is RebalancingExchangeSender.  It could probably use the same
treatment, and then afterward we could get rid of the DBSP asynchronous
operator concept entirely.

Signed-off-by: Ben Pfaff <blp@feldera.com>
---
 crates/dbsp/src/circuit/runtime.rs            |   3 +-
 .../src/operator/communication/exchange.rs    | 131 ++++--------------
 .../balance/accumulate_trace_balanced.rs      |   1 -
 3 files changed, 28 insertions(+), 107 deletions(-)

diff --git a/crates/dbsp/src/circuit/runtime.rs b/crates/dbsp/src/circuit/runtime.rs
index 8e24494cb96..0b0aec4ec77 100644
--- a/crates/dbsp/src/circuit/runtime.rs
+++ b/crates/dbsp/src/circuit/runtime.rs
@@ -1426,8 +1426,9 @@ where
                             .await;
 
                         exchange
-                            .receive_all(|data| serde_json::from_slice(&data).unwrap())
+                            .receive_all(|data| serde_json::from_slice(&data).unwrap(), None)
                             .await
+                            .0
                     })
                     .await
                     .ok_or(SchedulerError::Killed)
diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs
index b6c2e6d289f..75079d9520d 100644
--- a/crates/dbsp/src/operator/communication/exchange.rs
+++ b/crates/dbsp/src/operator/communication/exchange.rs
@@ -806,39 +806,6 @@ where
         sender * self.npeers + receiver
     }
 
-    /// True if all `receiver`'s incoming mailboxes contain data.
-    ///
-    /// Once this function returns true, a subsequent `try_receive_all`
-    /// operation is guaranteed for `receiver`.
-    fn ready_to_receive(&self, receiver: usize) -> bool {
-        debug_assert!(receiver < self.npeers);
-        self.receiver_counters[receiver].load(Ordering::Acquire) == self.npeers
-    }
-
-    /// Register callback to be invoked whenever the `ready_to_receive`
-    /// condition becomes true.
-    ///
-    /// The callback can be setup at most once (e.g., when a scheduler attaches
-    /// to the circuit) and cannot be unregistered.  Notifications delivered
-    /// before the callback is registered are lost.  The client should call
-    /// `ready_to_receive` after installing the callback to check
-    /// the status.
-    ///
-    /// After the callback has been registered, notifications are delivered with
-    /// at-least-once semantics: a notification is generated whenever the
-    /// status changes from not ready to ready, but spurious notifications
-    /// can occur occasionally.  The user must check the status explicitly
-    /// by calling `ready_to_receive` or be prepared that `receive_all`
-    /// can fail.
-    pub(crate) fn register_receiver_callback<F>(&self, receiver: usize, cb: F)
-    where
-        F: Fn() + Send + Sync + 'static,
-    {
-        debug_assert!(receiver < self.npeers);
-
-        self.receiver_callbacks[receiver].set_callback(cb);
-    }
-
     /// Locks and returns the mailbox for the sender/receiver pair.
     fn mailbox(&self, sender: usize, receiver: usize) -> MutexGuard<'_, Option<Mailbox<T>>> {
         self.mailboxes[self.mailbox_index(sender, receiver)]
@@ -1004,7 +971,14 @@ where
 
     /// Read all incoming messages for this worker, waiting for data to arrive
     /// as needed.
-    pub(crate) async fn receive_all<D>(&self, deserialize: D) -> Vec<T>
+    ///
+    /// When the data is ready, but before reading it, this method swaps
+    /// `start_wait_usecs` with 0 and returns the old value along with the data.
+    pub(crate) async fn receive_all<D>(
+        &self,
+        deserialize: D,
+        start_wait_usecs: Option<&AtomicU64>,
+    ) -> (Vec<T>, Option<u64>)
     where
         D: Fn(AlignedVec) -> T,
     {
@@ -1024,6 +998,11 @@ where
             }
         }
 
+        let start_wait_usecs = start_wait_usecs.and_then(|v| {
+            let start_wait_usecs = v.swap(0, Ordering::Acquire);
+            (start_wait_usecs != 0).then_some(start_wait_usecs)
+        });
+
         let mut data = Vec::with_capacity(self.npeers);
         for sender in 0..self.npeers {
             let mailbox = self.mailbox(sender, receiver).take().unwrap();
@@ -1035,7 +1014,7 @@ where
             }
         }
 
-        data
+        (data, start_wait_usecs)
     }
 }
 
@@ -1100,13 +1079,6 @@ where
 ///                    ExchangeSender  ExchangeReceiver
 /// ```
 ///
-/// `ExchangeSender` is an asynchronous operator., i.e.,
-/// [`ExchangeSender::is_async`] returns `true`.  It becomes schedulable
-/// ([`ExchangeSender::ready`] returns `true`) once all peers have retrieved
-/// values written by the operator in the previous clock cycle.  The scheduler
-/// should use [`ExchangeSender::register_ready_callback`] to get notified when
-/// the operator becomes schedulable.
-///
 /// `ExchangeSender` doesn't have a public constructor and must be instantiated
 /// using the [`new_exchange_operators`] function, which creates an
 /// [`ExchangeSender`]/[`ExchangeReceiver`] pair of operators and connects them
@@ -1375,18 +1347,10 @@ where
 /// peer.
 ///
 /// See [`ExchangeSender`] documentation for details.
-///
-/// `ExchangeReceiver` is an asynchronous operator., i.e.,
-/// [`ExchangeReceiver::is_async`] returns `true`.  It becomes schedulable
-/// ([`ExchangeReceiver::ready`] returns `true`) once all peers have sent values
-/// for this worker in the current clock cycle.  The scheduler should use
-/// [`ExchangeReceiver::register_ready_callback`] to get notified when the
-/// operator becomes schedulable.
 pub struct ExchangeReceiver<IF, T, L, D>
 where
     T: Send + 'static + Clone,
 {
-    worker_index: usize,
     location: OperatorLocation,
     init: IF,
     deserialize: D,
@@ -1406,7 +1370,6 @@ where
     T: Send + 'static + Clone + Debug,
 {
     pub(crate) fn new(
-        worker_index: usize,
         location: OperatorLocation,
         exchange: Arc<Exchange<(T, bool)>>,
         init: IF,
@@ -1414,10 +1377,7 @@ where
         deserialize: D,
         combine: L,
     ) -> Self {
-        debug_assert!(worker_index < Runtime::num_workers());
-
         Self {
-            worker_index,
             location,
             init,
             combine,
@@ -1430,10 +1390,6 @@ where
             total_wait_time: Arc::new(AtomicU64::new(0)),
         }
     }
-
-    fn is_ready(&self) -> bool {
-        self.exchange.ready_to_receive(self.worker_index)
-    }
 }
 
 impl<IF, T, L, D> Operator for ExchangeReceiver<IF, T, L, D>
@@ -1464,48 +1420,6 @@ where
         });
     }
 
-    fn is_async(&self) -> bool {
-        true
-    }
-
-    fn register_ready_callback<F>(&mut self, cb: F)
-    where
-        F: Fn() + Send + Sync + 'static,
-    {
-        let start_wait_usecs = self.start_wait_usecs.clone();
-        let total_wait_time = self.total_wait_time.clone();
-        let exchange = self.exchange.clone();
-        let worker_index = self.worker_index;
-
-        let cb = move || {
-            if exchange.ready_to_receive(worker_index) {
-                // The callback can be invoked multiple times per step.
-                // Reset start_wait_usecs to 0 to make sure we don't double-count.
-                let start = start_wait_usecs.swap(0, Ordering::Acquire);
-                if start != 0 {
-                    let end = current_time_usecs();
-                    if end > start {
-                        let wait_time_usecs = end - start;
-                        // if worker_index == 0 {
-                        //     info!(
-                        //         "{worker_index}: {} +{wait_time_usecs}",
-                        //         exchange.exchange_id()
-                        //     );
-                        // }
-                        total_wait_time.fetch_add(wait_time_usecs, Ordering::AcqRel);
-                    }
-                }
-            }
-            cb()
-        };
-        self.exchange
-            .register_receiver_callback(self.worker_index, cb)
-    }
-
-    fn ready(&self) -> bool {
-        self.is_ready()
-    }
-
     fn fixedpoint(&self, _scope: Scope) -> bool {
         true
     }
@@ -1550,7 +1464,16 @@ where
         };
 
         let mut combined = (self.init)();
-        let res = self.exchange.receive_all(deserialize).await;
+        let (res, start_wait_usecs) = self
+            .exchange
+            .receive_all(deserialize, Some(&self.start_wait_usecs))
+            .await;
+        if let Some(start_wait_usecs) = start_wait_usecs {
+            self.total_wait_time.fetch_add(
+                current_time_usecs().saturating_sub(start_wait_usecs),
+                Ordering::Release,
+            );
+        }
         for (data, flushed) in res {
             if flushed {
                 self.flush_count += 1;
@@ -1632,7 +1555,6 @@ where
         return None;
     }
     let runtime = Runtime::runtime().unwrap();
-    let worker_index = Runtime::worker_index();
 
     let exchange_id = runtime.sequence_next().try_into().unwrap();
     let start_wait_usecs = Arc::new(AtomicU64::new(0));
@@ -1644,7 +1566,6 @@ where
         partition,
     );
     let receiver = ExchangeReceiver::new(
-        worker_index,
         location,
         exchange,
         init,
@@ -1703,8 +1624,8 @@ mod tests {
                     })
                     .await;
 
-                let received = exchange
-                    .receive_all(|data| aligned_deserialize(&data[..]))
+                let (received, _) = exchange
+                    .receive_all(|data| aligned_deserialize(&data[..]), None)
                     .await;
 
                 let expected = (0..n_workers).map(|worker| (worker, round)).collect_vec();
diff --git a/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs b/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs
index 06b87e61ab6..785c658e7a5 100644
--- a/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs
+++ b/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs
@@ -151,7 +151,6 @@ where
                     let batch_factories_clone = batch_factories.clone();
 
                     let receiver = circuit.add_source(ExchangeReceiver::new(
-                        worker_index,
                         Some(location),
                         exchange.clone(),
                         || Vec::new(),

From 06b4db81488deefbac5535313a6945fdec1af1f7 Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@feldera.com>
Date: Wed, 10 Jun 2026 13:16:29 -0700
Subject: [PATCH 4/8] [dbsp] Skip exchange overhead for input operators during
 transaction commit.

Input operators never exchange any actual data once a circuit's transaction
starts to commit, but until now they still incurred the overhead and
latency of an exchange.  This commit fixes that for at least two input
operators (perhaps others need improvements too).

The most interesting part of this commit is the new builder for sharding,
`Sharder` in `crates/dbsp/src/operator/dynamic/communication/shard.rs`.
I don't think we've used this pattern before for constructing operators.
It seems appropriate in this case because sharding is now configurable
among multiple axes and it was getting unwieldy to add new methods for
each combination.

Signed-off-by: Ben Pfaff <blp@feldera.com>
---
 crates/dbsp/src/circuit/runtime.rs            |   5 +-
 crates/dbsp/src/operator/communication.rs     |   4 +-
 .../src/operator/communication/exchange.rs    |  98 ++++++--
 crates/dbsp/src/operator/dynamic.rs           |   1 +
 .../balance/accumulate_trace_balanced.rs      |   4 +-
 .../operator/dynamic/communication/gather.rs  |   3 +-
 .../operator/dynamic/communication/shard.rs   | 217 ++++++++++++------
 crates/dbsp/src/operator/dynamic/input.rs     |  11 +-
 .../operator/dynamic/time_series/waterline.rs |   4 +-
 crates/sqllib/src/string_interner.rs          |   3 +-
 10 files changed, 262 insertions(+), 88 deletions(-)

diff --git a/crates/dbsp/src/circuit/runtime.rs b/crates/dbsp/src/circuit/runtime.rs
index 0b0aec4ec77..9720ffe1dff 100644
--- a/crates/dbsp/src/circuit/runtime.rs
+++ b/crates/dbsp/src/circuit/runtime.rs
@@ -7,7 +7,7 @@ use crate::SchedulerError;
 use crate::circuit::checkpointer::Checkpointer;
 use crate::circuit::dbsp_handle::StepSize;
 use crate::error::Error as DbspError;
-use crate::operator::communication::Exchange;
+use crate::operator::communication::{Exchange, ExchangeActivity};
 use crate::storage::backend::StorageBackend;
 use crate::storage::file::format::Compression;
 use crate::storage::file::writer::Parameters;
@@ -1383,7 +1383,8 @@ where
         match Runtime::runtime() {
             Some(runtime) if Runtime::num_workers() > 1 => {
                 let exchange_id = runtime.sequence_next().try_into().unwrap();
-                let exchange = Exchange::with_runtime(&runtime, exchange_id);
+                let exchange =
+                    Exchange::with_runtime(&runtime, exchange_id, ExchangeActivity::AllSteps);
                 let identifier = Arc::new(format!("broadcast {name} (exchange {exchange_id})"));
 
                 Self::MultiThreaded {
diff --git a/crates/dbsp/src/operator/communication.rs b/crates/dbsp/src/operator/communication.rs
index ec1eb6b591a..a8ed35132ca 100644
--- a/crates/dbsp/src/operator/communication.rs
+++ b/crates/dbsp/src/operator/communication.rs
@@ -5,4 +5,6 @@ mod shard;
 pub(crate) use exchange::{
     Exchange, ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, pop_flushed,
 };
-pub use exchange::{ExchangeReceiver, ExchangeSender, Mailbox, new_exchange_operators};
+pub use exchange::{
+    ExchangeActivity, ExchangeReceiver, ExchangeSender, Mailbox, new_exchange_operators,
+};
diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs
index 75079d9520d..511c32780f3 100644
--- a/crates/dbsp/src/operator/communication/exchange.rs
+++ b/crates/dbsp/src/operator/communication/exchange.rs
@@ -703,6 +703,9 @@ pub(crate) struct Exchange<T> {
 
     /// The number of bytes serialized.
     deserialized_bytes: AtomicUsize,
+
+    /// When the exchange is active.
+    activity: ExchangeActivity,
 }
 
 // Stop Rust from complaining about unused field.
@@ -763,6 +766,7 @@ where
         clients: Arc<ExchangeClients>,
         exchange_id: ExchangeId,
         directory: &ExchangeDirectory,
+        activity: ExchangeActivity,
     ) -> Arc<Self> {
         let npeers = Runtime::num_workers();
         let mailboxes: Vec<Mutex<Option<Mailbox<T>>>> =
@@ -788,6 +792,7 @@ where
             mailboxes,
             deserialization_usecs: AtomicU64::new(0),
             deserialized_bytes: AtomicUsize::new(0),
+            activity,
         });
 
         directory.insert(exchange_id, exchange.clone());
@@ -816,7 +821,11 @@ where
     /// Create a new `Exchange` instance if an instance with the same id
     /// (created by another thread) does not yet exist within `runtime`.
     /// The number of peers will be set to `runtime.num_workers()`.
-    pub(crate) fn with_runtime(runtime: &Runtime, exchange_id: ExchangeId) -> Arc<Self> {
+    pub(crate) fn with_runtime(
+        runtime: &Runtime,
+        exchange_id: ExchangeId,
+        activity: ExchangeActivity,
+    ) -> Arc<Self> {
         // It's tempting to move the following calls to create the
         // `ExchangeDirectory` and `ExchangeClients` into `Exchange::new`, but
         // don't do it: all three of these access `runtime.local_store` and
@@ -826,7 +835,7 @@ where
         runtime
             .local_store()
             .entry(ExchangeCacheId::new(exchange_id))
-            .or_insert_with(|| Exchange::new(runtime, clients, exchange_id, &directory))
+            .or_insert_with(|| Exchange::new(runtime, clients, exchange_id, &directory, activity))
             .value()
             .clone()
     }
@@ -1041,6 +1050,19 @@ where
     }
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+enum Phase {
+    Active,
+    Flush,
+    Commit,
+}
+
+impl Phase {
+    fn is_inactive(&self, activity: ExchangeActivity) -> bool {
+        *self == Phase::Commit && activity == ExchangeActivity::InputOnly
+    }
+}
+
 /// Operator that partitions incoming data across all workers.
 ///
 /// This operator works in tandem with [`ExchangeReceiver`], which reassembles
@@ -1154,7 +1176,7 @@ where
 /// use dbsp::{
 ///     operator::{communication::new_exchange_operators, Generator},
 ///     circuit::{WorkerLocation, WorkerLocations},
-///     operator::communication::Mailbox,
+///     operator::communication::{ExchangeActivity, Mailbox},
 ///     Circuit, RootCircuit, Runtime,
 ///     storage::file::to_bytes_dyn,
 ///     trace::aligned_deserialize,
@@ -1190,6 +1212,7 @@ where
 ///             },
 ///             |data| aligned_deserialize(&data[..]),///             // Reassemble received values into a vector.
 ///             |v: &mut Vec<usize>, n| v.push(n),
+///             ExchangeActivity::AllSteps,
 ///         ).unwrap();
 ///
 ///         // Add exchange operators to the circuit.
@@ -1231,7 +1254,7 @@ where
     // Input batch sizes.
     input_batch_stats: BatchSizeStats,
 
-    flushed: bool,
+    phase: Phase,
 
     // The instant when the sender produced its outputs, and the
     // receiver starts waiting for all other workers to produce their
@@ -1258,7 +1281,7 @@ where
             outputs: Vec::with_capacity(Runtime::num_workers()),
             exchange,
             input_batch_stats: BatchSizeStats::new(),
-            flushed: false,
+            phase: Phase::Active,
             start_wait_usecs,
             phantom: PhantomData,
         }
@@ -1296,8 +1319,12 @@ where
         true
     }
 
+    fn start_transaction(&mut self) {
+        self.phase = Phase::Active;
+    }
+
     fn flush(&mut self) {
-        self.flushed = true;
+        self.phase = Phase::Flush;
     }
 }
 
@@ -1312,6 +1339,16 @@ where
     }
 
     async fn eval_owned(&mut self, input: D) {
+        if self.phase.is_inactive(self.exchange.activity) {
+            return;
+        };
+        let flushed = if self.phase == Phase::Flush {
+            self.phase = Phase::Commit;
+            true
+        } else {
+            false
+        };
+
         self.input_batch_stats.add_batch(input.num_entries_deep());
 
         debug_assert!(self.ready());
@@ -1322,16 +1359,14 @@ where
 
         let data = self.outputs.drain(..).map(|mailbox| match mailbox {
             Mailbox::Tx(mut data) => {
-                data.push(self.flushed as u8);
+                data.push(flushed as u8);
                 Mailbox::Tx(data)
             }
             Mailbox::Rx(_) => unreachable!(),
-            Mailbox::Plain(item) => Mailbox::Plain((item, self.flushed)),
+            Mailbox::Plain(item) => Mailbox::Plain((item, flushed)),
         });
 
         self.exchange.send_all(&self.global_node_id, data).await;
-
-        self.flushed = false;
     }
 
     fn input_preference(&self) -> OwnershipPreference {
@@ -1360,6 +1395,7 @@ where
     flush_complete: bool,
     start_wait_usecs: Arc<AtomicU64>,
     total_wait_time: Arc<AtomicU64>,
+    phase: Phase,
 
     // Output batch sizes.
     output_batch_stats: BatchSizeStats,
@@ -1388,6 +1424,7 @@ where
             output_batch_stats: BatchSizeStats::new(),
             start_wait_usecs,
             total_wait_time: Arc::new(AtomicU64::new(0)),
+            phase: Phase::Active,
         }
     }
 }
@@ -1424,6 +1461,10 @@ where
         true
     }
 
+    fn start_transaction(&mut self) {
+        self.phase = Phase::Active;
+    }
+
     fn flush(&mut self) {
         // println!("{} exchange_receiver::flush", Runtime::worker_index());
         self.flush_complete = false;
@@ -1456,7 +1497,10 @@ where
     D: Fn(AlignedVec) -> T + Send + Sync + 'static,
 {
     async fn eval(&mut self) -> O {
-        debug_assert!(self.ready());
+        if self.phase.is_inactive(self.exchange.activity) {
+            return (self.init)();
+        }
+
         let deserialize = |mut vec: AlignedVec| {
             let flushed = pop_flushed(&mut vec);
             let value = (self.deserialize)(vec);
@@ -1489,6 +1533,7 @@ where
 
             self.flush_complete = true;
             self.flush_count = 0;
+            self.phase = Phase::Commit;
         }
 
         self.output_batch_stats
@@ -1511,6 +1556,28 @@ impl TypedMapKey<LocalStoreMarker> for DirectoryId {
     type Value = ExchangeDirectory;
 }
 
+/// The microsteps during which an exchange is active.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum ExchangeActivity {
+    /// The exchange is active in every microstep.
+    ///
+    /// This includes pre-commit and commit microstep.
+    AllSteps,
+
+    /// The exchange is active only during pre-commit microsteps.
+    ///
+    /// This allows for optimizations for exchanges used for sharding data from
+    /// input operators, which don't exchange any data during commit.
+    ///
+    /// # Limitation
+    ///
+    /// The current implementation only works for operators that flush in the
+    /// same (micro)step in every worker.  This is true for input operators,
+    /// which flush as soon as the transaction starts committing, but it is not
+    /// necessarily true for other operators.
+    InputOnly,
+}
+
 /// Create an [`ExchangeSender`]/[`ExchangeReceiver`] operator pair.
 ///
 /// See [`ExchangeSender`] documentation for details and example usage.
@@ -1542,6 +1609,7 @@ pub fn new_exchange_operators<TI, TO, TE, IF, PL, CL, D>(
     partition: PL,
     deserialize: D,
     combine: CL,
+    activity: ExchangeActivity,
 ) -> Option<(ExchangeSender<TI, TE, PL>, ExchangeReceiver<IF, TE, CL, D>)>
 where
     TO: Clone,
@@ -1558,7 +1626,7 @@ where
 
     let exchange_id = runtime.sequence_next().try_into().unwrap();
     let start_wait_usecs = Arc::new(AtomicU64::new(0));
-    let exchange = Exchange::with_runtime(&runtime, exchange_id);
+    let exchange = Exchange::with_runtime(&runtime, exchange_id, activity);
     let sender = ExchangeSender::new(
         location,
         exchange.clone(),
@@ -1591,7 +1659,7 @@ mod tests {
         },
         operator::{
             Generator,
-            communication::{Mailbox, new_exchange_operators},
+            communication::{ExchangeActivity, Mailbox, new_exchange_operators},
         },
         storage::file::{to_bytes, to_bytes_dyn},
         trace::aligned_deserialize,
@@ -1612,7 +1680,8 @@ mod tests {
     // value `(sender, n)` to each receiver, where `sender` is the sender's
     // worker number in round `n`.
     fn circuit() {
-        let exchange = Exchange::with_runtime(&Runtime::runtime().unwrap(), 0);
+        let exchange =
+            Exchange::with_runtime(&Runtime::runtime().unwrap(), 0, ExchangeActivity::AllSteps);
         TOKIO.block_on(async {
             let sender = Runtime::worker_index();
             let n_workers = Runtime::num_workers();
@@ -1742,6 +1811,7 @@ mod tests {
                 },
                 |data| aligned_deserialize(&data[..]),
                 |v: &mut Vec<usize>, n| v.push(n),
+                ExchangeActivity::AllSteps,
             )
             .unwrap();
 
diff --git a/crates/dbsp/src/operator/dynamic.rs b/crates/dbsp/src/operator/dynamic.rs
index 5bf472b5f1e..f35b4db6108 100644
--- a/crates/dbsp/src/operator/dynamic.rs
+++ b/crates/dbsp/src/operator/dynamic.rs
@@ -34,6 +34,7 @@ pub mod time_series;
 pub mod trace;
 pub(crate) mod upsert;
 
+pub use communication::shard::Sharder;
 pub(crate) use communication::shard::shard_batch;
 
 /// The "standard" indexed Z-set type used by monomorphic
diff --git a/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs b/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs
index 785c658e7a5..ec2ca8fb3a4 100644
--- a/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs
+++ b/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs
@@ -20,7 +20,7 @@ use crate::{
     operator::{
         Z1,
         async_stream_operators::{StreamingTernarySinkOperator, StreamingTernarySinkWrapper},
-        communication::{Exchange, ExchangeReceiver},
+        communication::{Exchange, ExchangeActivity, ExchangeReceiver},
         dynamic::{
             accumulate_trace::{AccumulateBoundsId, AccumulateTraceAppend, AccumulateZ1Trace},
             balance::{
@@ -145,7 +145,7 @@ where
                     // during the current transaction. It will be set to `true` exactly once per transaction.
                     // Once all senders are flushed, the receiver can report itself as flushed too.
                     let exchange: Arc<Exchange<(B, bool)>> =
-                        Exchange::with_runtime(&runtime, exchange_id);
+                        Exchange::with_runtime(&runtime, exchange_id, ExchangeActivity::AllSteps);
 
                     // Exchange receiver
                     let batch_factories_clone = batch_factories.clone();
diff --git a/crates/dbsp/src/operator/dynamic/communication/gather.rs b/crates/dbsp/src/operator/dynamic/communication/gather.rs
index ce2c962ae82..40eb252178b 100644
--- a/crates/dbsp/src/operator/dynamic/communication/gather.rs
+++ b/crates/dbsp/src/operator/dynamic/communication/gather.rs
@@ -8,7 +8,7 @@ use crate::{
         runtime::{WorkerLocation, WorkerLocations},
     },
     circuit_cache_key,
-    operator::communication::{Mailbox, new_exchange_operators},
+    operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators},
     trace::{Batch, deserialize_indexed_wset, merge_batches, serialize_indexed_wset},
 };
 use arc_swap::ArcSwap;
@@ -159,6 +159,7 @@ where
                             },
                             move |data| deserialize_indexed_wset(&factories_clone, &data),
                             |batches: &mut Vec<B>, batch: B| batches.push(batch),
+                            ExchangeActivity::AllSteps,
                         )
                         .unwrap();
 
diff --git a/crates/dbsp/src/operator/dynamic/communication/shard.rs b/crates/dbsp/src/operator/dynamic/communication/shard.rs
index f5cd70af397..f1e212d6f71 100644
--- a/crates/dbsp/src/operator/dynamic/communication/shard.rs
+++ b/crates/dbsp/src/operator/dynamic/communication/shard.rs
@@ -16,7 +16,7 @@ use crate::{
     },
     circuit_cache_key,
     dynamic::{Data, DataTrait, DynPair, DynPairs, Factory},
-    operator::communication::{Mailbox, new_exchange_operators},
+    operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators},
     storage::file::SerializerInner,
     trace::{
         Batch, BatchReader, Builder, IndexedWSetSerializer, deserialize_indexed_wset, merge_batches,
@@ -35,68 +35,59 @@ fn all_workers() -> Range<usize> {
     0..Runtime::num_workers()
 }
 
-impl<C, IB> Stream<C, IB>
+/// Builder for sharding a batch or pairs.
+pub struct Sharder<'a, C, B> {
+    stream: &'a Stream<C, B>,
+    workers: Range<usize>,
+    activity: ExchangeActivity,
+}
+
+impl<'a, C, B> Sharder<'a, C, B>
 where
     C: Circuit,
-    IB: BatchReader<Time = ()> + Clone,
 {
-    /// See [`Stream::shard`].
-    #[track_caller]
-    pub fn dyn_shard(&self, factories: &IB::Factories) -> Stream<C, IB>
-    where
-        IB: Batch + Send,
-    {
-        // `shard_generic` returns `None` if there is only one worker thread
-        // and hence sharding is a no-op.  In this case, we simply return the
-        // input stream.  This allows us to use `shard` unconditionally without
-        // incurring any overhead in the single-threaded case.
-        self.dyn_shard_generic(factories)
-            .unwrap_or_else(|| self.clone())
+    /// Sets `workers` as the range of workers to shard to.  This is useful for
+    /// sharding to workers on one particular host in a multihost pipeline.  By
+    /// default, sharding happens across all workers.
+    pub fn with_workers(self, workers: Range<usize>) -> Self {
+        Self { workers, ..self }
     }
 
-    /// See [`Stream::shard_workers`].
+    /// Returns the sharder configured to be active during `activity`.  The
+    /// default is [ExchangeActivity::AllSteps]
+    pub fn with_activity(self, activity: ExchangeActivity) -> Self {
+        Self { activity, ..self }
+    }
+
+    /// Shards the stream with output to a batch of the same kind.
     #[track_caller]
-    pub fn dyn_shard_workers(
-        &self,
-        workers: Range<usize>,
-        factories: &IB::Factories,
-    ) -> Stream<C, IB>
+    pub fn shard(self, factories: &B::Factories) -> Stream<C, B>
     where
-        IB: Batch + Send,
+        B: Batch<Time = ()> + Clone + Send,
     {
-        // `shard_generic_workers` returns `None` if there is only one worker
-        // thread and hence sharding is a no-op.  In this case, we simply return
-        // the input stream.  This allows us to use `shard_workers`
-        // unconditionally without incurring any overhead in the single-threaded
-        // case.
-        self.dyn_shard_generic_workers(workers, factories)
-            .unwrap_or_else(|| self.clone())
+        self.shard_generic_ref(factories)
+            .unwrap_or_else(|| self.stream.clone())
     }
+}
 
-    /// Like [`Self::dyn_shard`], but can assemble the results into any output batch
-    /// type `OB`.
-    ///
-    /// Returns `None` when the circuit is not running inside a multithreaded
-    /// runtime or is running in a runtime with a single worker thread.
+impl<'a, C, IB> Sharder<'a, C, IB>
+where
+    C: Circuit,
+    IB: BatchReader<Time = ()> + Clone,
+{
+    /// Shards the stream with output to batch type `OB`.  If the stream doesn't
+    /// need to be sharded because there is only a single worker, returns
+    /// `None`.
     #[track_caller]
-    pub fn dyn_shard_generic<OB>(&self, factories: &OB::Factories) -> Option<Stream<C, OB>>
+    pub fn shard_generic<OB>(self, factories: &OB::Factories) -> Option<Stream<C, OB>>
     where
         OB: Batch<Key = IB::Key, Val = IB::Val, Time = (), R = IB::R> + Send,
     {
-        self.dyn_shard_generic_workers(all_workers(), factories)
+        self.shard_generic_ref(factories)
     }
 
-    /// Like [`Self::dyn_shard`], but can assemble the results into any output batch
-    /// type `OB`.
-    ///
-    /// Returns `None` when the circuit is not running inside a multithreaded
-    /// runtime or is running in a runtime with a single worker thread.
     #[track_caller]
-    pub fn dyn_shard_generic_workers<OB>(
-        &self,
-        workers: Range<usize>,
-        factories: &OB::Factories,
-    ) -> Option<Stream<C, OB>>
+    fn shard_generic_ref<OB>(&self, factories: &OB::Factories) -> Option<Stream<C, OB>>
     where
         OB: Batch<Key = IB::Key, Val = IB::Val, Time = (), R = IB::R> + Send,
     {
@@ -105,9 +96,10 @@ where
         }
         let location = Location::caller();
         let output = self
+            .stream
             .circuit()
             .cache_get_or_insert_with(
-                ShardId::new((self.stream_id(), workers.clone())),
+                ShardId::new((self.stream.stream_id(), self.workers.clone())),
                 move || {
                     // As a minor optimization, we reuse this array across all invocations
                     // of the sharding operator.
@@ -115,10 +107,10 @@ where
                     let factories_clone2 = factories.clone();
                     let factories_clone3 = factories.clone();
                     let factories_clone4 = factories.clone();
-                    let workers_clone = workers.clone();
-                    let workers_clone2 = workers.clone();
+                    let workers_clone = self.workers.clone();
+                    let workers_clone2 = self.workers.clone();
 
-                    let output = self.circuit().region("shard", || {
+                    let output = self.stream.circuit().region("shard", || {
                         let (sender, receiver) = new_exchange_operators(
                             Some(location),
                             || Vec::new(),
@@ -133,26 +125,30 @@ where
                             },
                             move |data| deserialize_indexed_wset(&factories_clone4, &data),
                             |batches: &mut Vec<OB>, batch: OB| batches.push(batch),
+                            self.activity,
                         )
                         .unwrap();
 
-                        self.circuit()
-                            .add_exchange(sender, receiver, self)
+                        self.stream
+                            .circuit()
+                            .add_exchange(sender, receiver, self.stream)
                             .apply_owned_named("merge shards", move |batches| {
                                 merge_batches(&factories_clone2, batches, &None, &None)
                             })
                     });
 
-                    self.circuit().cache_insert(
+                    self.stream.circuit().cache_insert(
                         ShardId::new((output.stream_id(), workers_clone2)),
                         output.clone(),
                     );
 
-                    self.circuit()
-                        .cache_insert(UnshardId::new(output.stream_id()), self.clone());
+                    self.stream
+                        .circuit()
+                        .cache_insert(UnshardId::new(output.stream_id()), self.stream.clone());
 
                     output.set_persistent_id(
-                        self.get_persistent_id()
+                        self.stream
+                            .get_persistent_id()
                             .map(|name| format!("{name}.shard"))
                             .as_deref(),
                     )
@@ -164,19 +160,20 @@ where
     }
 }
 
-impl<C, K, V> Stream<C, Vec<Box<DynPairs<K, V>>>>
+impl<'a, C, K, V> Sharder<'a, C, Vec<Box<DynPairs<K, V>>>>
 where
     C: Circuit,
     K: DataTrait + ?Sized,
     V: DataTrait + ?Sized,
 {
+    /// Shards a vector of pairs across workers.
     #[track_caller]
-    pub fn dyn_shard_pairs(
-        &self,
+    pub fn shard_pairs(
+        self,
         pairs_factory: &'static dyn Factory<DynPairs<K, V>>,
     ) -> Stream<C, Vec<Box<DynPairs<K, V>>>> {
-        if self.is_sharded() {
-            return self.clone();
+        if self.stream.is_sharded() {
+            return self.stream.clone();
         }
 
         let location = Location::caller();
@@ -186,19 +183,24 @@ where
             Vec::new,
             move |input_pairs: Vec<Box<DynPairs<K, V>>>,
                   output_pairs: &mut Vec<Mailbox<Box<DynPairs<K, V>>>>| {
-                shard_pairs(input_pairs, &all_workers(), output_pairs, pairs_factory);
+                shard_pairs(input_pairs, &self.workers, output_pairs, pairs_factory);
             },
             move |data| deserialize_pairs(&data, pairs_factory),
             |output_pairs: &mut Vec<Box<DynPairs<K, V>>>, batch: Box<DynPairs<K, V>>| {
                 output_pairs.push(batch);
             },
+            self.activity,
         )
         .unwrap();
 
-        let output = self.circuit().add_exchange(sender, receiver, self);
+        let output = self
+            .stream
+            .circuit()
+            .add_exchange(sender, receiver, self.stream);
 
         output.set_persistent_id(
-            self.get_persistent_id()
+            self.stream
+                .get_persistent_id()
                 .map(|name| format!("{name}.shard"))
                 .as_deref(),
         );
@@ -206,6 +208,93 @@ where
     }
 }
 
+impl<C, B> Stream<C, B> {
+    /// Constructs a [Sharder] builder for sharding a batch or pairs with
+    /// flexible configuration.
+    pub fn dyn_sharder(&self) -> Sharder<'_, C, B> {
+        Sharder {
+            stream: self,
+            activity: ExchangeActivity::AllSteps,
+            workers: all_workers(),
+        }
+    }
+}
+
+impl<C, IB> Stream<C, IB>
+where
+    C: Circuit,
+    IB: BatchReader<Time = ()> + Clone,
+{
+    /// See [`Stream::shard`].
+    #[track_caller]
+    pub fn dyn_shard(&self, factories: &IB::Factories) -> Stream<C, IB>
+    where
+        IB: Batch + Send,
+    {
+        self.dyn_sharder().shard(factories)
+    }
+
+    /// See [`Stream::shard_workers`].
+    #[track_caller]
+    pub fn dyn_shard_workers(
+        &self,
+        workers: Range<usize>,
+        factories: &IB::Factories,
+    ) -> Stream<C, IB>
+    where
+        IB: Batch + Send,
+    {
+        self.dyn_sharder().with_workers(workers).shard(factories)
+    }
+
+    /// Like [`Self::dyn_shard`], but can assemble the results into any output batch
+    /// type `OB`.
+    ///
+    /// Returns `None` when the circuit is not running inside a multithreaded
+    /// runtime or is running in a runtime with a single worker thread.
+    #[track_caller]
+    pub fn dyn_shard_generic<OB>(&self, factories: &OB::Factories) -> Option<Stream<C, OB>>
+    where
+        OB: Batch<Key = IB::Key, Val = IB::Val, Time = (), R = IB::R> + Send,
+    {
+        self.dyn_sharder().shard_generic(factories)
+    }
+
+    /// Like [`Self::dyn_shard`], but can assemble the results into any output batch
+    /// type `OB`.
+    ///
+    /// Returns `None` when the circuit is not running inside a multithreaded
+    /// runtime or is running in a runtime with a single worker thread.
+    #[track_caller]
+    pub fn dyn_shard_generic_workers<OB>(
+        &self,
+        workers: Range<usize>,
+        factories: &OB::Factories,
+    ) -> Option<Stream<C, OB>>
+    where
+        OB: Batch<Key = IB::Key, Val = IB::Val, Time = (), R = IB::R> + Send,
+    {
+        self.dyn_sharder()
+            .with_workers(workers)
+            .shard_generic(factories)
+    }
+}
+
+impl<C, K, V> Stream<C, Vec<Box<DynPairs<K, V>>>>
+where
+    C: Circuit,
+    K: DataTrait + ?Sized,
+    V: DataTrait + ?Sized,
+{
+    #[track_caller]
+    pub fn dyn_shard_pairs(
+        &self,
+        pairs_factory: &'static dyn Factory<DynPairs<K, V>>,
+    ) -> Stream<C, Vec<Box<DynPairs<K, V>>>> {
+        self.dyn_sharder().shard_pairs(pairs_factory)
+    }
+}
+
 fn deserialize_pairs<K, V>(
     data: &[u8],
     pairs_factory: &'static dyn Factory<DynPairs<K, V>>,
diff --git a/crates/dbsp/src/operator/dynamic/input.rs b/crates/dbsp/src/operator/dynamic/input.rs
index ed79ff25197..e655f60e120 100644
--- a/crates/dbsp/src/operator/dynamic/input.rs
+++ b/crates/dbsp/src/operator/dynamic/input.rs
@@ -10,6 +10,7 @@ use crate::{
     },
     operator::{
         Input, InputHandle, Update,
+        communication::ExchangeActivity,
         dynamic::{
             input_upsert::{
                 DynUpdate, InputUpsertFactories, InputUpsertWithWaterlineFactories, PatchFunc,
@@ -390,7 +391,11 @@ impl RootCircuit {
             // This adds small overhead to tables that don't get materialized and hence don't need to get sharded.
             // If this proves to be a problem in practice, we can add a variant of this function that doesn't shard
             // its output stream for use with non-materializes tables.
-            let stream = self.add_source(input).dyn_shard(&factories.zset_factories);
+            let stream = self
+                .add_source(input)
+                .dyn_sharder()
+                .with_activity(ExchangeActivity::InputOnly)
+                .shard(&factories.zset_factories);
 
             let zset_handle = <CollectionHandle<DynPair<K, DynUnit>, DynZWeight>>::new(
                 factories.pair_factory,
@@ -474,7 +479,9 @@ impl RootCircuit {
         //       |--->[some other operator]
         let stream = self
             .add_source(input)
-            .dyn_shard(&factories.indexed_zset_factories);
+            .dyn_sharder()
+            .with_activity(ExchangeActivity::InputOnly)
+            .shard(&factories.indexed_zset_factories);
 
         let zset_handle = <CollectionHandle<K, DynPair<V, DynZWeight>>>::new(
             factories.pair_factory,
diff --git a/crates/dbsp/src/operator/dynamic/time_series/waterline.rs b/crates/dbsp/src/operator/dynamic/time_series/waterline.rs
index 10a2bd40769..1b919acb79f 100644
--- a/crates/dbsp/src/operator/dynamic/time_series/waterline.rs
+++ b/crates/dbsp/src/operator/dynamic/time_series/waterline.rs
@@ -5,7 +5,7 @@ use size_of::SizeOf;
 use crate::circuit::checkpointer::Checkpoint;
 use crate::circuit::runtime::{WorkerLocation, WorkerLocations};
 use crate::dynamic::{DynData, DynUnit};
-use crate::operator::communication::Mailbox;
+use crate::operator::communication::{ExchangeActivity, Mailbox};
 use crate::operator::dynamic::{MonoIndexedZSet, MonoZSet};
 use crate::{
     Circuit, NumEntries, RootCircuit, Stream,
@@ -111,6 +111,7 @@ where
                         *result = waterline;
                     }
                 },
+                ExchangeActivity::AllSteps,
             );
 
             match exchange {
@@ -191,6 +192,7 @@ where
                     let old_result = clone_box(result);
                     least_upper_bound(&old_result, &waterline, result.as_mut());
                 },
+                ExchangeActivity::AllSteps,
             );
 
             match exchange {
diff --git a/crates/sqllib/src/string_interner.rs b/crates/sqllib/src/string_interner.rs
index b4c2f79b4d9..60747c0fefe 100644
--- a/crates/sqllib/src/string_interner.rs
+++ b/crates/sqllib/src/string_interner.rs
@@ -10,7 +10,7 @@ use dbsp::{
     Circuit, DynZWeight, OrdZSet, RootCircuit, Runtime, Stream, ZWeight,
     circuit::{LocalStoreMarker, WorkerLocation, WorkerLocations},
     dynamic::{DowncastTrait, DynData},
-    operator::communication::{Mailbox, new_exchange_operators},
+    operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators},
     storage::file::to_bytes,
     trace::{
         BatchReader, BatchReaderFactories, Cursor, OrdIndexedWSet as DynOrdIndexedWSet,
@@ -315,6 +315,7 @@ pub fn build_string_interner(
                 snapshot.extend(remote_snapshot);
             }
         },
+        ExchangeActivity::AllSteps,
     );
     let by_id = match exchange {
         Some((sender, receiver)) => interned_strings

From 0e8e08d2d83bf6377fd335d8f346eed0653f7d3e Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@feldera.com>
Date: Fri, 12 Jun 2026 13:30:20 -0700
Subject: [PATCH 5/8] [dbsp] Segregate inter-host exchange channels by message
 size and type.

Until now, when one host in a DBSP pipeline sends another one a message
as part of synchronous or streaming exchange, it sent it over a single
host-to-host TCP connection.  This means that synchronous exchange has to
wait for streaming exchange messages, which increases latency.  This
commit uses different TCP connections for streaming and synchronous
exchange, which helps.  In addition, it uses separate TCP connections for
small and large synchronous exchange; we speculate that giving small
synchronous messages priority can help with broadcast and consensus
messages (this is unproven).

Signed-off-by: Ben Pfaff <blp@feldera.com>
---
 crates/dbsp/src/operator/communication.rs     |  3 +-
 .../src/operator/communication/exchange.rs    | 68 ++++++++++++++++---
 .../operator/dynamic/sharded_accumulator.rs   | 15 ++--
 3 files changed, 67 insertions(+), 19 deletions(-)

diff --git a/crates/dbsp/src/operator/communication.rs b/crates/dbsp/src/operator/communication.rs
index a8ed35132ca..73ab4ef33cc 100644
--- a/crates/dbsp/src/operator/communication.rs
+++ b/crates/dbsp/src/operator/communication.rs
@@ -3,7 +3,8 @@ mod gather;
 mod shard;
 
 pub(crate) use exchange::{
-    Exchange, ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, pop_flushed,
+    Exchange, ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, MessageType,
+    pop_flushed,
 };
 pub use exchange::{
     ExchangeActivity, ExchangeReceiver, ExchangeSender, Mailbox, new_exchange_operators,
diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs
index 511c32780f3..23197f55d50 100644
--- a/crates/dbsp/src/operator/communication/exchange.rs
+++ b/crates/dbsp/src/operator/communication/exchange.rs
@@ -23,6 +23,7 @@ use crate::{
 };
 use binrw::{BinRead, BinWrite};
 use crossbeam_utils::CachePadded;
+use enum_map::{Enum, EnumMap};
 use feldera_samply::Span;
 use feldera_storage::fbuf::FBuf;
 use itertools::Itertools;
@@ -239,6 +240,45 @@ struct ExchangeMessage {
     data: Vec<FBuf>,
 }
 
+/// Distinguishes messages by size.
+///
+/// We segregate big and small messages into separate queues so that simple
+/// broadcast and consensus messages don't get delayed behind bigger messages.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Enum)]
+pub(crate) enum MessageSize {
+    /// A big message.
+    Big,
+
+    /// A small message.
+    Small,
+}
+
+impl MessageSize {
+    /// Constructs `MessageSize` from a count of `bytes`.
+    pub(crate) fn from_bytes(bytes: usize) -> Self {
+        if bytes <= 4096 {
+            Self::Small
+        } else {
+            Self::Big
+        }
+    }
+}
+
+/// Categorizes an [ExchangeMessage].
+///
+/// We segregate messages in different categories into different queues so that
+/// messages in one category do not delay those in other categories.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Enum)]
+pub(crate) enum MessageType {
+    /// Messages sent via synchronous exchange.
+    Synchronous(
+        /// Message size.
+        MessageSize,
+    ),
+    /// Messages sent via streaming exchange.
+    Streaming,
+}
+
 pub struct ExchangeClient {
     tx: ByteBoundedSender,
 }
@@ -484,9 +524,11 @@ pub struct ExchangeClients {
     /// tries to send data to one.
     listener: OnceCell<Option<ExchangeListener>>,
 
-    /// Maps from a range of worker IDs to the RPC client used to contact those
+    /// Maps from a range of worker IDs to the RPC clients used to contact those
     /// workers.  Only worker IDs for remote workers appear in the map.
-    clients: Vec<(Host, OnceCell<ExchangeClient>)>,
+    ///
+    /// We use one RPC client per [MessageType] per [Host].
+    clients: Vec<(Host, EnumMap<MessageType, OnceCell<ExchangeClient>>)>,
 }
 
 impl ExchangeClients {
@@ -509,14 +551,14 @@ impl ExchangeClients {
             clients: runtime
                 .layout()
                 .other_hosts()
-                .map(|host| (host.clone(), OnceCell::new()))
+                .map(|host| (host.clone(), Default::default()))
                 .collect(),
         }
     }
 
     /// Returns a client for `worker`, which must be a remote worker ID, first
     /// establishing a connection if there isn't one yet.
-    pub async fn connect(&self, worker: usize) -> &ExchangeClient {
+    pub async fn connect(&self, worker: usize, message_type: MessageType) -> &ExchangeClient {
         self.listener
             .get_or_init(|| async {
                 if let Some(runtime) = self.runtime.upgrade()
@@ -540,7 +582,8 @@ impl ExchangeClients {
             .iter()
             .find(|(host, _client)| host.workers.contains(&worker))
             .unwrap();
-        cell.get_or_init(|| ExchangeClient::new(host.address, &host.workers))
+        cell[message_type]
+            .get_or_init(|| ExchangeClient::new(host.address, &host.workers))
             .await
     }
 }
@@ -963,16 +1006,19 @@ where
                         })
                         .collect_vec();
 
+                    let message_type = MessageType::Synchronous(MessageSize::from_bytes(
+                        items.iter().map(|fbuf| fbuf.len()).sum(),
+                    ));
+
                     // We discard the return value that could allow us to wait
                     // for the channel tx buffer to drain, because exchange is
                     // synchronous, meaning that it will drain before we send
                     // the next message.
-                    let _ = self.clients.connect(receivers.start).await.send(
-                        global_node_id.clone(),
-                        self.exchange_id,
-                        sender,
-                        items,
-                    );
+                    let _ = self
+                        .clients
+                        .connect(receivers.start, message_type)
+                        .await
+                        .send(global_node_id.clone(), self.exchange_id, sender, items);
                 }
             }
         }
diff --git a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
index 5f6f1ee62e1..51ba069782b 100644
--- a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
+++ b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
@@ -28,7 +28,8 @@ use crate::{
     circuit_cache_key,
     operator::{
         communication::{
-            ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, pop_flushed,
+            ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, MessageType,
+            pop_flushed,
         },
         dynamic::{
             accumulator::{Accumulation, EnableCount},
@@ -312,12 +313,12 @@ where
                             item.push(flush as u8);
                         }
                         let this = self.clone();
-                        if let Some(waiter) = this.clients.connect(receivers.start).await.send(
-                            name.clone(),
-                            this.exchange_id,
-                            sender,
-                            items,
-                        ) {
+                        if let Some(waiter) = this
+                            .clients
+                            .connect(receivers.start, MessageType::Streaming)
+                            .await
+                            .send(name.clone(), this.exchange_id, sender, items)
+                        {
                             remote_waiters.push(waiter);
                         }
                     }

From e3cfded90ba58c617fd4fbc86d1e400b9aefe740 Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@feldera.com>
Date: Wed, 17 Jun 2026 11:18:25 -0700
Subject: [PATCH 6/8] [dbsp] Simplify
 `ChildCircuit::add_exchange_with_preference`.

This should be equivalent code, just shorter.  Just noticed this while
reading code, there is nothing important here.

Signed-off-by: Ben Pfaff <blp@feldera.com>
---
 crates/dbsp/src/circuit/circuit_builder.rs | 31 +++++-----------------
 1 file changed, 6 insertions(+), 25 deletions(-)

diff --git a/crates/dbsp/src/circuit/circuit_builder.rs b/crates/dbsp/src/circuit/circuit_builder.rs
index 91f59bccf5a..e07d24b8860 100644
--- a/crates/dbsp/src/circuit/circuit_builder.rs
+++ b/crates/dbsp/src/circuit/circuit_builder.rs
@@ -3939,31 +3939,12 @@ where
         SndOp: SinkOperator<I>,
         RcvOp: SourceOperator<O>,
     {
-        let sender_id = self.add_node(|id| {
-            self.log_circuit_event(&CircuitEvent::operator(
-                GlobalNodeId::child_of(self, id),
-                sender.name(),
-                sender.location(),
-            ));
-
-            let node = SinkNode::new(sender, input_stream.clone(), self.clone(), id);
-            self.connect_stream(input_stream, id, input_preference);
-            (node, id)
-        });
-
-        let output_stream = self.add_node(|id| {
-            self.log_circuit_event(&CircuitEvent::operator(
-                GlobalNodeId::child_of(self, id),
-                receiver.name(),
-                receiver.location(),
-            ));
-
-            let node = SourceNode::new(receiver, self.clone(), id);
-            let output_stream = node.output_stream();
-            (node, output_stream)
-        });
-
-        self.add_dependency(sender_id, output_stream.local_node_id());
+        let sender_id = self.add_sink_with_preference(sender, input_stream, input_preference);
+        let output_stream = self.add_source(receiver);
+        self.add_dependency(
+            sender_id.local_node_id().unwrap(),
+            output_stream.local_node_id(),
+        );
         output_stream
     }
 

From b1474fd2443e85daf1631d149051f7185f96c6cc Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@feldera.com>
Date: Wed, 17 Jun 2026 12:59:20 -0700
Subject: [PATCH 7/8] [dbsp] Break sharded accumulator waiting for merging into
 separate operator.

Before this commit, the sharded accumulator waited for spines to merge in
the ShardedAccumulatorSender operator before it returned a value from its
eval() method.  This prevented any other work from happening in the
receiver and any other part of the circuit downstream from the receiver,
because there is an artificial dependency between the sender and receiver
created by add_exchange_operator().

One way to solve the problem is to remove that dependency (it is not
essential for streaming exchange) but it still makes some sense to have it;
it does make sense to receive after sending.  Another reason is that
ExchangeActivity::InputOnly currently expects all of the workers to flush
in the same step, but without the dependency I am not confident without
further thought that that would still be the case, since it would become
nondeterministic which operator (sender or receiver) would run first.

Instead, this commit breaks off the waiting part of
ShardedAccumulatorSender into a new operator ShardedAccumulatorWaiter,
which just does the wait.  It adds a dependency on
ShardedAccumulatorReceiver for that operator, so that it runs after
everything has been sent and received in a step.  The new waiter operator
does no work, it only keeps the step from finishing before the spine used
for accumulating data has merged to an acceptable number of batches.  This
way, the sender runs quickly and the wait does not block the receiver.

Remote waiters are a different problem.  This commit handles remote
waiters by adding a single operator to wait for all of them at the very
end of a step (by making the operator depend on all the sending
operators).

This isn't a problem in synchronous exchange (or perhaps it is a
different problem) because there it is the receiver, not the sender,
that adds the batch to the spine.

Signed-off-by: Ben Pfaff <blp@feldera.com>
---
 .../src/operator/communication/exchange.rs    |  14 ++
 .../operator/dynamic/sharded_accumulator.rs   | 183 ++++++++++++++----
 2 files changed, 164 insertions(+), 33 deletions(-)

diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs
index 23197f55d50..08fa320fc54 100644
--- a/crates/dbsp/src/operator/communication/exchange.rs
+++ b/crates/dbsp/src/operator/communication/exchange.rs
@@ -382,6 +382,10 @@ impl ExchangeClient {
             })
             .expect("remote exchange failed")
     }
+
+    pub async fn wait(&self) {
+        self.tx.bound.wait().await;
+    }
 }
 
 /// Uniquely identifies an `Exchange` or `ShardedAccumulator` within a circuit.
@@ -586,6 +590,16 @@ impl ExchangeClients {
             .get_or_init(|| ExchangeClient::new(host.address, &host.workers))
             .await
     }
+
+    pub async fn wait(&self) {
+        for (_, clients) in &self.clients {
+            for client in clients.values() {
+                if let Some(client) = client.get() {
+                    client.wait().await;
+                }
+            }
+        }
+    }
 }
 
 struct CallbackInner {
diff --git a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
index 51ba069782b..38b066325de 100644
--- a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
+++ b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs
@@ -6,6 +6,7 @@ use std::{
     panic::Location,
     pin::Pin,
     sync::{Arc, Mutex, MutexGuard},
+    time::Instant,
 };
 
 use feldera_samply::Span;
@@ -16,7 +17,7 @@ use size_of::{HumanBytes, SizeOf, TotalSize};
 use crate::{
     Circuit, NumEntries, Runtime, Scope, Stream,
     circuit::{
-        GlobalNodeId, OwnershipPreference, StepSize, WorkerLocation, WorkerLocations,
+        GlobalNodeId, NodeId, OwnershipPreference, StepSize, WorkerLocation, WorkerLocations,
         circuit_builder::StreamId,
         metadata::{
             ALLOCATED_MEMORY_BYTES, BatchSizeStats, INPUT_BATCHES_STATS, MEMORY_ALLOCATIONS_COUNT,
@@ -43,6 +44,8 @@ circuit_cache_key!(local StreamingExchangeCacheId<B: Batch>(ExchangeId => Arc<Sh
 
 circuit_cache_key!(ShardedAccumulatorId<C, B: Batch>((StreamId, Range<usize>) => Accumulation<Stream<C, Option<Spine<B>>>>));
 
+circuit_cache_key!(ShardedAccumulatorRemoteWaiterId(() => NodeId));
+
 impl<C, B> Stream<C, B>
 where
     C: Circuit,
@@ -76,6 +79,21 @@ where
             && runtime.layout().n_workers() > 1
             && runtime.get_step_size() == StepSize::Microsteps
         {
+            let remote_waiter_node_id = if runtime.layout().is_multihost() {
+                let clients = ExchangeClients::for_runtime(&runtime);
+                Some(*self.circuit().cache_get_or_insert_with(
+                    ShardedAccumulatorRemoteWaiterId::new(()),
+                    move || {
+                        let waiter = self
+                            .circuit()
+                            .add_source(ShardedAccumulatorRemoteWaiter::new(clients.clone()));
+                        waiter.local_node_id()
+                    },
+                ))
+            } else {
+                None
+            };
+
             self.circuit()
                 .cache_get_or_insert_with(
                     ShardedAccumulatorId::new((self.stream_id(), workers.clone())),
@@ -89,7 +107,13 @@ where
                             factories,
                         );
                         let enable_count = exchange.enable_count.clone();
-                        let stream = self
+                        let local_waiter =
+                            self.circuit()
+                                .add_source(ShardedAccumulatorLocalWaiter::new(
+                                    Some(Location::caller()),
+                                    exchange.clone(),
+                                ));
+                        let receiver = self
                             .circuit()
                             .add_exchange(
                                 ShardedAccumulatorSender::new(
@@ -100,8 +124,14 @@ where
                                 self,
                             )
                             .mark_sharded_workers(workers.clone());
+                        self.circuit()
+                            .add_dependency(receiver.local_node_id(), local_waiter.local_node_id());
+                        if let Some(remote_waiter_node_id) = remote_waiter_node_id {
+                            self.circuit()
+                                .add_dependency(receiver.local_node_id(), remote_waiter_node_id);
+                        }
                         Accumulation {
-                            stream,
+                            stream: receiver,
                             enable_count,
                         }
                     },
@@ -135,6 +165,7 @@ where
     /// The RPC clients to contact remote hosts.
     clients: Arc<ExchangeClients>,
 
+    /// One [Rxq] for each of `local_workers`.
     rxq: Vec<Mutex<Rxq<B>>>,
 
     enable_count: EnableCount,
@@ -261,7 +292,6 @@ where
         let worker_locations = WorkerLocations::for_layout(layout);
         let mut data = batches.into_iter();
         let mut remote_waiters = Vec::new();
-        let mut local_waiters = Vec::new();
         let mut serialized_bytes = 0;
         for receivers in layout.all_hosts() {
             match worker_locations[receivers.start] {
@@ -272,18 +302,7 @@ where
                             .expect("data should include one item per peer")
                             .into_plain()
                             .expect("local data should not be serialized");
-                        if self.deliver(&self.factories, sender, receiver, item, flush)
-                            && !flush
-                            && let Some(waiter) = self
-                                .rxq(receiver)
-                                .spines
-                                .back()
-                                .unwrap()
-                                .spine
-                                .backpressure_waiter()
-                        {
-                            local_waiters.push((receiver, waiter));
-                        }
+                        self.deliver(&self.factories, sender, receiver, item, flush);
                     }
                 }
                 WorkerLocation::Remote => {
@@ -326,23 +345,6 @@ where
             }
         }
 
-        if !local_waiters.is_empty() {
-            let _span = Span::new("local send wait")
-                .with_category("Exchange")
-                .with_tooltip(|| {
-                    format!(
-                        "{name} wait for batches to merge in {} receive queues (for workers {})",
-                        local_waiters.len(),
-                        local_waiters
-                            .iter()
-                            .map(|(receiver, _waiter)| receiver)
-                            .format(", ")
-                    )
-                });
-            for (_receiver, waiter) in local_waiters {
-                waiter.await;
-            }
-        }
         if !remote_waiters.is_empty() {
             let _span = Span::new("remote send wait")
                 .with_category("Exchange")
@@ -372,6 +374,43 @@ where
     }
 }
 
+impl<B> ShardedAccumulator<B>
+where
+    B: Batch,
+{
+    async fn wait(&self, name: Arc<String>) {
+        let start = Instant::now();
+        let mut local_waiters = Vec::new();
+        for (rxq, worker) in self.rxq.iter().zip(self.local_workers.clone()) {
+            // This is intentionally two separate statements to avoid holding
+            // the lock while waiting.
+            let waiter = rxq
+                .lock()
+                .unwrap()
+                .spines
+                .front()
+                .and_then(|entry| entry.spine.backpressure_waiter());
+            if let Some(waiter) = waiter {
+                local_waiters.push(worker);
+                waiter.await;
+            }
+        }
+        if !local_waiters.is_empty() {
+            Span::new("local send wait")
+                .with_start(start)
+                .with_category("Exchange")
+                .with_tooltip(|| {
+                    format!(
+                        "{name} wait for batches to merge in {} receive queues (for workers {})",
+                        local_waiters.len(),
+                        local_waiters.iter().format(", ")
+                    )
+                })
+                .record();
+        }
+    }
+}
+
 impl<B> ExchangeDelivery for ShardedAccumulator<B>
 where
     B: Batch<Time = ()>,
@@ -664,6 +703,84 @@ where
     }
 }
 
+struct ShardedAccumulatorRemoteWaiter {
+    clients: Arc<ExchangeClients>,
+}
+
+impl Operator for ShardedAccumulatorRemoteWaiter {
+    fn name(&self) -> std::borrow::Cow<'static, str> {
+        Cow::Borrowed("ShardedAccumulatorRemoteWaiter")
+    }
+
+    fn fixedpoint(&self, _scope: crate::circuit::Scope) -> bool {
+        true
+    }
+}
+
+impl SourceOperator<()> for ShardedAccumulatorRemoteWaiter {
+    async fn eval(&mut self) {
+        self.clients.wait().await;
+    }
+}
+
+impl ShardedAccumulatorRemoteWaiter {
+    fn new(clients: Arc<ExchangeClients>) -> Self {
+        Self { clients }
+    }
+}
+
+struct ShardedAccumulatorLocalWaiter<B>
+where
+    B: Batch,
+{
+    exchange: Arc<ShardedAccumulator<B>>,
+    location: OperatorLocation,
+    name: OperatorName,
+}
+
+impl<B> ShardedAccumulatorLocalWaiter<B>
+where
+    B: Batch,
+{
+    fn new(location: OperatorLocation, exchange: Arc<ShardedAccumulator<B>>) -> Self {
+        Self {
+            exchange,
+            location,
+            name: OperatorName::new("ShardedAccumulatorLocalWaiter"),
+        }
+    }
+}
+
+impl<B> Operator for ShardedAccumulatorLocalWaiter<B>
+where
+    B: Batch,
+{
+    fn name(&self) -> std::borrow::Cow<'static, str> {
+        Cow::Borrowed("ShardedAccumulatorLocalWaiter")
+    }
+
+    fn location(&self) -> OperatorLocation {
+        self.location
+    }
+
+    fn init(&mut self, global_id: &GlobalNodeId) {
+        self.name.init(global_id);
+    }
+
+    fn fixedpoint(&self, _scope: crate::circuit::Scope) -> bool {
+        true
+    }
+}
+
+impl<B> SourceOperator<()> for ShardedAccumulatorLocalWaiter<B>
+where
+    B: Batch,
+{
+    async fn eval(&mut self) {
+        self.exchange.wait(self.name.get()).await;
+    }
+}
+
 struct ShardedAccumulatorReceiver<B>
 where
     B: Batch,

From c87d78ee543c8b78388b371ab963fda1ac1945c8 Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@feldera.com>
Date: Wed, 15 Jul 2026 13:46:20 -0700
Subject: [PATCH 8/8] [dbsp] Improve comment.

Signed-off-by: Ben Pfaff <blp@feldera.com>
---
 crates/dbsp/src/operator/communication/exchange.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs
index 08fa320fc54..71f22513b01 100644
--- a/crates/dbsp/src/operator/communication/exchange.rs
+++ b/crates/dbsp/src/operator/communication/exchange.rs
@@ -1043,6 +1043,8 @@ where
     ///
     /// When the data is ready, but before reading it, this method swaps
     /// `start_wait_usecs` with 0 and returns the old value along with the data.
+    /// This allows the caller to obtain the waiting time incurred just after it
+    /// became ready.
     pub(crate) async fn receive_all<D>(
         &self,
         deserialize: D,