feldera · ryzhyk · Jul 15, 2026 · Jul 3, 2026 · mythical-fred · Jul 3, 2026
diff --git a/crates/dbsp/src/operator/output.rs b/crates/dbsp/src/operator/output.rs
@@ -21,7 +21,8 @@ use std::{
     hash::{Hash, Hasher},
     marker::PhantomData,
     mem::transmute,
-    sync::Arc,
+    ops::Range,
+    sync::{Arc, Mutex},
 };
 use typedmap::TypedMapKey;
 
@@ -200,6 +201,97 @@ where
     type Value = OutputHandle<T>;
 }
 
+/// `TypedMapKey` entry used to share the pending-cohort slots of an
+/// [`AccumulateOutput`] operator across the workers of one host.
+struct CohortId<T> {
+    id: usize,
+    // `fn() -> T` keeps the key `Send + Sync` regardless of `T`: the key
+    // names a type but never holds a value of it.
+    _marker: PhantomData<fn() -> T>,
+}
+
+// Implement `Hash`, `Eq` manually to avoid `T: Hash` type bound.
+impl<T> Hash for CohortId<T> {
+    fn hash<H>(&self, state: &mut H)
+    where
+        H: Hasher,
+    {
+        self.id.hash(state);
+    }
+}
+
+impl<T> PartialEq for CohortId<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id
+    }
+}
+
+impl<T> Eq for CohortId<T> {}
+
+impl<T> CohortId<T> {
+    fn new(id: usize) -> Self {
+        Self {
+            id,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<T> TypedMapKey<LocalStoreMarker> for CohortId<T>
+where
+    T: 'static,
+{
+    type Value = CohortSlots<T>;
+}
+
+/// Pending per-worker outputs for the transaction in progress.
+///
+/// Slot `i` stores the output of the host's `i`-th worker until every worker
+/// on this host has produced its output for the transaction. Remote workers
+/// have no slots: they cannot reach this host's store and publish through
+/// their own host's instance.
+struct CohortSlots<T>(Arc<Mutex<Vec<Option<T>>>>);
+
+impl<T> Clone for CohortSlots<T> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+impl<T> CohortSlots<T> {
+    fn new(num_workers: usize) -> Self {
+        assert_ne!(num_workers, 0);
+        Self(Arc::new(Mutex::new(
+            (0..num_workers).map(|_| None).collect(),
+        )))
+    }
+
+    /// Store `value` in `slot` as its worker's output for the transaction in
+    /// progress.
+    ///
+    /// Returns the complete cohort, one value per slot in slot order, when
+    /// `value` fills the last empty slot; returns `None` otherwise.
+    fn put(&self, slot: usize, value: T) -> Option<Vec<T>> {
+        let mut pending = self.0.lock().unwrap();
+
+        // The upstream accumulator emits exactly once per transaction per
+        // worker, so the slot must be free.
+        debug_assert!(pending[slot].is_none());
+        pending[slot] = Some(value);
+
+        if pending.iter().all(Option::is_some) {
+            Some(
+                pending
+                    .iter_mut()
+                    .map(|slot| slot.take().unwrap())
+                    .collect(),
+            )
+        } else {
+            None
+        }
+    }
+}
+
 struct OutputHandleInternal<T> {
     mailbox: Vec<Mailbox<Option<T>>>,
 }
@@ -479,12 +571,34 @@ where
     }
 }
 
+/// Sink operator that publishes an accumulated stream through an
+/// [`OutputHandle`].
+///
+/// Every worker's accumulator emits exactly once per transaction, but when a
+/// transaction commit spans several steps, workers can emit in different
+/// steps. Writing each emission to its mailbox immediately would let a reader
+/// observe a partial cross-worker output for a transaction. Instead, the
+/// operator parks emissions in [`CohortSlots`] shared by the workers of one
+/// host and publishes the whole cohort to the mailboxes when the last of
+/// these workers emits, so a reader sees either all of a transaction's
+/// outputs or none.
+///
+/// In a multihost runtime, the cohort covers one host: each host's workers
+/// share host-local [`CohortSlots`] and publish to host-local mailboxes,
+/// which remote workers cannot reach. A reader observes a host's outputs
+/// atomically; distinct hosts publish independently.
 pub struct AccumulateOutput<B>
 where
     B: Batch,
 {
     global_id: GlobalNodeId,
-    mailbox: Mailbox<Option<SpineSnapshot<B>>>,
+    handle: OutputHandle<SpineSnapshot<B>>,
+    /// Global indices of this host's workers: cohort slot `i` belongs to
+    /// worker `local_workers.start + i`.
+    local_workers: Range<usize>,
+    /// This worker's cohort slot: its offset within `local_workers`.
+    slot: usize,
+    pending: CohortSlots<SpineSnapshot<B>>,
     output_batch_stats: BatchSizeStats,
 }
 
@@ -494,17 +608,51 @@ where
 {
     pub fn new() -> (Self, OutputHandle<SpineSnapshot<B>>) {
         let handle = OutputHandle::new();
-        let mailbox = handle.mailbox(Runtime::worker_index()).clone();
+
+        // The slots live in the host-local store, out of reach of remote
+        // workers, so the cohort covers only this host's workers.
+        let (local_workers, pending) = match Runtime::runtime() {
+            None => (0..1, CohortSlots::new(1)),
+            Some(runtime) => {
+                let local_workers = runtime.layout().local_workers();
+                let cohort_id = runtime.sequence_next();
+
+                let pending = runtime
+                    .local_store()
+                    .entry(CohortId::new(cohort_id))
+                    .or_insert_with(|| CohortSlots::new(local_workers.len()))
+                    .value()
+                    .clone();
+
+                (local_workers, pending)
+            }
+        };
 
         let output = Self {
             global_id: GlobalNodeId::root(),
-            mailbox,
+            handle: handle.clone(),
+            local_workers,
+            slot: Runtime::local_worker_offset(),
+            pending,
             output_batch_stats: BatchSizeStats::new(),
         };
 
         (output, handle)
     }
 
+    /// Write this worker's output; publish the cohort if it is now complete.
+    ///
+    /// Only the worker that completes a cohort publishes it, so there is
+    /// never more than one publisher per cohort. Mailboxes are indexed by
+    /// global worker index, slots by local offset.
+    fn put_and_publish(&self, snapshot: SpineSnapshot<B>) {
+        if let Some(cohort) = self.pending.put(self.slot, snapshot) {
+            for (worker, snapshot) in self.local_workers.clone().zip(cohort) {
+                self.handle.mailbox(worker).set(Some(snapshot));
+            }
+        }
+    }
+
     fn checkpoint_file(base: &StoragePath, persistent_id: &str) -> StoragePath {
         base.child(format!("accumulate-output-{}.dat", persistent_id))
     }
@@ -567,14 +715,16 @@ where
     async fn eval(&mut self, val: &Option<Spine<B>>) {
         if let Some(val) = val {
             self.output_batch_stats.add_batch(val.len());
-            self.mailbox.set(Some(val.ro_snapshot()));
+            // Park even empty outputs: cohort completion requires one
+            // emission per worker per transaction.
+            self.put_and_publish(val.ro_snapshot());
         }
     }
 
     async fn eval_owned(&mut self, val: Option<Spine<B>>) {
         if let Some(val) = val {
             self.output_batch_stats.add_batch(val.len());
-            self.mailbox.set(Some(val.ro_snapshot()));
+            self.put_and_publish(val.ro_snapshot());
         }
     }
 
@@ -687,8 +837,97 @@ where
 
 #[cfg(test)]
 mod test {
+    use super::CohortSlots;
     use crate::{Runtime, typed_batch::OrdZSet, utils::Tup2};
 
+    /// Parking releases a cohort only when every worker has contributed, and
+    /// the slots are reusable for the next cohort.
+    #[test]
+    fn test_cohort_slots() {
+        let slots = CohortSlots::new(3);
+
+        // Partial cohort: nothing is released.
+        assert_eq!(slots.put(0, "a0"), None);
+        assert_eq!(slots.put(2, "a2"), None);
+
+        // The last worker's output releases the whole cohort.
+        assert_eq!(slots.put(1, "a1"), Some(vec!["a0", "a1", "a2"]));
+
+        // The slots are empty again and serve the next cohort.
+        assert_eq!(slots.put(1, "b1"), None);
+        assert_eq!(slots.put(2, "b2"), None);
+        assert_eq!(slots.put(0, "b0"), Some(vec!["b0", "b1", "b2"]));
+    }
+
+    /// End-to-end: the accumulated output surfaces exactly once per
+    /// transaction, as a complete cohort with one output per worker, and
+    /// never mid-transaction or as a partial cohort mid-commit.
+    ///
+    /// Drives transactions with the manual `start_transaction` ->
+    /// `start_commit_transaction` -> `step` API so the handle can be
+    /// inspected after every step, including each step of a multi-step
+    /// commit; `DBSPHandle::transaction` would only return after the commit
+    /// completed, when the cohort is trivially complete.
+    #[test]
+    fn test_accumulate_output() {
+        use crate::trace::BatchReader;
+
+        const WORKERS: usize = 8;
+
+        let (mut dbsp, (input, output)) = Runtime::init_circuit(WORKERS, |circuit| {
+            let (zset, zset_handle) = circuit.add_input_zset::<u64>();
+            let output = zset.accumulate_output();
+
+            Ok((zset_handle, output))
+        })
+        .unwrap();
+
+        for round in 0..20u64 {
+            // Odd rounds run an empty transaction, which must still publish
+            // a complete, empty cohort.
+            let expected = if round % 2 == 0 { 20_000 } else { 0 };
+
+            dbsp.start_transaction().unwrap();
+
+            if expected > 0 {
+                let mut tuples = (round * 100_000..round * 100_000 + 20_000)
+                    .map(|k| Tup2(k, 1i64))
+                    .collect::<Vec<_>>();
+                input.append(&mut tuples);
+            }
+
+            // Steps inside an open transaction publish nothing.
+            dbsp.step().unwrap();
+            assert!(output.take_from_all().is_empty());
+            dbsp.step().unwrap();
+            assert!(output.take_from_all().is_empty());
+
+            dbsp.start_commit_transaction().unwrap();
+
+            // Every commit step yields either nothing or one complete
+            // cohort, and the whole transaction yields exactly one.
+            let mut cohorts = 0;
+            loop {
+                let committed = dbsp.step().unwrap();
+
+                let batches = output.take_from_all();
+                if !batches.is_empty() {
+                    assert_eq!(batches.len(), WORKERS);
+                    assert_eq!(batches.iter().map(|b| b.len()).sum::<usize>(), expected);
+                    cohorts += 1;
+                }
+
+                if committed {
+                    break;
+                }
+            }
+            assert_eq!(cohorts, 1);
+            assert!(output.take_from_all().is_empty());
+        }
+
+        dbsp.kill().unwrap();
+    }
+
     #[test]
     fn test_output_handle() {
         let (mut dbsp, (input, output)) = Runtime::init_circuit(4, |circuit| {