mz_compute/sink/

materialized_view_v2.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Sync Timely operator implementation of the MV sink.
//!
//! This module provides an alternative implementation of `persist_sink` that uses sync Timely
//! operators communicating with Tokio tasks via channels, instead of async Timely operators.
//! Gated behind the `ENABLE_SYNC_MV_SINK` dyncfg.
//!
//! See the [main module](super::materialized_view) for the operator graph and design docs.
//!
//! ### Channel ordering requirements
//!
//! Each operator splits state across a Timely thread (which observes inputs and frontiers) and a
//! Tokio task (which owns persist I/O state). They communicate via `mpsc` command channels, which
//! preserve send order on a single sender. Each operator instance constructs its own
//! `(tx, rx)` pair inside `render` and never clones the sender, so there is exactly one producer
//! per channel — sends are totally ordered. Different worker instances of the same operator
//! never share a channel, so cross-worker ordering is not a concern. The correctness of the
//! operators relies on a few ordering invariants between the messages sent within a single Timely
//! activation:
//!
//!  * **`mint`**: the `persist_watch` Tokio task is the sole producer of persist-frontier updates
//!    and emits them in monotonically increasing order, terminated by the empty frontier. The
//!    Timely closure drains the receiver each activation; processing in receive order is therefore
//!    sufficient. No cross-channel ordering is needed because `mint` only has the one channel.
//!
//!  * **`write`**: per-activation, the Timely closure first appends all observed input data into
//!    a single `WriteCommand::Batch` and only then sends a `WriteCommand::WriteBatch` (issued from
//!    `maybe_start_batch` after frontier checks). The Tokio task processes commands FIFO, so a
//!    `WriteBatch` is guaranteed to see every `Batch` from the same activation already applied to
//!    the corrections buffer. Reversing this order would let the task write a batch that is
//!    missing updates the Timely closure already observed.
//!
//!  * **`append`**: per-activation, the Timely closure forwards messages in the order
//!    `Description` → `Batch` → `BatchesFrontier`. The first two carry the data the task needs
//!    to absorb; `BatchesFrontier` is the trigger that allows `maybe_append_batches` to fire.
//!    Sending `BatchesFrontier` *after* its corresponding `Batch` messages ensures the task does
//!    not append a batch description before all batches contributing to it have been absorbed.
//!    If the order were reversed, `maybe_append_batches` could fire on an incomplete `batches`
//!    set and miss writes.

use std::any::Any;
use std::cell::RefCell;
use std::rc::Rc;
use std::sync::Arc;

use differential_dataflow::{Hashable, VecCollection};
use mz_compute_types::dyncfgs::{
    MV_SINK_ADVANCE_PERSIST_FRONTIERS, MV_SINK_SHARED_BATCH_MAX_BYTES,
};
use mz_dyncfg::{ConfigSet, ConfigValHandle};
use mz_ore::cast::CastFrom;
use mz_persist_client::batch::{Batch, ProtoBatch};
use mz_persist_client::write::WriteHandle;
use mz_repr::{Diff, GlobalId, Row, Timestamp};
use mz_storage_types::StorageDiff;
use mz_storage_types::sources::SourceData;
use timely::PartialOrder;
use timely::dataflow::channels::pact::{Exchange, Pipeline};
use timely::dataflow::operators::generic::OutputBuilder;
use timely::dataflow::operators::generic::builder_rc::OperatorBuilder as OperatorBuilderRc;
use timely::dataflow::operators::vec::Broadcast;
use timely::dataflow::operators::{Capability, CapabilitySet};
use timely::progress::Antichain;
use timely::progress::frontier::AntichainRef;
use tokio::sync::{mpsc, watch};
use tracing::trace;

use crate::compute_state::ComputeState;
use crate::render::StartSignal;
use crate::render::errors::DataflowErrorSer;
use crate::sink::correction::{ChannelLogging, Correction, CorrectionLogger};
use crate::sink::materialized_view::{
    BatchDescription, BatchesStream, DescsStream, DesiredStreams, OkErr, PersistApi,
    PersistStreams, SharedSinkFrontier, advance, operator_name, persist_source,
};
use crate::sink::shared_batch::SharedBatchWriter;

/// Renders an MV sink writing the given desired collection into the `target` persist collection.
///
/// This is the sync Timely operator implementation, using Tokio tasks for I/O.
pub(super) fn persist_sink<'s>(
    sink_id: GlobalId,
    target: &mz_storage_types::controller::CollectionMetadata,
    ok_collection: VecCollection<'s, Timestamp, Row, Diff>,
    err_collection: VecCollection<'s, Timestamp, DataflowErrorSer, Diff>,
    as_of: Antichain<Timestamp>,
    compute_state: &mut ComputeState,
    start_signal: StartSignal,
    read_only_rx: watch::Receiver<bool>,
) -> Rc<dyn Any> {
    let scope = ok_collection.scope();
    let desired = OkErr::new(ok_collection.inner, err_collection.inner);

    // Read back the persist shard.
    let (persist, persist_token) =
        persist_source(scope, sink_id, target.clone(), compute_state, start_signal);

    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}"),
    };

    let (desired, descs, sink_frontier) = mint::render(
        sink_id,
        persist_api.clone(),
        as_of.clone(),
        read_only_rx.clone(),
        desired,
    );

    // Broadcast batch descriptions to all workers, regardless of whether or not they are
    // responsible for the append, to give them a chance to clean up any outdated state they
    // might still hold.
    let descs = descs.broadcast();

    let batches = write::render(
        sink_id,
        persist_api.clone(),
        as_of,
        desired,
        persist,
        descs.clone(),
        read_only_rx,
        Rc::clone(&compute_state.worker_config),
    );

    append::render(sink_id, persist_api, descs, batches);

    // Report sink frontier updates to the `ComputeState`.
    let collection = compute_state.expect_collection_mut(sink_id);
    collection.sink_write_frontier = Some(sink_frontier);

    Rc::new(persist_token)
}

/// Implementation of the `mint` operator.
mod mint {
    use super::*;
    use timely::progress::frontier::AntichainRef;

