From 18f295c6a336a4c7685a3a89409c3593aa31b6b2 Mon Sep 17 00:00:00 2001
From: Moritz Hoffmann <antiguru@gmail.com>
Date: Tue, 29 Sep 2026 10:48:02 +0200
Subject: [PATCH] Build an operator's trace with a caller-supplied factory

`arrange_core` and `reduce_with_tactic` construct their trace with
`Trace::new` and move it into a `TraceAgent`, after which the only way to
reach the trace instance is `TraceAgent::trace_box_unstable`. A caller whose
trace shares state with something outside the operator, for example a
handle through which other threads observe its batches, has no stable way to
obtain that state.

`arrange_core_with_trace` and `reduce_with_tactic_and_trace` take a factory
with the signature of `Trace::new` and call it once, where the operator
builds its trace. The caller can take a handle from the trace before
returning it. `arrange_core` and `reduce_with_tactic` delegate with
`Trace::new`, so existing callers are unchanged.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015tLhSbZdXrTSK2KwSocT59
---
 differential-dataflow/CHANGELOG.md            |  4 +
 .../src/operators/arrange/arrangement.rs      | 26 ++++++-
 differential-dataflow/src/operators/reduce.rs | 27 ++++++-
 differential-dataflow/tests/trace_factory.rs  | 75 +++++++++++++++++++
 4 files changed, 128 insertions(+), 4 deletions(-)
 create mode 100644 differential-dataflow/tests/trace_factory.rs

diff --git a/differential-dataflow/CHANGELOG.md b/differential-dataflow/CHANGELOG.md
index 256bb80d2..b619cb213 100644
--- a/differential-dataflow/CHANGELOG.md
+++ b/differential-dataflow/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- `arrange_core_with_trace` and `reduce_with_tactic_and_trace`, which build the operator's trace with a caller-supplied factory in place of `Trace::new`, so a caller can keep a handle into the trace before `TraceAgent` takes ownership of it.
+
 ## [0.25.1](https://github.com/TimelyDataflow/differential-dataflow/compare/differential-dataflow-v0.25.0...differential-dataflow-v0.25.1) - 2026-07-15
 
 ### Other
diff --git a/differential-dataflow/src/operators/arrange/arrangement.rs b/differential-dataflow/src/operators/arrange/arrangement.rs
index 457b760d9..83af43648 100644
--- a/differential-dataflow/src/operators/arrange/arrangement.rs
+++ b/differential-dataflow/src/operators/arrange/arrangement.rs
@@ -20,7 +20,8 @@
 use timely::dataflow::operators::{Enter, vec::Map};
 use timely::order::PartialOrder;
 use timely::dataflow::{Scope, Stream};
-use timely::dataflow::operators::generic::Operator;
+use timely::dataflow::operators::generic::{Operator, OperatorInfo};
+use timely::scheduling::Activator;
 use timely::dataflow::channels::pact::{ParallelizationContract, Pipeline};
 use timely::progress::Timestamp;
 use timely::progress::Antichain;
@@ -310,6 +311,27 @@ pub fn arrange_core<'scope, P, C, Ba, Tr>(
     name: &str,
     batcher: impl FnOnce(Option<Logger>, usize) -> Ba,
 ) -> Arranged<'scope, TraceAgent<Tr>>
+where
+    C: Container + Clone + 'static,
+    P: ParallelizationContract<Tr::Time, C>,
+    Ba: Batcher<C, Time = Tr::Time, Output: Into<Tr::Batch>> + 'static,
+    Tr: Trace+'static,
+{
+    arrange_core_with_trace(stream, pact, name, batcher, Tr::new)
+}
+
+/// Arranges a stream of updates like [`arrange_core`], building the trace with `trace`.
+///
+/// `trace` receives the arguments [`Trace::new`] would, and is called once, on the operator's
+/// construction. A caller that needs to reach the trace after [`TraceAgent`] takes ownership of it,
+/// for example through a handle the trace shares, can obtain that handle here.
+pub fn arrange_core_with_trace<'scope, P, C, Ba, Tr>(
+    stream: Stream<'scope, Tr::Time, C>,
+    pact: P,
+    name: &str,
+    batcher: impl FnOnce(Option<Logger>, usize) -> Ba,
+    trace: impl FnOnce(OperatorInfo, Option<Logger>, Option<Activator>) -> Tr,
+) -> Arranged<'scope, TraceAgent<Tr>>
 where
     C: Container + Clone + 'static,
     P: ParallelizationContract<Tr::Time, C>,
