Migrate production spines to Arc batches for cross-thread arrangement sharing

Cargo.toml

@@ -644,6 +644,12 @@ debug-assertions = true
 # merged), after which point it becomes impossible to build that historical
 # version of Materialize.
 [patch.crates-io]
+# Arc'd batches for cross-thread arrangement sharing, from the branch behind
+# TimelyDataflow/differential-dataflow#807. Prototype plumbing: repoint at a
+# released version (and drop the patch) once that PR lands upstream.
+differential-dataflow = { git = "https://github.com/antiguru/differential-dataflow", branch = "claude/spines-differential-arc-j93mho" }
+differential-dogs3 = { git = "https://github.com/antiguru/differential-dataflow", branch = "claude/spines-differential-arc-j93mho" }
+
 # Waiting on https://github.com/sfackler/rust-postgres/pull/752.
 postgres = { git = "https://github.com/MaterializeInc/rust-postgres" }
 tokio-postgres = { git = "https://github.com/MaterializeInc/rust-postgres" }

src/compute/src/extensions/arrange.rs

@@ -8,7 +8,8 @@
 // by the Apache License, Version 2.0.
 
 use std::collections::BTreeMap;
-use std::rc::{Rc, Weak};
+use std::rc::Rc;
+use std::sync::{Arc, Weak};
 
 use differential_dataflow::difference::Semigroup;
 use differential_dataflow::lattice::Lattice;
@@ -241,9 +242,9 @@ pub trait ArrangementSize {
 /// * `logic`: Closure that calculates the heap size/capacity/allocations for a batch. The return
 ///    value are size and capacity in bytes, and number of allocations, all in absolute values.
 fn log_arrangement_size_inner<'scope, B, L>(
-    arranged: Arranged<'scope, TraceAgent<Spine<Rc<B>>>>,
+    arranged: Arranged<'scope, TraceAgent<Spine<Arc<B>>>>,
     mut logic: L,
-) -> Arranged<'scope, TraceAgent<Spine<Rc<B>>>>
+) -> Arranged<'scope, TraceAgent<Spine<Arc<B>>>>
 where
     B: Batch + 'static,
     L: FnMut(&B) -> (usize, usize, usize) + 'static,
@@ -282,8 +283,8 @@ where
                 input.for_each(|time, data| {
                     for batch in data.iter() {
                         batches
-                            .entry(Rc::as_ptr(batch))
-                            .or_insert_with(|| (Rc::downgrade(batch), logic(batch)));
+                            .entry(Arc::as_ptr(batch))
+                            .or_insert_with(|| (Arc::downgrade(batch), logic(batch)));
                     }
                     output.session(&time).give_container(data);
                 });
@@ -293,8 +294,8 @@ where
 
                 trace.borrow().trace().map_batches(|batch| {
                     batches
-                        .entry(Rc::as_ptr(batch))
-                        .or_insert_with(|| (Rc::downgrade(batch), logic(batch)));
+                        .entry(Arc::as_ptr(batch))
+                        .or_insert_with(|| (Arc::downgrade(batch), logic(batch)));
                 });
 
                 let (mut size, mut capacity, mut allocations) = (0, 0, 0);

src/compute/src/typedefs.rs