    /// Render the `mint` operator.
    ///
    /// The parameters passed in are:
    ///  * `sink_id`: The `GlobalId` of the sink export.
    ///  * `persist_api`: An object providing access to the output persist shard.
    ///  * `as_of`: The first time for which the sink may produce output.
    ///  * `read_only_rx`: A receiver that reports the sink is in read-only mode.
    ///  * `desired`: The ok/err streams that should be sinked to persist.
    pub fn render<'s>(
        sink_id: GlobalId,
        persist_api: PersistApi,
        as_of: Antichain<Timestamp>,
        mut read_only_rx: watch::Receiver<bool>,
        desired: DesiredStreams<'s>,
    ) -> (DesiredStreams<'s>, DescsStream<'s>, SharedSinkFrontier) {
        let scope = desired.ok.scope();
        let worker_id = scope.index();
        let worker_count = scope.peers();

        // Determine the active worker for the mint operator.
        let active_worker_id = usize::cast_from(sink_id.hashed()) % scope.peers();

        let sink_frontier = Rc::new(RefCell::new(Antichain::from_elem(Timestamp::MIN)));
        let shared_frontier = Rc::clone(&sink_frontier);

        let name = operator_name(sink_id, "mint");
        let mut builder = OperatorBuilderRc::new(name, scope.clone());
        let info = builder.operator_info();

        // Create outputs (before inputs, so no input connections yet).
        let (ok_output, ok_stream) = builder.new_output();
        let (err_output, err_stream) = builder.new_output();
        let (desc_output, desc_stream) = builder.new_output();

        let mut ok_output = OutputBuilder::from(ok_output);
        let mut err_output = OutputBuilder::from(err_output);
        let mut desc_output = OutputBuilder::from(desc_output);

        // desired_ok -> output 0 (ok passthrough)
        let mut desired_ok_input = builder.new_input_connection(
            desired.ok,
            Pipeline,
            [(0, Antichain::from_elem(Default::default()))],
        );
        // desired_err -> output 1 (err passthrough)
        let mut desired_err_input = builder.new_input_connection(
            desired.err,
            Pipeline,
            [(1, Antichain::from_elem(Default::default()))],
        );

        // Set up background tasks and state for the active worker only.
        let mut task_handles = Vec::new();
        let read_only = *read_only_rx.borrow_and_update();
        let mut state = None;
        if worker_id == active_worker_id {
            // Spawn a Tokio task to watch the persist shard's upper frontier.
            //
            // We collect the persist frontier from a write handle directly, rather than
            // inspecting the `persist` stream, because the latter has two annoying glitches:
            //  (a) It starts at the shard's read frontier, not its write frontier.
            //  (b) It can lag behind if there are spikes in ingested data.
            //
            // The decoupling from the `persist` stream is load-bearing: that stream can fall
            // arbitrarily behind the shard upper during snapshot replay or write spikes. Using
            // it would delay both (1) the controller-visible sink frontier (`shared_frontier`),
            // which previously caused a CrossJoin feature-bench regression where the controller
            // held a finished MV dataflow open waiting for the empty frontier, and (2) the
            // `state.persist_frontier` that gates batch-description minting, stalling mint until
            // the read-back stream catches up. The goal isn't tick-granular descriptions per
            // se — it's avoiding the stream-induced stall. See 5eab5ff896 for the original
            // regression and rationale.
            //
            // The task sends the empty frontier as its final message before exiting. The
            // operator drops `persist_rx` once it receives the empty frontier.
            let (persist_tx, persist_rx) = mpsc::unbounded_channel();
            let sync_activator = scope.worker().sync_activator_for(info.address.to_vec());
            let handle = mz_ore::task::spawn(
                || operator_name(sink_id, "mint::persist_watch"),
                async move {
                    let mut writer = persist_api.open_writer().await;
                    let mut frontier = Antichain::from_elem(Timestamp::MIN);
                    loop {
                        writer.wait_for_upper_past(&frontier).await;
                        frontier = writer.upper().clone();
                        if persist_tx.send(frontier.clone()).is_err() {
                            return;
                        }
                        if sync_activator.activate().is_err() {
                            return;
                        }
                        if frontier.is_empty() {
                            return;
                        }
                    }
                },
            );
            task_handles.push(handle.abort_on_drop());

            // Spawn a Tokio task to wake the operator when read-only mode changes.
            if read_only {
                let sync_activator = scope.worker().sync_activator_for(info.address.to_vec());
                let mut rx = read_only_rx.clone();
                let handle = mz_ore::task::spawn(
                    || format!("mv_sink({sink_id})::mint::read_only_watch"),
                    async move {
                        let _ = rx.changed().await;
                        let _ = sync_activator.activate();
                    },
                );
                task_handles.push(handle.abort_on_drop());
            }

            state = Some(State::new(
                sink_id,
                worker_count,
                as_of,
                read_only,
                persist_rx,
            ));
        }

        builder.build(move |capabilities| {
            // Passing through the `desired` streams only requires data capabilities, so we can
            // immediately drop their initial capabilities here.
            let [_, _, desc_cap]: [_; 3] =
                capabilities.try_into().expect("one capability per output");

            let mut cap_set = if state.is_some() {
                Some(CapabilitySet::from_elem(desc_cap))
            } else {
                drop(desc_cap);
                shared_frontier.borrow_mut().clear();
                None
            };

            move |frontiers| {
                // Keep task handles alive so they are aborted when the operator is dropped.
                let _ = &task_handles;

                // Pass through desired data.
                let mut ok_out = ok_output.activate();
                desired_ok_input.for_each(|cap, data| {
                    ok_out.session(&cap).give_container(data);
                });
                let mut err_out = err_output.activate();
                desired_err_input.for_each(|cap, data| {
                    err_out.session(&cap).give_container(data);
                });

                let Some(state) = &mut state else {
                    // Non-active worker: just pass through data.
                    return;
                };
                let cap_set = cap_set.as_mut().unwrap();

                // Track desired frontiers.
                state.advance_desired_ok_frontier(frontiers[0].frontier());
                state.advance_desired_err_frontier(frontiers[1].frontier());

                state.drain_persist_rx(&shared_frontier);

                // Check read-only mode.
                if state.read_only && read_only_rx.has_changed().unwrap_or(false) {
                    if !*read_only_rx.borrow_and_update() {
                        state.allow_writes();
                    }
                }

                // Try to mint a batch description.
                let mut desc_out = desc_output.activate();
                if let Some(desc) = state.maybe_mint_batch_description() {
                    let lower_ts = *desc.lower.as_option().expect("not empty");
                    let cap = cap_set.delayed(&lower_ts);
                    desc_out.session(&cap).give(desc);

                    // We only emit strictly increasing `lower`s, so we can let our output frontier
                    // advance beyond the current `lower`.
                    cap_set.downgrade([lower_ts.step_forward()]);
                } else {
                    // The next emitted `lower` will be at least the `persist` frontier, so we can
                    // advance our output frontier as far.
                    let _ = cap_set.try_downgrade(state.persist_frontier.iter());
                }
            }
        });

        let desired_output_streams = OkErr::new(ok_stream, err_stream);
        (desired_output_streams, desc_stream, sink_frontier)
    }

    /// State maintained by the `mint` operator.
    struct State {
        sink_id: GlobalId,
        /// The number of workers in the Timely cluster.
        worker_count: usize,
        /// The frontiers of the `desired` inputs.
        desired_frontiers: OkErr<Antichain<Timestamp>, Antichain<Timestamp>>,
        /// The frontier of the target persist shard.
        persist_frontier: Antichain<Timestamp>,
        /// Receiver for persist frontier updates from the Tokio persist_watch task.
        ///
        /// Dropped once the empty frontier is received (the task's shutdown signal).
        persist_rx: Option<mpsc::UnboundedReceiver<Antichain<Timestamp>>>,
        /// The append worker for the next batch description, chosen in round-robin fashion.
        next_append_worker: usize,
        /// The last `lower` we have emitted in a batch description, if any. Whenever the
        /// `persist_frontier` moves beyond this frontier, we need to mint a new description.
        last_lower: Option<Antichain<Timestamp>>,
        /// Whether we are operating in read-only mode.
        ///
        /// In read-only mode, minting of batch descriptions is disabled.
        read_only: bool,
    }

    impl State {
        fn new(
            sink_id: GlobalId,
            worker_count: usize,
            as_of: Antichain<Timestamp>,
            read_only: bool,
            persist_rx: mpsc::UnboundedReceiver<Antichain<Timestamp>>,
        ) -> Self {
            // Initializing `persist_frontier` to the `as_of` ensures that the first minted batch
            // description will have a `lower` of `as_of` or beyond, and thus that we don't spend
            // effort writing out snapshots of data that is already in the shard.
            let persist_frontier = as_of;

            Self {
                sink_id,
                worker_count,
                desired_frontiers: OkErr::new_frontiers(),
                persist_frontier,
                persist_rx: Some(persist_rx),
                next_append_worker: 0,
                last_lower: None,
                read_only,
            }
        }

        fn trace<S: AsRef<str>>(&self, message: S) {
            let message = message.as_ref();
            trace!(
                sink_id = %self.sink_id,
                desired_frontier = ?self.desired_frontiers.frontier().elements(),
                persist_frontier = ?self.persist_frontier.elements(),
                last_lower = ?self.last_lower,
                message,
            );
        }

        fn advance_desired_ok_frontier(&mut self, frontier: AntichainRef<Timestamp>) {
            if advance(&mut self.desired_frontiers.ok, frontier) {
                self.trace("advanced `desired` ok frontier");
            }
        }

        fn advance_desired_err_frontier(&mut self, frontier: AntichainRef<Timestamp>) {
            if advance(&mut self.desired_frontiers.err, frontier) {
                self.trace("advanced `desired` err frontier");
            }
        }

        fn advance_persist_frontier(&mut self, frontier: AntichainRef<Timestamp>) {
            if advance(&mut self.persist_frontier, frontier) {
                self.trace("advanced `persist` frontier");
            }
        }

        /// Drain persist frontier updates from the Tokio task.
        ///
        /// Frontiers from the `persist_watch` task are monotonically increasing, so only the
        /// most recent one matters. We drain all queued messages and apply just the latest,
        /// avoiding redundant `advance_persist_frontier`/`trace!` calls when several updates
        /// arrived between activations.
        ///
        /// The task sends the empty frontier as its final message before exiting. Once
        /// received, we drop the receiver.
        fn drain_persist_rx(&mut self, shared_frontier: &RefCell<Antichain<Timestamp>>) {
            let Some(mut rx) = self.persist_rx.take() else {
                return;
            };
            let mut latest: Option<Antichain<Timestamp>> = None;
            let mut closed = false;
            loop {
                match rx.try_recv() {
                    Ok(frontier) => {
                        let done = frontier.is_empty();
                        latest = Some(frontier);
                        if done {
                            closed = true;
                            break;
                        }
                    }
                    Err(mpsc::error::TryRecvError::Empty) => break,
                    Err(mpsc::error::TryRecvError::Disconnected) => {
                        panic!("mint persist_watch task unexpectedly gone");
                    }
                }
            }
            if let Some(frontier) = latest {
                shared_frontier.borrow_mut().clone_from(&frontier);
                self.advance_persist_frontier(frontier.borrow());
            }
            if !closed {
                self.persist_rx = Some(rx);
            }
        }

        fn allow_writes(&mut self) {
            if self.read_only {
                self.read_only = false;
                self.trace("switched to write mode");
            }
        }

        fn maybe_mint_batch_description(&mut self) -> Option<BatchDescription> {
            let desired_frontier = self.desired_frontiers.frontier();
            let persist_frontier = &self.persist_frontier;

            // We only mint new batch descriptions when:
            //  1. We are _not_ in read-only mode.
            //  2. The `desired` frontier is ahead of the `persist` frontier.
            //  3. The `persist` frontier advanced since we last emitted a batch description.
            let desired_ahead = PartialOrder::less_than(persist_frontier, desired_frontier);
            let persist_advanced = self.last_lower.as_ref().map_or(true, |lower| {
                PartialOrder::less_than(lower, persist_frontier)
            });

            if self.read_only || !desired_ahead || !persist_advanced {
                return None;
            }

            let lower = persist_frontier.clone();
            let upper = desired_frontier.clone();
            let append_worker = self.next_append_worker;
            let desc = BatchDescription::new(lower, upper, append_worker);

            self.next_append_worker = (append_worker + 1) % self.worker_count;
            self.last_lower = Some(desc.lower.clone());

            self.trace(format!("minted batch description: {desc:?}"));
            Some(desc)
        }
    }
}

