compute: barrier discipline for MV sink shared batches

misc/python/materialize/mzcompose/__init__.py

@@ -305,6 +305,11 @@ def get_variable_system_parameters(
             "true",
             ["true", "false"],
         ),
+        VariableSystemParameter(
+            "enable_compute_sync_mv_sink_shared_batches",
+            "true",
+            ["true", "false"],
+        ),
         VariableSystemParameter(
             "enable_password_auth",
             "true",
@@ -617,6 +622,10 @@ def get_default_system_parameters(
 # all. Only add it in UNINTERESTING_SYSTEM_PARAMETERS if none of the above
 # apply.
 UNINTERESTING_SYSTEM_PARAMETERS = [
+    # The shared-batch barrier adds an unbounded cross-worker wait, too risky to randomize across
+    # the whole CI suite. It is covered by a dedicated testdrive test and the parallel-workload flag
+    # flipper. (The best-effort shared_batches flag stays in get_variable_system_parameters.)
+    "enable_compute_sync_mv_sink_shared_batches_barrier",
     "enable_compute_half_join2",
     "enable_mz_join_core",
     "linear_join_yielding",

misc/python/materialize/parallel_workload/action.py

@@ -2882,6 +2882,12 @@ def __init__(
             BOOLEAN_FLAG_VALUES
         )
         self.flags_with_values["enable_eager_delta_joins"] = BOOLEAN_FLAG_VALUES
+        self.flags_with_values["enable_compute_sync_mv_sink_shared_batches"] = (
+            BOOLEAN_FLAG_VALUES
+        )
+        self.flags_with_values["enable_compute_sync_mv_sink_shared_batches_barrier"] = (
+            BOOLEAN_FLAG_VALUES
+        )
         self.flags_with_values["enable_public_metrics_endpoint"] = BOOLEAN_FLAG_VALUES
         self.flags_with_values["persist_batch_structured_key_lower_len"] = [
             "0",

src/compute-types/src/dyncfgs.rs

@@ -262,6 +262,27 @@ pub const ENABLE_SYNC_MV_SINK: Config<bool> = Config::new(
     ParameterScope::Environment,
 );
 
+/// Coalesce the parts written by the workers in one process into a single shared batch per batch
+/// interval, instead of each worker writing its own batch.
+pub const ENABLE_SYNC_MV_SINK_SHARED_BATCHES: Config<bool> = Config::new(
+    "enable_compute_sync_mv_sink_shared_batches",
+    false,
+    "Coalesce the parts written by the workers in one process into a single shared batch per \
+     batch interval in the MV sink.",
+    ParameterScope::Environment,
+);
+
+/// When shared batches are enabled, make the workers in one process barrier on each batch interval
+/// so they coalesce into exactly one batch, rather than best-effort. The barrier waits for every
+/// process-local worker to report, so there is no timeout.
+pub const ENABLE_SYNC_MV_SINK_SHARED_BATCHES_BARRIER: Config<bool> = Config::new(
+    "enable_compute_sync_mv_sink_shared_batches_barrier",
+    false,
+    "Barrier the workers in one process onto a single shared batch per interval in the MV sink, \
+     rather than coalescing best-effort.",
+    ParameterScope::Environment,
+);
+
 /// Whether rendering should use the new MV sink correction buffer implementation.
 pub const ENABLE_CORRECTION_V2: Config<bool> = Config::new(
     "enable_compute_correction_v2",
@@ -688,6 +709,8 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&ENABLE_ERROR_DISTINCT)
         .add(&ENABLE_MZ_JOIN_CORE)
         .add(&ENABLE_SYNC_MV_SINK)
+        .add(&ENABLE_SYNC_MV_SINK_SHARED_BATCHES)
+        .add(&ENABLE_SYNC_MV_SINK_SHARED_BATCHES_BARRIER)
         .add(&ENABLE_CORRECTION_V2)
         .add(&CORRECTION_V2_CHAIN_PROPORTIONALITY)
         .add(&CORRECTION_V2_CHUNK_SIZE)

src/compute/src/compute_state.rs

@@ -52,6 +52,7 @@ use mz_persist_types::PersistLocation;
 use mz_persist_types::codec_impls::UnitSchema;
 use mz_repr::fixed_length::ExtendDatums;
 use mz_repr::{DatumVec, Diff, GlobalId, Row, RowArena, Timestamp};
+use mz_storage_operators::persist::SharedBatches;
 use mz_storage_operators::stats::StatsCursor;
 use mz_storage_types::StorageDiff;
 use mz_storage_types::controller::CollectionMetadata;
@@ -114,6 +115,9 @@ pub struct ComputeState {
     /// A process-global cache of (blob_uri, consensus_uri) -> PersistClient.
     /// This is intentionally shared between workers.
     pub persist_clients: Arc<PersistClientCache>,
+    // A process-global cache of shared persist batch writers. This allows
+    // coalescing writes across local workers without explicit exchanges.
+    pub persist_batches: SharedBatches,
     /// Context necessary for rendering txn-wal operators.
     pub txns_ctx: TxnsContext,
     /// History of commands received by this workers and all its peers.
@@ -178,6 +182,7 @@ impl ComputeState {
     /// Construct a new `ComputeState`.
     pub fn new(
         persist_clients: Arc<PersistClientCache>,
+        persist_batches: SharedBatches,
         txns_ctx: TxnsContext,
         metrics: WorkerMetrics,
         tracing_handle: Arc<TracingHandle>,
@@ -198,6 +203,7 @@ impl ComputeState {
             peek_stash_persist_location: None,
             compute_logger: None,
             persist_clients,
+            persist_batches,
             txns_ctx,
             command_history,
             max_result_size: u64::MAX,

src/compute/src/compute_state/peek_stash.rs

@@ -155,14 +155,13 @@ impl StashingPeek {
         //
         // TODO: We _could_ work around the above by teaching the bare columnar
         // Row encoder about zero-column rows.
-        let mut batch_builder = client
-            .batch_builder::<SourceData, (), Timestamp, i64>(
-                shard_id,
-                write_schemas,
-                lower,
-                Some(batch_max_runs),
-            )
-            .await;
+        let mut batch_builder = client.batch_builder::<SourceData, (), Timestamp, i64>(
+            shard_id,
+            "peek_stash",
+            write_schemas,
+            lower,
+            Some(batch_max_runs),
+        );
 
         let mut num_rows: u64 = 0;
 

src/compute/src/server.rs

@@ -29,6 +29,7 @@ use mz_ore::halt;
 use mz_ore::metrics::MetricsRegistry;
 use mz_ore::tracing::TracingHandle;
 use mz_persist_client::cache::PersistClientCache;
+use mz_storage_operators::persist::SharedBatches;
 use mz_storage_types::connections::ConnectionContext;
 use mz_timely_util::capture::EventLink;
 use mz_txn_wal::operator::TxnsContext;
@@ -129,6 +130,8 @@ pub(crate) type StorageTimelyLogReader =
 struct Config {
     /// `persist` client cache.
     pub persist_clients: Arc<PersistClientCache>,
+    /// Shared persist batch state.
+    pub persist_batches: SharedBatches,
     /// Context necessary for rendering txn-wal operators.
     pub txns_ctx: TxnsContext,
     /// A process-global handle to tracing configuration.
@@ -173,6 +176,7 @@ pub async fn serve(
     mz_timely_util::pool_config::metrics::register(metrics_registry);
 
     let config = Config {
+        persist_batches: SharedBatches::new(),
         persist_clients,
         txns_ctx,
         tracing_handle,
@@ -305,6 +309,7 @@ struct Worker<'w> {
     /// A process-global cache of (blob_uri, consensus_uri) -> PersistClient.
     /// This is intentionally shared between workers
     persist_clients: Arc<PersistClientCache>,
+    persist_batches: SharedBatches,
     /// Context necessary for rendering txn-wal operators.
     txns_ctx: TxnsContext,
     /// A process-global handle to tracing configuration.
@@ -361,6 +366,7 @@ impl ClusterSpec for Config {
             metrics,
             context: self.context.clone(),
             persist_clients: Arc::clone(&self.persist_clients),
+            persist_batches: self.persist_batches.clone(),
             txns_ctx: self.txns_ctx.clone(),
             compute_state: None,
             tracing_handle: Arc::clone(&self.tracing_handle),
@@ -498,6 +504,7 @@ impl<'w> Worker<'w> {
         if matches!(&cmd, ComputeCommand::CreateInstance(_)) {
             self.compute_state = Some(ComputeState::new(
                 Arc::clone(&self.persist_clients),
+                self.persist_batches.clone(),
                 self.txns_ctx.clone(),
                 self.metrics.clone(),
                 Arc::clone(&self.tracing_handle),

src/compute/src/sink/materialized_view.rs

@@ -134,6 +134,7 @@ use mz_persist_client::write::WriteHandle;
 use mz_persist_client::{Diagnostics, PersistClient};
 use mz_persist_types::codec_impls::UnitSchema;
 use mz_repr::{Diff, GlobalId, Row, Timestamp};
+use mz_storage_operators::persist::{SharedBatchId, SharedBatches};
 use mz_storage_types::StorageDiff;
 use mz_storage_types::controller::CollectionMetadata;
 use mz_storage_types::sources::SourceData;
@@ -267,6 +268,7 @@ where
 
     let persist_api = PersistApi {
         persist_clients: Arc::clone(&compute_state.persist_clients),
+        persist_batches: compute_state.persist_batches.clone(),
         collection: target.clone(),
         shard_name: sink_id.to_string(),
         purpose: format!("MV sink {sink_id}"),
@@ -351,6 +353,7 @@ pub(super) fn advance(
 #[derive(Clone)]
 pub(super) struct PersistApi {
     pub(super) persist_clients: Arc<PersistClientCache>,
+    pub(super) persist_batches: SharedBatches,
     pub(super) collection: CollectionMetadata,
     pub(super) shard_name: String,
     pub(super) purpose: String,
@@ -442,6 +445,8 @@ pub(super) struct BatchDescription {
     pub(super) lower: Antichain<Timestamp>,
     pub(super) upper: Antichain<Timestamp>,
     pub(super) append_worker: usize,
+    /// Identifies the shared batch all workers building this description contribute to.
+    pub(super) shared_id: SharedBatchId,
 }
 
 impl BatchDescription {
@@ -455,6 +460,7 @@ impl BatchDescription {
             lower,
             upper,
             append_worker,
+            shared_id: SharedBatchId::new(),
         }
     }
 }