@@ -29,32 +29,32 @@ pub use crate::typedefs::spines::{ColKeySpine, ColValSpine};
 pub use mz_row_spine::{RowRowSpine, RowSpine, RowValBatcher, RowValSpine};
 
 pub(crate) mod spines {
-    use std::rc::Rc;
+    use std::sync::Arc;
 
     use columnation::Columnation;
     use differential_dataflow::trace::implementations::ord_neu::{
         OrdKeyBatch, OrdKeyBuilder, OrdValBatch, OrdValBuilder,
     };
     use differential_dataflow::trace::implementations::spine_fueled::Spine;
     use differential_dataflow::trace::implementations::{Layout, Update};
-    use differential_dataflow::trace::rc_blanket_impls::RcBuilder;
+    use differential_dataflow::trace::arc_blanket_impls::ArcBuilder;
     use mz_timely_util::columnation::ColumnationStack;
 
     use mz_row_spine::OffsetOptimized;
 
     use crate::typedefs::{KeyBatcher, KeyValBatcher};
 
     /// A spine for generic keys and values.
-    pub type ColValSpine<K, V, T, R> = Spine<Rc<OrdValBatch<MzStack<((K, V), T, R)>>>>;
+    pub type ColValSpine<K, V, T, R> = Spine<Arc<OrdValBatch<MzStack<((K, V), T, R)>>>>;
     pub type ColValBatcher<K, V, T, R> = KeyValBatcher<K, V, T, R>;
     pub type ColValBuilder<K, V, T, R> =
-        RcBuilder<OrdValBuilder<MzStack<((K, V), T, R)>, ColumnationStack<((K, V), T, R)>>>;
+        ArcBuilder<OrdValBuilder<MzStack<((K, V), T, R)>, ColumnationStack<((K, V), T, R)>>>;
 
     /// A spine for generic keys
-    pub type ColKeySpine<K, T, R> = Spine<Rc<OrdKeyBatch<MzStack<((K, ()), T, R)>>>>;
+    pub type ColKeySpine<K, T, R> = Spine<Arc<OrdKeyBatch<MzStack<((K, ()), T, R)>>>>;
     pub type ColKeyBatcher<K, T, R> = KeyBatcher<K, T, R>;
     pub type ColKeyBuilder<K, T, R> =
-        RcBuilder<OrdKeyBuilder<MzStack<((K, ()), T, R)>, ColumnationStack<((K, ()), T, R)>>>;
+        ArcBuilder<OrdKeyBuilder<MzStack<((K, ()), T, R)>, ColumnationStack<((K, ()), T, R)>>>;
 
     /// A layout based on chunked timely stacks
     pub struct MzStack<U: Update> {

src/row-spine/Cargo.toml

@@ -19,5 +19,8 @@ mz-repr = { path = "../repr" }
 mz-timely-util = { path = "../timely-util", default-features = false }
 timely.workspace = true
 
+[dev-dependencies]
+mz-ore = { path = "../ore", default-features = false, features = ["columnation"] }
+
 [features]
 default = ["mz-ore/default", "mz-timely-util/default"]

src/row-spine/src/lib.rs

@@ -29,15 +29,15 @@ pub static DICTIONARY_COMPRESSION: std::sync::atomic::AtomicBool =
 
 /// Spines specialized to contain `Row` types in keys and values.
 mod spines {
-    use std::rc::Rc;
+    use std::sync::Arc;
 
     use columnation::Columnation;
     use differential_dataflow::trace::implementations::Layout;
     use differential_dataflow::trace::implementations::Update;
     use differential_dataflow::trace::implementations::merge_batcher::MergeBatcher;
     use differential_dataflow::trace::implementations::ord_neu::{OrdKeyBatch, OrdValBatch};
     use differential_dataflow::trace::implementations::spine_fueled::Spine;
-    use differential_dataflow::trace::rc_blanket_impls::RcBuilder;
+    use differential_dataflow::trace::arc_blanket_impls::ArcBuilder;
     use mz_repr::Row;
     use mz_timely_util::columnation::{ColInternalMerger, ColumnationStack};
 
@@ -48,9 +48,9 @@ mod spines {
     type KeyValBatcher<K, V, T, D> = MergeBatcher<ColInternalMerger<(K, V), T, D>>;
     type KeyBatcher<K, T, D> = KeyValBatcher<K, (), T, D>;
 
-    pub type RowRowSpine<T, R> = Spine<Rc<OrdValBatch<RowRowLayout<((Row, Row), T, R)>>>>;
+    pub type RowRowSpine<T, R> = Spine<Arc<OrdValBatch<RowRowLayout<((Row, Row), T, R)>>>>;
     pub type RowRowBatcher<T, R> = KeyValBatcher<Row, Row, T, R>;
-    pub type RowRowBuilder<T, R> = RcBuilder<crate::dictionary::builders::RowRowBuilder<T, R>>;
+    pub type RowRowBuilder<T, R> = ArcBuilder<crate::dictionary::builders::RowRowBuilder<T, R>>;
 
     /// `RowRowBuilder` variant that consumes [`Column`] chunks. Pairs with
     /// [`Col2ValPagedBatcher`] for the spillable arrange path. Installs a
@@ -61,21 +61,21 @@ mod spines {
     /// [`Col2ValPagedBatcher`]: mz_timely_util::columnar::Col2ValPagedBatcher
     /// [`Column`]: mz_timely_util::columnar::Column
     pub type RowRowColPagedBuilder<T, R> =
-        RcBuilder<crate::dictionary::builders::RowRowColPagedBuilder<T, R>>;
+        ArcBuilder<crate::dictionary::builders::RowRowColPagedBuilder<T, R>>;
 
-    pub type RowValSpine<V, T, R> = Spine<Rc<OrdValBatch<RowValLayout<((Row, V), T, R)>>>>;
+    pub type RowValSpine<V, T, R> = Spine<Arc<OrdValBatch<RowValLayout<((Row, V), T, R)>>>>;
     pub type RowValBatcher<V, T, R> = KeyValBatcher<Row, V, T, R>;
     pub type RowValBuilder<V, T, R> =
-        RcBuilder<crate::dictionary::builders::RowValBuilder<V, T, R>>;
+        ArcBuilder<crate::dictionary::builders::RowValBuilder<V, T, R>>;
 
-    pub type RowSpine<T, R> = Spine<Rc<OrdKeyBatch<RowLayout<((Row, ()), T, R)>>>>;
+    pub type RowSpine<T, R> = Spine<Arc<OrdKeyBatch<RowLayout<((Row, ()), T, R)>>>>;
     pub type RowBatcher<T, R> = KeyBatcher<Row, T, R>;
-    pub type RowBuilder<T, R> = RcBuilder<crate::dictionary::builders::RowBuilder<T, R>>;
+    pub type RowBuilder<T, R> = ArcBuilder<crate::dictionary::builders::RowBuilder<T, R>>;
 
-    pub type ValRowSpine<K, T, R> = Spine<Rc<OrdValBatch<ValRowLayout<((K, Row), T, R)>>>>;
+    pub type ValRowSpine<K, T, R> = Spine<Arc<OrdValBatch<ValRowLayout<((K, Row), T, R)>>>>;
     pub type ValRowBatcher<K, T, R> = KeyValBatcher<K, Row, T, R>;
     pub type ValRowBuilder<K, T, R> =
-        RcBuilder<crate::dictionary::builders::ValRowBuilder<K, T, R>>;
+        ArcBuilder<crate::dictionary::builders::ValRowBuilder<K, T, R>>;
 
     /// `ValRowBuilder` variant that consumes [`Column`] chunks. Pairs with
     /// `Col2ValPagedBatcher<K, Row, T, R>` for the spillable arrange path where
@@ -86,7 +86,7 @@ mod spines {
     ///
     /// [`Column`]: mz_timely_util::columnar::Column
     pub type ValRowColPagedBuilder<K, T, R> =
-        RcBuilder<crate::dictionary::builders::ValRowColPagedBuilder<K, T, R>>;
+        ArcBuilder<crate::dictionary::builders::ValRowColPagedBuilder<K, T, R>>;
 
     /// A layout based on timely stacks
     pub struct RowRowLayout<U: Update<Key = Row, Val = Row>> {
@@ -155,10 +155,28 @@ mod spines {
 #[cfg(test)]
 mod tests {
     use crate::DatumContainer;
+    use crate::spines::{RowLayout, RowRowLayout, RowValLayout};
     use differential_dataflow::trace::implementations::BatchContainer;
+    use differential_dataflow::trace::implementations::ord_neu::{OrdKeyBatch, OrdValBatch};
     use mz_repr::adt::date::Date;
     use mz_repr::adt::interval::Interval;
-    use mz_repr::{Datum, Row, SqlScalarType};
+    use mz_repr::{Datum, Diff, Row, SqlScalarType, Timestamp};
+    use mz_timely_util::columnation::ColumnationStack;
+
+    fn assert_send_sync<T: Send + Sync>() {}
+
+    /// The batch types backing our spines must stay `Send + Sync`, so that batches
+    /// can be shared across threads (for example behind an `Arc`) to serve reads
+    /// from outside the worker that maintains the trace. This holds because the
+    /// backing containers bottom out in `Vec`s, lgalloc regions, and `CompactBytes`,
+    /// all of which are thread-safe.
+    #[mz_ore::test]
+    fn batches_are_send_sync() {
+        assert_send_sync::<OrdValBatch<RowRowLayout<((Row, Row), Timestamp, Diff)>>>();
+        assert_send_sync::<OrdValBatch<RowValLayout<((Row, Row), Timestamp, Diff)>>>();
+        assert_send_sync::<OrdKeyBatch<RowLayout<((Row, ()), Timestamp, Diff)>>>();
+        assert_send_sync::<ColumnationStack<((Row, Row), Timestamp, Diff)>>();
+    }
 
     #[mz_ore::test]
     #[cfg_attr(miri, ignore)] // unsupported operation: integer-to-pointer casts and `ptr::with_exposed_provenance` are not supported

src/storage/src/render/sinks.rs

@@ -15,7 +15,7 @@ use std::time::{Duration, Instant};
 use differential_dataflow::operators::arrange::{Arrange, Arranged, TraceAgent};
 use differential_dataflow::trace::TraceReader;
 use differential_dataflow::trace::implementations::ord_neu::{
-    OrdValBatcher, OrdValSpine, RcOrdValBuilder,
+    ArcOrdValBuilder, ArcOrdValSpine, OrdValBatcher,
 };
 use differential_dataflow::{AsCollection, Hashable, VecCollection};
 use mz_persist_client::operators::shard_source::SnapshotMode;
@@ -35,7 +35,7 @@ use crate::storage_state::StorageState;
 /// The concrete trace type produced internally when arranging a sink's input.
 /// The sink never sees this directly — only the batches flowing through it —
 /// but it's the anchor for the batch type in [`SinkBatchStream`].
-pub(crate) type SinkTrace = TraceAgent<OrdValSpine<Option<Row>, Row, Timestamp, Diff>>;
+pub(crate) type SinkTrace = TraceAgent<ArcOrdValSpine<Option<Row>, Row, Timestamp, Diff>>;
 
 /// Stream of arrangement batches handed to [`SinkRender::render_sink`].
 ///
@@ -148,7 +148,7 @@ fn arrange_sink_input<'scope>(
     // Allow access to `arrange_named` because we cannot access Mz's wrapper
     // from here. TODO(database-issues#5046): Revisit with cluster unification.
     #[allow(clippy::disallowed_methods)]
-    let Arranged {stream, trace: _} = keyed.arrange_named::<OrdValBatcher<_, _, _, _>, RcOrdValBuilder<_, _, _, _>, OrdValSpine<_, _, _, _>>("Arrange Sink");
+    let Arranged {stream, trace: _} = keyed.arrange_named::<OrdValBatcher<_, _, _, _>, ArcOrdValBuilder<_, _, _, _>, ArcOrdValSpine<_, _, _, _>>("Arrange Sink");
     stream
 }