/// Implementation of the `write` operator.
mod write {
    use super::*;

    use mz_timely_util::activator::ArcActivator;

    /// Commands sent from the Timely operator to the Tokio write task.
    enum WriteCommand {
        /// A coalesced batch of work gathered during a single operator activation.
        ///
        /// The Timely closure accumulates all updates observed across the four data inputs plus
        /// any frontier advancement and forced-consolidation flag, and sends a single
        /// `WriteCommand::Batch` per activation. This keeps the channel overhead independent of
        /// the number of Timely chunks processed per activation.
        Batch(BatchUpdates),
        /// Write a batch with the given description. The task drains corrections and writes
        /// them to persist.
        WriteBatch(BatchDescription),
    }

    /// The payload of a coalesced [`WriteCommand::Batch`].
    struct BatchUpdates {
        /// Positive contributions from the `desired` ok input.
        desired_ok: Vec<(Row, Timestamp, Diff)>,
        /// Positive contributions from the `desired` err input.
        desired_err: Vec<(DataflowErrorSer, Timestamp, Diff)>,
        /// Negative contributions from the `persist` ok input.
        persist_ok: Vec<(Row, Timestamp, Diff)>,
        /// Negative contributions from the `persist` err input.
        persist_err: Vec<(DataflowErrorSer, Timestamp, Diff)>,
        /// The new persist frontier, if it advanced this activation.
        persist_frontier: Option<Antichain<Timestamp>>,
        /// Whether a consolidation of the corrections buffer should be forced.
        force_consolidation: bool,
    }

    impl BatchUpdates {
        fn new() -> Self {
            Self {
                desired_ok: Vec::new(),
                desired_err: Vec::new(),
                persist_ok: Vec::new(),
                persist_err: Vec::new(),
                persist_frontier: None,
                force_consolidation: false,
            }
        }

        /// Returns true if there is no work in this batch.
        fn is_empty(&self) -> bool {
            self.desired_ok.is_empty()
                && self.desired_err.is_empty()
                && self.persist_ok.is_empty()
                && self.persist_err.is_empty()
                && self.persist_frontier.is_none()
                && !self.force_consolidation
        }
    }

    /// The correction buffers owned by the Tokio write task.
    type Corrections = OkErr<Correction<Row>, Correction<DataflowErrorSer>>;

    /// A response from the Tokio write task back to the Timely operator.
    struct WriteResponse {
        /// The written batches, empty if this worker has none to deliver for the interval.
        batches: Vec<ProtoBatch>,
    }