@@ -349,7 +371,7 @@ where
         let mut capabilities = Antichain::<Capability<Tr::Time>>::new();
 
         let activator = Some(scope.activator_for(std::rc::Rc::clone(&info.address)));
-        let mut empty_trace = Tr::new(info.clone(), logger.clone(), activator);
+        let mut empty_trace = trace(info.clone(), logger.clone(), activator);
         // If there is default exertion logic set, install it.
         if let Some(exert_logic) = scope.worker().config().get::<trace::ExertionLogic>("differential/default_exert_logic").cloned() {
             empty_trace.set_exert_logic(exert_logic);
diff --git a/differential-dataflow/src/operators/reduce.rs b/differential-dataflow/src/operators/reduce.rs
index e93f54826..24d0ef438 100644
--- a/differential-dataflow/src/operators/reduce.rs
+++ b/differential-dataflow/src/operators/reduce.rs
@@ -14,8 +14,11 @@ use timely::progress::Timestamp;
 use timely::dataflow::operators::Operator;
 use timely::dataflow::operators::CapabilitySet;
 use timely::dataflow::channels::pact::Pipeline;
+use timely::dataflow::operators::generic::OperatorInfo;
+use timely::scheduling::Activator;
 
 use crate::operators::arrange::{Arranged, TraceAgent};
+use crate::logging::Logger;
 use crate::trace::{Span, ExertionLogic, Trace, TraceReader};
 
 /// Sort and deduplicate a list. Shared by the cursor and proxy tactics, which
@@ -75,7 +78,27 @@ pub use crate::operators::cursor::reduce::reduce_trace;
 /// `TraceReader` of its input and `Trace` of its output, never `Navigable`: it extracts batches via
 /// `spans_through`, and building cursors over them (if that is how the reduce proceeds) is the
 /// tactic's concern.
-pub fn reduce_with_tactic<'scope, Tr1, Tr2, T>(trace: Arranged<'scope, Tr1>, name: &str, mut tactic: T) -> Arranged<'scope, TraceAgent<Tr2>>
+pub fn reduce_with_tactic<'scope, Tr1, Tr2, T>(trace: Arranged<'scope, Tr1>, name: &str, tactic: T) -> Arranged<'scope, TraceAgent<Tr2>>
+where
+    Tr1: TraceReader + 'static,
+    Tr2: Trace<Time = Tr1::Time> + 'static,
+    T: ReduceTactic<Tr1::Time, Tr1::Batch, Tr2::Batch> + 'static,
+{
+    reduce_with_tactic_and_trace(trace, name, tactic, Tr2::new)
+}
+
+/// Drives a key-wise reduction like [`reduce_with_tactic`], building the output trace with
+/// `output`.
+///
+/// `output` receives the arguments [`Trace::new`] would, and is called once, on the operator's
+/// construction. See [`crate::operators::arrange::arrangement::arrange_core_with_trace`] for why a
+/// caller would supply it.
+pub fn reduce_with_tactic_and_trace<'scope, Tr1, Tr2, T>(
+    trace: Arranged<'scope, Tr1>,
+    name: &str,
+    mut tactic: T,
+    output: impl FnOnce(OperatorInfo, Option<Logger>, Option<Activator>) -> Tr2,
+) -> Arranged<'scope, TraceAgent<Tr2>>
 where
     Tr1: TraceReader + 'static,
     Tr2: Trace<Time = Tr1::Time> + 'static,
@@ -95,7 +118,7 @@ where
             let logger = scope.worker().logger_for::<crate::logging::DifferentialEventBuilder>("differential/arrange").map(Into::into);
 
             let activator = Some(scope.activator_for(std::rc::Rc::clone(&operator_info.address)));