    /// Render the `write` operator.
    ///
    /// The parameters passed in are:
    ///  * `sink_id`: The `GlobalId` of the sink export.
    ///  * `persist_api`: An object providing access to the output persist shard.
    ///  * `as_of`: The first time for which the sink may produce output.
    ///  * `desired`: The ok/err streams that should be sinked to persist.
    ///  * `persist`: The ok/err streams read back from the output persist shard.
    ///  * `descs`: The stream of batch descriptions produced by the `mint` operator.
    pub fn render<'s>(
        sink_id: GlobalId,
        persist_api: PersistApi,
        as_of: Antichain<Timestamp>,
        desired: DesiredStreams<'s>,
        persist: PersistStreams<'s>,
        descs: DescsStream<'s>,
        mut read_only_rx: watch::Receiver<bool>,
        worker_config: Rc<ConfigSet>,
    ) -> BatchesStream<'s> {
        let scope = desired.ok.scope();
        let worker_id = scope.index();

        let name = operator_name(sink_id, "write");
        let mut builder = OperatorBuilderRc::new(name, scope.clone());
        let info = builder.operator_info();

        // Set up correction buffer logging. CorrectionLogger is not Send (uses timely
        // loggers), so the Tokio task uses ChannelLogging to send events back to the
        // Timely thread for application by the CorrectionLogger.
        let mut channel_logging = None;
        let mut correction_logger = None;
        if let (Some(compute_logger), Some(differential_logger)) = (
            scope.worker().logger_for("materialize/compute"),
            scope.worker().logger_for("differential/arrange"),
        ) {
            let operator_info = builder.operator_info();
            let (tx, rx) = mpsc::unbounded_channel();
            channel_logging = Some(ChannelLogging::new(tx));
            correction_logger = Some(CorrectionLogger::new(
                compute_logger,
                differential_logger.into(),
                operator_info.global_id,
                operator_info.address.to_vec(),
                rx,
            ));
        }

        // It is important that we exchange the `desired` and `persist` data the same way, so
        // updates that cancel each other out end up on the same worker.
        let exchange_ok = |(d, _, _): &(Row, Timestamp, Diff)| d.hashed();
        let exchange_err = |(d, _, _): &(DataflowErrorSer, Timestamp, Diff)| d.hashed();

        // Data inputs are created before the output, so they are not connected to it.
        let mut desired_ok_input = builder.new_input(desired.ok, Exchange::new(exchange_ok));
        let mut desired_err_input = builder.new_input(desired.err, Exchange::new(exchange_err));
        let mut persist_ok_input = builder.new_input(persist.ok, Exchange::new(exchange_ok));
        let mut persist_err_input = builder.new_input(persist.err, Exchange::new(exchange_err));
        let mut descs_input = builder.new_input(descs, Pipeline);

        // Only descs (input 4) is connected to the batches output.
        let (batches_output, batches_output_stream) =
            builder.new_output_connection([(4, Antichain::from_elem(Default::default()))]);
        let mut batches_output = OutputBuilder::from(batches_output);

        // Obtain SinkMetrics synchronously from the persist client cache, rather than through
        // a WriteHandle, to avoid async I/O on the Timely thread.
        let sink_metrics = persist_api.persist_clients.metrics().sink.clone();

        // Construct corrections on the Timely thread (reads ConfigSet), then move to the
        // Tokio task. The ChannelLogging sends events back to the Timely thread.
        let worker_metrics = sink_metrics.for_worker(worker_id);
        let mut corrections: Corrections = OkErr::new(
            Correction::new(
                sink_metrics.clone(),
                worker_metrics.clone(),
                channel_logging.clone(),
                &worker_config,
            ),
            Correction::new(
                sink_metrics.clone(),
                worker_metrics,
                channel_logging,
                &worker_config,
            ),
        );

        // Read `MV_SINK_ADVANCE_PERSIST_FRONTIERS` exactly once and reuse the captured value for
        // both the Tokio-side `corrections.since` initialization below and the Timely-side
        // `persist_frontiers` initialization in `State::new`. Re-reading the dyncfg per init site
        // would let the value flip between reads and produce the very inconsistency this fix
        // addresses: `persist_frontiers = as_of` (gate open) with `corrections.since = MIN`
        // (snapshot updates not advanced) reproduces the original `UpdateNotBeyondLower` panic.
        let advance_persist_frontiers_at_startup =
            MV_SINK_ADVANCE_PERSIST_FRONTIERS.get(&worker_config);

        let shared_batch_max_bytes = MV_SINK_SHARED_BATCH_MAX_BYTES.handle(&worker_config);

        // Mirror the persist-frontier initialization performed by `State::new` below. With the
        // flag enabled, `State` advances its Timely-side `persist_frontiers` to `as_of`, opening
        // the `maybe_start_batch` write gate (`desc.lower <= persist_frontiers.frontier()`)
        // immediately for the first description minted with `lower = as_of`. The corrections
        // buffer lives on the Tokio task and only learns of frontier advancements through
        // `WriteCommand::Batch { persist_frontier, .. }`, which the Timely closure populates only
        // when an input frontier actually moves. On startup, the input frontiers begin at
        // `Timestamp::MIN`, so no `Batch` carries `persist_frontier` until the persist input
        // catches up — yet a `WriteBatch(desc)` with `desc.lower = as_of` can already be sent.
        // Snapshot-replay updates inserted in the meantime stay at their original timestamps
        // (`Correction::insert` rounds to `max(t, since)` and `since == MIN`), and slip into the
        // batch, tripping persist's `UpdateNotBeyondLower` invariant. Advancing
        // `corrections.since` here keeps the Tokio side in lockstep with the Timely side, the
        // same invariant `materialized_view::write::State::new` upholds via
        // `apply_persist_frontier_advancement`.
        if advance_persist_frontiers_at_startup {
            corrections.ok.advance_since(as_of.clone());
            corrections.err.advance_since(as_of.clone());
        }

        // Channels for commands and responses.
        let (cmd_tx, mut cmd_rx) = mpsc::unbounded_channel::<WriteCommand>();
        let (resp_tx, mut resp_rx) = mpsc::unbounded_channel::<WriteResponse>();

        // Spawn Tokio task that owns the batch-writing context and corrections buffer.
        let (activator, activation_ack) = ArcActivator::new(scope, &info);
        let write_task_handle = {
            mz_ore::task::spawn(
                || operator_name(sink_id, "write::batch_writer"),
                async move {
                    let batch_writer = BatchWriter::new(&persist_api, shared_batch_max_bytes).await;

                    while let Some(cmd) = cmd_rx.recv().await {
                        corrections =
                            apply_command(sink_id, corrections, &batch_writer, cmd, &resp_tx).await;
                        // Activate the operator to drain logging events and process batch responses.
                        // ArcActivator suppresses redundant activations, so this is cheap.
                        activator.activate();
                    }
                },
            )
            .abort_on_drop()
        };

        // In read-only mode, wake the operator when writes are allowed so the forced
        // consolidation stops re-arming and the `WriteBatch` path takes over promptly. Without
        // this the operator only learns of the transition on its next input activation, which an
        // idle sink may not get for a while.
        let read_only_watch_handle = (*read_only_rx.borrow()).then(|| {
            let sync_activator = scope.worker().sync_activator_for(info.address.to_vec());
            let mut rx = read_only_rx.clone();
            mz_ore::task::spawn(
                || operator_name(sink_id, "write::read_only_watch"),
                async move {
                    let _ = rx.changed().await;
                    let _ = sync_activator.activate();
                },
            )
            .abort_on_drop()
        });

        builder.build(move |capabilities| {
            // We will use the data capabilities from the `descs` input to produce output, so no
            // need to hold onto the initial capabilities.
            drop(capabilities);

            let read_only = *read_only_rx.borrow_and_update();
            let mut state = State::new(
                sink_id,
                worker_id,
                as_of,
                advance_persist_frontiers_at_startup,
                read_only,
            );

            // Whether a batch write is currently in flight in the Tokio task.
            let mut batch_in_flight: Option<(BatchDescription, Capability<Timestamp>)> = None;

            // CorrectionLogger lives on the Timely thread and drains events from
            // the channel each activation. On drop, it drains remaining events and
            // retracts all logged state.
            let mut correction_logger = correction_logger;

            move |frontiers| {
                // Keep task handles alive so they are aborted when the operator is dropped.
                let _ = &write_task_handle;
                let _ = &read_only_watch_handle;

                // Acknowledge activation so the Tokio task can activate us again.
                activation_ack.ack();

                // Drain logging events from the Tokio task's ChannelLogging.
                if let Some(logger) = &mut correction_logger {
                    logger.apply_events();
                }
                // Coalesce all work from this activation into a single command. This keeps
                // per-chunk channel overhead low, which matters for hydration when many small
                // Timely chunks arrive in a single activation sweep.
                let mut batch = BatchUpdates::new();
                desired_ok_input.for_each(|_cap, data| {
                    batch.desired_ok.append(data);
                });
                desired_err_input.for_each(|_cap, data| {
                    batch.desired_err.append(data);
                });
                persist_ok_input.for_each(|_cap, data| {
                    batch.persist_ok.append(data);
                });
                persist_err_input.for_each(|_cap, data| {
                    batch.persist_err.append(data);
                });

                // Accept batch descriptions.
                descs_input.for_each(|cap, data| {
                    let cap = cap.retain(0);
                    for desc in data.drain(..) {
                        state.absorb_batch_description(desc, cap.clone());
                    }
                });

                // Track frontiers. Include the new persist frontier in the coalesced batch for
                // `advance_since`.
                state.advance_desired_ok_frontier(frontiers[0].frontier());
                state.advance_desired_err_frontier(frontiers[1].frontier());
                if state.advance_persist_ok_frontier(frontiers[2].frontier())
                    | state.advance_persist_err_frontier(frontiers[3].frontier())
                {
                    batch.persist_frontier = Some(state.persist_frontiers.frontier().to_owned());
                }
                // Track read-only mode so the forced consolidation stops re-arming once writes
                // are allowed and the `WriteBatch` path takes over driving consolidation.
                if state.read_only && read_only_rx.has_changed().unwrap_or(false) {
                    if !*read_only_rx.borrow_and_update() {
                        state.set_read_only(false);
                    }
                }

                if state.should_force_consolidation() {
                    batch.force_consolidation = true;
                }

                if !batch.is_empty() {
                    cmd_tx
                        .send(WriteCommand::Batch(batch))
                        .expect("write task unexpectedly gone");
                }

                // Try to receive batch results from the Tokio task.
                loop {
                    match resp_rx.try_recv() {
                        Ok(resp) => {
                            if let Some((desc, cap)) = batch_in_flight.take() {
                                if resp.batches.is_empty() {
                                    state.trace("skipping empty batch");
                                } else {
                                    state.trace(format!("wrote {} batches", resp.batches.len()));
                                    let mut out = batches_output.activate();
                                    let mut session = out.session(&cap);
                                    for batch in resp.batches {
                                        session.give((desc.clone(), batch));
                                    }
                                }
                            }
                        }
                        Err(mpsc::error::TryRecvError::Empty) => break,
                        Err(mpsc::error::TryRecvError::Disconnected) => {
                            panic!("write task unexpectedly gone");
                        }
                    }
                }

                // If no batch in flight, try to write a new batch.
                if batch_in_flight.is_none() {
                    if let Some((desc, cap)) = state.maybe_start_batch(&cmd_tx) {
                        batch_in_flight = Some((desc, cap));
                    }
                }
            }
        });

        batches_output_stream
    }

    /// How many updates the batch read-back hands over per chunk.
    const READ_BACK_CHUNK: usize = 1024;
    /// How many chunks may sit between the blocking read-back and the async task feeding the
    /// batch builder. Together with [`READ_BACK_CHUNK`] this bounds the handoff to a few thousand
    /// buffered updates.
    const READ_BACK_CHUNKS_IN_FLIGHT: usize = 4;

    /// The batch-writing context owned by the Tokio write task.
    ///
    /// The `writer` cleans up a finished batch when the response channel has already gone away.
    struct BatchWriter {
        writer: WriteHandle<SourceData, (), Timestamp, StorageDiff>,
        shared: SharedBatchWriter,
    }

    impl BatchWriter {
        async fn new(
            persist_api: &PersistApi,
            shared_batch_max_bytes: ConfigValHandle<usize>,
        ) -> Self {
            let writer = persist_api.open_writer().await;
            let shared = SharedBatchWriter::new(persist_api, &writer, shared_batch_max_bytes).await;
            Self { writer, shared }
        }
    }

    /// Apply a single command to the task state, returning the correction buffers.
    ///
    /// `desired` updates enter `corrections` as positive contributions and `persist` updates as
    /// negative contributions, so the buffer contains `desired - persist`, i.e. the updates that
    /// need to be written to bring the shard in line with `desired`.
    ///
    /// Correction maintenance is CPU-bound and unbounded in duration: an insert can merge chains
    /// spanning the whole buffer and a consolidation sorts it, neither with an await point in
    /// between. Running that inline occupies a Tokio worker thread for the entire time, which
    /// stops it from polling every other task scheduled on it. It therefore runs on a blocking
    /// thread, which the OS can preempt. The buffers move into the blocking closure and back out
    /// again because this task owns them exclusively.
    async fn apply_command(
        sink_id: GlobalId,
        mut corrections: Corrections,
        batch_writer: &BatchWriter,
        cmd: WriteCommand,
        resp_tx: &mpsc::UnboundedSender<WriteResponse>,
    ) -> Corrections {
        match cmd {
            WriteCommand::Batch(batch) => apply_batch(sink_id, corrections, batch).await,
            WriteCommand::WriteBatch(desc) => {
                // Reading the updates back clones a row per update and, with the pager enabled,
                // pages chunks back in, so it stalls the same way the consolidation does and runs
                // on the same blocking thread: the closure keeps the buffers, consolidates, and
                // streams the consolidated updates back in chunks that this task feeds to the
                // persist batch builder.
                //
                // Not `block_in_place`: that parks a worker's core on a blocking-pool thread for
                // the whole write, while the batch's part uploads are spawned tasks that need a
                // live core. With one write per sink and worker, enough concurrent sinks exhaust
                // the pool's thread cap and the runtime deadlocks. `spawn_blocking` queues
                // instead of stranding a core, and this task stays cancellable at every chunk
                // boundary.
                let upper = desc.upper.clone();
                let (updates_tx, mut updates_rx) = mpsc::channel(READ_BACK_CHUNKS_IN_FLIGHT);
                let read_back =
                    mz_ore::task::spawn_blocking(
                        || operator_name(sink_id, "write::consolidate"),
                        move || {
                            corrections.ok.consolidate_before(&upper);
                            corrections.err.consolidate_before(&upper);

                            let oks = corrections
                                .ok
                                .consolidated_updates_before(&upper)
                                .map(|(d, t, r)| ((SourceData(Ok(d)), ()), t, r.into_inner()));
                            let errs = corrections.err.consolidated_updates_before(&upper).map(
                                |(d, t, r)| {
                                    ((SourceData(Err(d.deserialize())), ()), t, r.into_inner())
                                },
                            );

                            let mut updates = oks.chain(errs).peekable();
                            while updates.peek().is_some() {
                                let mut chunk = Vec::with_capacity(READ_BACK_CHUNK);
                                chunk.extend(updates.by_ref().take(READ_BACK_CHUNK));
                                // A closed channel means the write task is gone, so the batch it
                                // asked for is moot and there is no point in pulling the rest of
                                // the buffer.
                                if updates_tx.blocking_send(chunk).is_err() {
                                    break;
                                }
                            }

                            // The iterators borrow the correction buffers, so they must end before
                            // the buffers move back out.
                            drop(updates);
                            corrections
                        },
                    );

                let batches = write_batches(&batch_writer.shared, &desc, &mut updates_rx).await;
                let corrections = read_back.await;

                if let Err(err) = resp_tx.send(WriteResponse { batches }) {
                    for proto_batch in err.0.batches {
                        let batch = batch_writer
                            .writer
                            .batch_from_transmittable_batch(proto_batch);
                        batch.delete().await;
                    }
                }

                corrections
            }
        }
    }

    /// Write the updates for `desc` through `writer`, returning the batches this worker delivers.
    async fn write_batches(
        writer: &SharedBatchWriter,
        desc: &BatchDescription,
        updates_rx: &mut mpsc::Receiver<Vec<((SourceData, ()), Timestamp, StorageDiff)>>,
    ) -> Vec<ProtoBatch> {
        // Start the write only once there are updates, since `maybe_start_batch` requests a write
        // for every ready interval, including idle ones. A worker that never joins a shared batch
        // cannot be the last holder responsible for delivering it.
        let mut write = None;
        while let Some(chunk) = updates_rx.recv().await {
            let write = write.get_or_insert_with(|| writer.start(desc));
            for (kv, t, d) in &chunk {
                write.push(kv, *t, *d).await;
            }
        }
        match write {
            Some(write) => write.finish().await,
            None => Vec::new(),
        }
    }

    /// Apply a coalesced batch of updates to the correction buffers, on a blocking thread.
    ///
    /// See [`apply_command`] for why this work must not run on a Tokio worker thread.
    async fn apply_batch(
        sink_id: GlobalId,
        mut corrections: Corrections,
        mut batch: BatchUpdates,
    ) -> Corrections {
        mz_ore::task::spawn_blocking(
            || operator_name(sink_id, "write::apply_batch"),
            move || {
                // Stands in for the page-in storm this code produces when the process is under
                // memory pressure: the buffer's chunks come back one blocking read at a time, from
                // inside the sorts and merges below, where no await point can be placed. The
                // failpoint's `sleep` action blocks its thread the same way, so it has to sit on
                // the same side of the `spawn_blocking` boundary as the work it stands for.
                //
                // NOTE: activate it with the `FAILPOINTS` env var on the clusterd process, e.g.
                // `FAILPOINTS=mv_sink_correction=sleep(30000)`. The `failpoints` session variable
                // only reaches environmentd, never a cluster.
                fail::fail_point!("mv_sink_correction");

                // Apply the same logical sequence of operations that the per-chunk commands
                // used to: positive desired inserts, negated persist inserts, then optional
                // frontier advancement and forced consolidation.
                if !batch.desired_ok.is_empty() {
                    corrections.ok.insert(&mut batch.desired_ok);
                }
                if !batch.desired_err.is_empty() {
                    corrections.err.insert(&mut batch.desired_err);
                }
                if !batch.persist_ok.is_empty() {
                    corrections.ok.insert_negated(&mut batch.persist_ok);
                }
                if !batch.persist_err.is_empty() {
                    corrections.err.insert_negated(&mut batch.persist_err);
                }
                if let Some(frontier) = batch.persist_frontier {
                    // We will only emit times at or after the `persist` frontier, so now is a good
                    // time to advance the times of stashed updates.
                    corrections.ok.advance_since(frontier.clone());
                    corrections.err.advance_since(frontier);
                }
                if batch.force_consolidation {
                    corrections.ok.consolidate_at_since();
                    corrections.err.consolidate_at_since();
                }
                corrections
            },
        )
        .await
    }

    /// State maintained by the `write` operator on the Timely thread.
    struct State {
        sink_id: GlobalId,
        worker_id: usize,
        /// The frontiers of the `desired` inputs.
        desired_frontiers: OkErr<Antichain<Timestamp>, Antichain<Timestamp>>,
        /// The frontiers of the `persist` inputs.
        ///
        /// Note that this is _not_ the same as the write frontier of the output persist shard! It
        /// usually is, but during snapshot processing, these frontiers will start at the shard's
        /// read frontier, so they can be beyond its write frontier. This is important as it means
        /// we must not discard batch descriptions based on these persist frontiers: A batch
        /// description might still be valid even if its `lower` is before the persist frontiers we
        /// observe.
        persist_frontiers: OkErr<Antichain<Timestamp>, Antichain<Timestamp>>,
        /// The current valid batch description and associated output capability, if any.
        batch_description: Option<(BatchDescription, Capability<Timestamp>)>,
        /// A request to force a consolidation of corrections once both `desired_frontiers` and
        /// `persist_frontiers` become greater than the given frontier.
        ///
        /// Normally we force a consolidation whenever we write a batch, but there are periods
        /// (like read-only mode) when that doesn't happen, and we need to manually force
        /// consolidation instead. In read-only mode this is re-armed after each firing so it
        /// keeps sweeping forward as the frontiers advance (see `should_force_consolidation`).
        force_consolidation_after: Option<Antichain<Timestamp>>,
        /// Whether the sink is in read-only mode. While read-only the `write` operator mints no
        /// batches, so the `WriteBatch` path never sweeps `consolidate_before(upper)` forward;
        /// the forced consolidation stands in for it and is re-armed as long as this holds.
        read_only: bool,
    }

    impl State {
        fn new(
            sink_id: GlobalId,
            worker_id: usize,
            as_of: Antichain<Timestamp>,
            advance_persist_frontiers_at_startup: bool,
            read_only: bool,
        ) -> Self {
            // Force a consolidation of corrections after the snapshot updates have been fully
            // processed, to ensure we get rid of those as quickly as possible.
            let force_consolidation_after = Some(as_of.clone());

            let mut state = Self {
                sink_id,
                worker_id,
                desired_frontiers: OkErr::new_frontiers(),
                persist_frontiers: OkErr::new_frontiers(),
                batch_description: None,
                force_consolidation_after,
                read_only,
            };

            // Immediately advance the persist frontier tracking to the `as_of`.
            // This is important to ensure the persist sink doesn't get stuck if the output shard's
            // initial frontier is less than the `as_of`. The `mint` operator first emits a batch
            // description with `lower = as_of`, and the `write` operator only emits a batch when
            // its observed persist frontier is >= the batch description's `lower`, which (assuming
            // no other writers) would be never if we didn't advance the observed persist frontier
            // to the `as_of`.
            //
            // The `must_use` bool returned by the advance helpers signals that a corresponding
            // `corrections.advance_since` must be queued to the Tokio task. We drop it here
            // because the Tokio-side `corrections.since` is initialized to the same `as_of` in
            // `write::render` before the task is spawned (using the same captured flag value),
            // keeping both sides in lockstep without an additional channel send.
            if advance_persist_frontiers_at_startup {
                let _ = state.advance_persist_ok_frontier(as_of.borrow());
                let _ = state.advance_persist_err_frontier(as_of.borrow());
            }

            state
        }

        fn trace<S: AsRef<str>>(&self, message: S) {
            let message = message.as_ref();
            trace!(
                sink_id = %self.sink_id,
                worker = %self.worker_id,
                desired_frontier = ?self.desired_frontiers.frontier().elements(),
                persist_frontier = ?self.persist_frontiers.frontier().elements(),
                batch_description = ?self.batch_description.as_ref().map(|(d, _)| d),
                message,
            );
        }

        fn advance_desired_ok_frontier(&mut self, frontier: AntichainRef<Timestamp>) {
            if advance(&mut self.desired_frontiers.ok, frontier) {
                self.trace("advanced `desired` ok frontier");
            }
        }

        fn advance_desired_err_frontier(&mut self, frontier: AntichainRef<Timestamp>) {
            if advance(&mut self.desired_frontiers.err, frontier) {
                self.trace("advanced `desired` err frontier");
            }
        }

        /// Returns true if the persist frontier advanced.
        ///
        /// The caller must propagate a `true` return value into the next `WriteCommand::Batch`'s
        /// `persist_frontier` field so the Tokio task advances `corrections.since` accordingly.
        /// Dropping the bool leaves the Tokio-side `since` lagging behind `persist_frontiers`,
        /// which can let updates with timestamps below `desc.lower` slip into a written batch.
        #[must_use = "advance_persist_ok_frontier's return value gates a `corrections.advance_since` send to the Tokio task"]
        fn advance_persist_ok_frontier(&mut self, frontier: AntichainRef<Timestamp>) -> bool {
            if advance(&mut self.persist_frontiers.ok, frontier) {
                self.trace("advanced `persist` ok frontier");
                true
            } else {
                false
            }
        }

        /// Returns true if the persist frontier advanced.
        ///
        /// The caller must propagate a `true` return value into the next `WriteCommand::Batch`'s
        /// `persist_frontier` field so the Tokio task advances `corrections.since` accordingly.
        /// Dropping the bool leaves the Tokio-side `since` lagging behind `persist_frontiers`,
        /// which can let updates with timestamps below `desc.lower` slip into a written batch.
        #[must_use = "advance_persist_err_frontier's return value gates a `corrections.advance_since` send to the Tokio task"]
        fn advance_persist_err_frontier(&mut self, frontier: AntichainRef<Timestamp>) -> bool {
            if advance(&mut self.persist_frontiers.err, frontier) {
                self.trace("advanced `persist` err frontier");
                true
            } else {
                false
            }
        }

        /// Check if a forced consolidation should be triggered.
        fn should_force_consolidation(&mut self) -> bool {
            let Some(request) = &self.force_consolidation_after else {
                return false;
            };

            let desired_frontier = self.desired_frontiers.frontier();
            let persist_frontier = self.persist_frontiers.frontier();
            if PartialOrder::less_than(request, desired_frontier)
                && PartialOrder::less_than(request, persist_frontier)
            {
                self.trace("requesting correction consolidation");
                // In read-only mode the sink mints no batches, so the `WriteBatch` path never
                // sweeps `consolidate_before(upper)` forward and a single forced consolidation
                // only covers the band below the `since` at the moment it fires, leaving the
                // forward region uncancelled for the whole read-only window (CLU-131). Re-arm at
                // the current persist frontier so the forced consolidation keeps sweeping forward
                // as `desired`/`persist` advance, mirroring the read-write path. Once writes are
                // allowed, the `WriteBatch` path takes over and we disarm.
                self.force_consolidation_after = if self.read_only {
                    Some(persist_frontier.to_owned())
                } else {
                    None
                };
                true
            } else {
                false
            }
        }

        /// Update read-only mode. While read-only the forced consolidation re-arms itself; once
        /// writes are allowed the still-armed request fires one final time and then disarms, after
        /// which the `WriteBatch` path drives consolidation.
        fn set_read_only(&mut self, read_only: bool) {
            self.read_only = read_only;
        }

        fn absorb_batch_description(&mut self, desc: BatchDescription, cap: Capability<Timestamp>) {
            // Enforce monotonicity: drop descriptions whose `lower` regresses below the one we
            // already hold. The `mint` operator only emits strictly increasing `lower`s
            // (invariant 1), so a regression means this description is outdated. We cannot use
            // `persist_frontiers` for the same check, because during snapshot processing those
            // frontiers can be ahead of the shard's write frontier and a still-valid description
            // may have a `lower` below them.
            if let Some((prev, _)) = &self.batch_description {
                if PartialOrder::less_than(&desc.lower, &prev.lower) {
                    self.trace(format!("skipping outdated batch description: {desc:?}"));
                    return;
                }
            }

            self.batch_description = Some((desc, cap));
            self.trace("set batch description");
        }

        /// Check if a batch can be written and send a write command to the Tokio task if so.
        fn maybe_start_batch(
            &mut self,
            cmd_tx: &mpsc::UnboundedSender<WriteCommand>,
        ) -> Option<(BatchDescription, Capability<Timestamp>)> {
            let (desc, _cap) = self.batch_description.as_ref()?;

            // We can write a new batch if we have seen all `persist` updates before `lower` and
            // all `desired` updates before `upper`.
            let persist_ready =
                PartialOrder::less_equal(&desc.lower, self.persist_frontiers.frontier());
            let desired_ready =
                PartialOrder::less_equal(&desc.upper, self.desired_frontiers.frontier());
            if !persist_ready || !desired_ready {
                return None;
            }

            self.trace("write batch description");
            let (desc, cap) = self.batch_description.take()?;
            cmd_tx
                .send(WriteCommand::WriteBatch(desc.clone()))
                .expect("write task unexpectedly gone");
            Some((desc, cap))
        }
    }

    #[cfg(test)]
    mod tests {
        use std::sync::Arc;
        use std::sync::atomic::{AtomicU64, Ordering};
        use std::time::{Duration, Instant};

        use mz_ore::metrics::MetricsRegistry;
        use mz_persist_client::cfg::PersistConfig;
        use mz_persist_client::metrics::Metrics;

        use super::*;

        /// One stall per worker thread would leave scheduling to chance, so oversubscribe: with
        /// more stalls than workers every worker is guaranteed to pick one up.
        const WORKER_THREADS: usize = 2;
        const STALLS: usize = WORKER_THREADS * 2;
        /// How long a stalled `apply_batch` blocks its thread.
        const STALL: Duration = Duration::from_millis(1_000);
        /// The largest tick gap the canary may observe. Generous enough to absorb scheduling noise
        /// on a loaded CI host, while far below the `STALL` an inline correction pass produces.
        const MAX_GAP: Duration = Duration::from_millis(500);
        const TICK: Duration = Duration::from_millis(10);

        fn corrections() -> Corrections {
            let registry = MetricsRegistry::new();
            let metrics = Metrics::new(&PersistConfig::new_for_tests(), &registry);
            let sink_metrics = metrics.sink.clone();
            let worker_metrics = sink_metrics.for_worker(0);
            let config = mz_dyncfgs::all_dyncfgs();
            OkErr::new(
                Correction::new(sink_metrics.clone(), worker_metrics.clone(), None, &config),
                Correction::new(sink_metrics, worker_metrics, None, &config),
            )
        }

        /// Tick every [`TICK`], recording the largest gap between consecutive ticks in millis.
        ///
        /// A gap only opens up if no worker thread was available to poll this task.
        async fn canary(max_gap_millis: Arc<AtomicU64>) {
            let mut last = Instant::now();
            loop {
                tokio::time::sleep(TICK).await;
                let now = Instant::now();
                let gap = u64::try_from(now.duration_since(last).as_millis()).unwrap_or(u64::MAX);
                max_gap_millis.fetch_max(gap, Ordering::Relaxed);
                last = now;
            }
        }

        /// Correction maintenance must not occupy a Tokio worker thread.
        ///
        /// Run inline, a pass over a large buffer pins a worker for its whole duration. Under
        /// memory pressure it is blocked in the kernel paging chunks back in rather than
        /// computing, which is why yielding cannot fix it.
        ///
        /// Stalls every worker thread inside `apply_batch` and asserts an unrelated task keeps
        /// getting polled throughout.
        #[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 2))]
        #[cfg_attr(miri, ignore)] // depends on real thread scheduling
        async fn correction_work_does_not_stall_the_runtime() {
            let max_gap_millis = Arc::new(AtomicU64::new(0));
            let _canary = mz_ore::task::spawn(|| "canary", canary(Arc::clone(&max_gap_millis)))
                .abort_on_drop();

            // Let the canary settle, then discard the gaps observed while it did.
            tokio::time::sleep(TICK * 5).await;
            max_gap_millis.store(0, Ordering::Relaxed);

            let action = format!("sleep({})", STALL.as_millis());
            fail::cfg("mv_sink_correction", &action).expect("valid failpoint action");

            let stalls: Vec<_> = (0..STALLS)
                .map(|i| {
                    // A forced consolidation is the operation that hurt: it sweeps the whole
                    // buffer, so its cost is unbounded by the size of the incoming batch.
                    let batch = BatchUpdates {
                        force_consolidation: true,
                        ..BatchUpdates::new()
                    };
                    let sink_id = GlobalId::User(u64::cast_from(i));
                    mz_ore::task::spawn(|| "stall", apply_batch(sink_id, corrections(), batch))
                })
                .collect();
            for stall in stalls {
                let _ = stall.await;
            }

            fail::remove("mv_sink_correction");

            let max_gap = Duration::from_millis(max_gap_millis.load(Ordering::Relaxed));
            assert!(
                max_gap < MAX_GAP,
                "runtime went unpolled for {max_gap:?}, so correction work is occupying a \
                 worker thread",
            );
        }
    }
}

/// Implementation of the `append` operator.
mod append {
    use super::*;

    /// Commands sent from the Timely operator to the Tokio append task.
    enum AppendCommand {
        /// A new batch description has been received.
        Description(BatchDescription),
        /// A written batch has been received.
        Batch(ProtoBatch),
        /// The batches frontier has advanced.
        BatchesFrontier(Antichain<Timestamp>),
    }

    /// Render the `append` operator.
    ///
    /// The parameters passed in are:
    ///  * `sink_id`: The `GlobalId` of the sink export.
    ///  * `persist_api`: An object providing access to the output persist shard.
    ///  * `descs`: The stream of batch descriptions produced by the `mint` operator.
    ///  * `batches`: The stream of written batches produced by the `write` operator.
    pub fn render<'s>(
        sink_id: GlobalId,
        persist_api: PersistApi,
        descs: DescsStream<'s>,
        batches: BatchesStream<'s>,
    ) {
        let scope = descs.scope();
        let worker_id = scope.index();

        let name = operator_name(sink_id, "append");
        let mut builder = OperatorBuilderRc::new(name, scope.clone());
        let mut descs_input = builder.new_input(descs, Pipeline);
        let batch_exchange =
            Exchange::new(|(desc, _): &(BatchDescription, _)| u64::cast_from(desc.append_worker));
        let mut batches_input = builder.new_input(batches, batch_exchange);

        // Channel for commands to the Tokio append task.
        let (cmd_tx, mut cmd_rx) = mpsc::unbounded_channel::<AppendCommand>();

        // Spawn Tokio task that owns the append state machine.
        let append_task_handle =
            mz_ore::task::spawn(|| operator_name(sink_id, "append"), async move {
                let writer = persist_api.open_writer().await;
                let mut state = State::new(sink_id, worker_id, writer);

                while let Some(cmd) = cmd_rx.recv().await {
                    match cmd {
                        AppendCommand::Description(desc) => {
                            state.absorb_batch_description(desc).await;
                            state.maybe_append_batches().await;
                        }
                        AppendCommand::Batch(batch) => {
                            state.absorb_batch(batch).await;
                        }
                        AppendCommand::BatchesFrontier(frontier) => {
                            state.advance_batches_frontier(frontier.borrow());
                            state.maybe_append_batches().await;
                        }
                    }
                }
            })
            .abort_on_drop();

        builder.build(move |_capabilities| {
            let mut prev_batches_frontier = Antichain::from_elem(Timestamp::MIN);

            move |frontiers| {
                // Keep task handle alive so it is aborted when the operator is dropped.
                let _ = &append_task_handle;

                // Forward batch descriptions to the Tokio task.
                descs_input.for_each(|_cap, data| {
                    for desc in data.drain(..) {
                        cmd_tx
                            .send(AppendCommand::Description(desc))
                            .expect("append task unexpectedly gone");
                    }
                });

                // Forward batches to the Tokio task.
                batches_input.for_each(|_cap, data| {
                    for (_desc, batch) in data.drain(..) {
                        // The batch description is only used for routing and we ignore it
                        // here since we already get one from `descs_input`.
                        cmd_tx
                            .send(AppendCommand::Batch(batch))
                            .expect("append task unexpectedly gone");
                    }
                });

                // Forward batches frontier advancements *after* the per-activation
                // `Description`/`Batch` sends above. The Tokio task drains commands FIFO and only
                // calls `maybe_append_batches` on `Description`/`BatchesFrontier`; if a frontier
                // advance arrived before its batches, the task could append an incomplete set.
                // See module-level docs for the full ordering invariant.
                let new_batches_frontier = frontiers[1].frontier();
                if PartialOrder::less_than(&prev_batches_frontier.borrow(), &new_batches_frontier) {
                    prev_batches_frontier.clear();
                    prev_batches_frontier.extend(new_batches_frontier.iter().cloned());
                    cmd_tx
                        .send(AppendCommand::BatchesFrontier(
                            new_batches_frontier.to_owned(),
                        ))
                        .expect("append task unexpectedly gone");
                }
            }
        });
    }

    /// State maintained by the `append` Tokio task.
    struct State {
        sink_id: GlobalId,
        worker_id: usize,
        persist_writer: WriteHandle<SourceData, (), Timestamp, StorageDiff>,
        /// The current input frontier of `batches`.
        batches_frontier: Antichain<Timestamp>,
        /// The greatest observed `lower` from both `descs` and `batches`.
        lower: Antichain<Timestamp>,
        /// The batch description for `lower`, if any.
        batch_description: Option<BatchDescription>,
        /// Batches received for `lower`.
        batches: Vec<Batch<SourceData, (), Timestamp, StorageDiff>>,
    }

    impl State {
        fn new(
            sink_id: GlobalId,
            worker_id: usize,
            persist_writer: WriteHandle<SourceData, (), Timestamp, StorageDiff>,
        ) -> Self {
            Self {
                sink_id,
                worker_id,
                persist_writer,
                batches_frontier: Antichain::from_elem(Timestamp::MIN),
                lower: Antichain::from_elem(Timestamp::MIN),
                batch_description: None,
                batches: Default::default(),
            }
        }

        fn trace<S: AsRef<str>>(&self, message: S) {
            let message = message.as_ref();
            trace!(
                sink_id = %self.sink_id,
                worker = %self.worker_id,
                batches_frontier = ?self.batches_frontier.elements(),
                lower = ?self.lower.elements(),
                batch_description = ?self.batch_description,
                message,
            );
        }

        fn advance_batches_frontier(&mut self, frontier: AntichainRef<Timestamp>) {
            if advance(&mut self.batches_frontier, frontier) {
                self.trace("advanced `batches` frontier");
            }
        }

        /// Advance the current `lower`.
        ///
        /// Discards all currently stashed batches and batch descriptions, assuming that they are
        /// now invalid.
        async fn advance_lower(&mut self, frontier: Antichain<Timestamp>) {
            assert!(PartialOrder::less_than(&self.lower, &frontier));

            self.lower = frontier;
            self.batch_description = None;

            // Remove stashed batches, cleaning up those we didn't append.
            for batch in self.batches.drain(..) {
                batch.delete().await;
            }

            self.trace("advanced `lower`");
        }

        /// Absorb the given batch description into the state, provided it is not outdated.
        async fn absorb_batch_description(&mut self, desc: BatchDescription) {
            if PartialOrder::less_than(&self.lower, &desc.lower) {
                self.advance_lower(desc.lower.clone()).await;
            } else if &self.lower != &desc.lower {
                self.trace(format!("skipping outdated batch description: {desc:?}"));
                return;
            }

            if desc.append_worker == self.worker_id {
                self.batch_description = Some(desc);
                self.trace("set batch description");
            }
        }

        /// Absorb the given batch into the state, provided it is not outdated.
        async fn absorb_batch(&mut self, batch: ProtoBatch) {
            let batch = self.persist_writer.batch_from_transmittable_batch(batch);
            if PartialOrder::less_than(&self.lower, batch.lower()) {
                self.advance_lower(batch.lower().clone()).await;
            } else if &self.lower != batch.lower() {
                self.trace(format!(
                    "skipping outdated batch: ({:?}, {:?})",
                    batch.lower().elements(),
                    batch.upper().elements(),
                ));

                // Ensure the batch's data gets properly cleaned up before dropping it.
                batch.delete().await;
                return;
            }

            self.batches.push(batch);
            self.trace("absorbed a batch");
        }

        async fn maybe_append_batches(&mut self) {
            let batches_complete = PartialOrder::less_than(&self.lower, &self.batches_frontier);
            if !batches_complete {
                return;
            }

            let Some(desc) = self.batch_description.take() else {
                return;
            };

            let new_lower = match self.append_batches(desc).await {
                Ok(shard_upper) => {
                    self.trace("appended a batch");
                    shard_upper
                }
                Err(shard_upper) => {
                    // Failing the append is expected in the presence of concurrent replicas. There
                    // is nothing special to do here: The self-correcting feedback mechanism
                    // ensures that we observe the concurrent changes, compute their consequences,
                    // and append them at a future time.
                    self.trace(format!(
                        "append failed due to `lower` mismatch: {:?}",
                        shard_upper.elements(),
                    ));
                    shard_upper
                }
            };

            self.advance_lower(new_lower).await;
        }

        /// Append the current `batches` to the output shard.
        ///
        /// Returns whether the append was successful or not, and the current shard upper in either
        /// case.
        ///
        /// This method advances the shard upper to the batch `lower` if necessary. This is the
        /// mechanism that brings the shard upper to the sink as-of when appending the initial
        /// batch.
        ///
        /// An alternative mechanism for bringing the shard upper to the sink as-of would be making
        /// a single append at operator startup. The reason we are doing it here instead is that it
        /// simplifies the implementation of read-only mode. In read-only mode we have to defer any
        /// persist writes, including the initial upper bump. Having only a single place that
        /// performs writes makes it easy to ensure we are doing that correctly.
        async fn append_batches(
            &mut self,
            desc: BatchDescription,
        ) -> Result<Antichain<Timestamp>, Antichain<Timestamp>> {
            let (lower, upper) = (desc.lower, desc.upper);
            let mut to_append: Vec<_> = self.batches.iter_mut().collect();

            loop {
                let result = self
                    .persist_writer
                    .compare_and_append_batch(&mut to_append, lower.clone(), upper.clone(), true)
                    .await
                    .expect("valid usage");

                match result {
                    Ok(()) => return Ok(upper),
                    Err(mismatch) if PartialOrder::less_than(&mismatch.current, &lower) => {
                        advance_shard_upper(&mut self.persist_writer, lower.clone()).await;

                        // At this point the shard's since and upper are likely the same, a state
                        // that is likely to hit edge-cases in logic reasoning about frontiers.
                        fail::fail_point!("mv_advanced_upper");
                    }
                    Err(mismatch) => return Err(mismatch.current),
                }
            }
        }
    }

    /// Advance the frontier of the given writer's shard to at least the given `upper`.
    async fn advance_shard_upper(
        persist_writer: &mut WriteHandle<SourceData, (), Timestamp, StorageDiff>,
        upper: Antichain<Timestamp>,
    ) {
        let empty_updates: &[((SourceData, ()), Timestamp, StorageDiff)] = &[];
        let lower = Antichain::from_elem(Timestamp::MIN);
        persist_writer
            .append(empty_updates, lower, upper)
            .await
            .expect("valid usage")
            .expect("should always succeed");
    }
}