-            let mut empty = Tr2::new(operator_info.clone(), logger.clone(), activator);
+            let mut empty = output(operator_info.clone(), logger.clone(), activator);
             // If there is default exert logic set, install it.
             if let Some(exert_logic) = scope.worker().config().get::<ExertionLogic>("differential/default_exert_logic").cloned() {
                 empty.set_exert_logic(exert_logic);
diff --git a/differential-dataflow/tests/trace_factory.rs b/differential-dataflow/tests/trace_factory.rs
new file mode 100644
index 000000000..86371832f
--- /dev/null
+++ b/differential-dataflow/tests/trace_factory.rs
@@ -0,0 +1,75 @@
+//! The `_with_trace` operator variants build their trace with the supplied factory.
+
+use std::cell::Cell;
+use std::rc::Rc;
+
+use timely::dataflow::channels::pact::Exchange;
+use timely::dataflow::operators::ToStream;
+
+use differential_dataflow::AsCollection;
+use differential_dataflow::hashable::Hashable;
+use differential_dataflow::operators::arrange::arrangement::arrange_core_with_trace;
+use differential_dataflow::operators::cursor::reduce::CursorTactic;
+use differential_dataflow::operators::reduce::reduce_with_tactic_and_trace;
+use differential_dataflow::trace::Trace;
+use differential_dataflow::trace::implementations::{ValBatcher, ValBuilder, ValSpine};
+
+type Spine = ValSpine<u64, u64, u64, isize>;
+
+#[test]
+fn arrange_builds_its_trace_with_the_factory() {
+    timely::example(|scope| {
+        let calls = Rc::new(Cell::new(0));
+        let built_for = Rc::new(Cell::new(None));
+        let (calls_in, built_for_in) = (Rc::clone(&calls), Rc::clone(&built_for));
+        let stream = vec![((1u64, 2u64), 0u64, 1isize)].to_stream(scope);
+        let exchange = Exchange::new(|update: &((u64, u64), u64, isize)| (update.0).0.hashed().into());
+        let arranged = arrange_core_with_trace::<_, _, _, Spine>(
+            stream,
+            exchange,
+            "Arrange",
+            ValBatcher::new,
+            move |info, logger, activator| {
+                calls_in.set(calls_in.get() + 1);
+                built_for_in.set(Some(info.global_id));
+                Spine::new(info, logger, activator)
+            },
+        );
+        assert_eq!(calls.get(), 1);
+        assert_eq!(built_for.get(), Some(arranged.trace.operator().global_id));
+    });
+}
+
+#[test]
+fn reduce_builds_its_output_trace_with_the_factory() {
+    timely::example(|scope| {
+        let calls = Rc::new(Cell::new(0));
+        let built_for = Rc::new(Cell::new(None));
+        let (calls_in, built_for_in) = (Rc::clone(&calls), Rc::clone(&built_for));
+        let input = vec![((1u64, 2u64), 0u64, 1isize)]
+            .to_stream(scope)
+            .as_collection()
+            .arrange_by_key();
+        let tactic = CursorTactic::<_, _, ValBuilder<u64, u64, u64, isize>, _, _>::new(
+            |_key: &u64, input: &[(&u64, isize)], _output: &mut Vec<(u64, isize)>, change: &mut Vec<(u64, isize)>| {
+                change.extend(input.iter().map(|(v, r)| (**v, *r)));
+            },
+            |vec: &mut Vec<((u64, u64), u64, isize)>, key: &u64, upds: &mut Vec<(u64, u64, isize)>| {
+                vec.clear();
+                vec.extend(upds.drain(..).map(|(v, t, r)| ((*key, v), t, r)));
+            },
+        );
+        let reduced = reduce_with_tactic_and_trace::<_, Spine, _>(
+            input,
+            "Reduce",
+            tactic,
+            move |info, logger, activator| {
+                calls_in.set(calls_in.get() + 1);
+                built_for_in.set(Some(info.global_id));
+                Spine::new(info, logger, activator)
+            },
+        );
+        assert_eq!(calls.get(), 1);
+        assert_eq!(built_for.get(), Some(reduced.trace.operator().global_id));
+    });
+}