Skip to main content

mz_compute/sink/
materialized_view.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! A dataflow sink that writes input records to a persist shard.
11//!
12//! This implementation is both parallel and self-correcting.
13//!
14//!  * parallel: Multiple workers can participate in writing updates for the same times, letting
15//!    sink throughput scale with the number of workers allocated to the replica.
16//!  * self-correcting: The sink continually compares the contents of the persist shard with the
17//!    contents of the input collection and writes down the difference. If the persist shard ends
18//!    up with undesired contents for any reason, this is corrected the next time the sink manages
19//!    to append to the shard.
20//!
21//! ### Operators
22//!
23//! The persist sink consists of a graph of operators.
24//!
25//!    desired                    persist <---------------.
26//!       |                          |                    |
27//!       |                          |                    |
28//!       |---------------------.    |                    |
29//!       |                     |    |                    |
30//!       |                     |    |                    |
31//!       v                     v    v                    |
32//!   +--------+              +--------+              +--------+
33//!   |  mint  | --descs-.--> | write  | --batches--> | append |
34//!   +--------+          \   +--------+          .-> +--------+
35//!                        \_____________________/
36//!
37//!  * `mint` mints batch descriptions, i.e., `(lower, upper)` bounds of batches that should be
38//!    written. The persist API requires that all workers write batches with the same bounds, so
39//!    they can be appended as a single logical batch. To ensure this, the `mint` operator only
40//!    runs on a single worker that broadcasts minted descriptions to all workers. Batch bounds are
41//!    picked based on the frontiers of the `desired` stream and the output persist shard.
42//!  * `write` stages batch data in persist, based on the batch descriptions received from the
43//!    `mint` operator, but without appending it to the persist shard. This is a multi-worker
44//!    operator, with each worker writing batches of the data that arrives at its local inputs. To
45//!    do so it reads from the `desired` and `persist` streams and produces the difference between
46//!    them to write back out, ensuring that the final contents of the persist shard match
47//!    `desired`.
48//!  * `append` appends the batches minted by `mint` and written by `write` to the persist shard.
49//!    This is a multi-worker operator, where workers are responsible for different subsets of
50//!    batch descriptions. If a worker is responsible for a given batch description, it waits for
51//!    all workers to stage their batches for that batch description, then appends all the batches
52//!    together as a single logical batch.
53//!
54//! Note that while the above graph suggests that `mint` and `write` both receive copies of the
55//! `desired` stream, the actual implementation passes that stream through `mint` and lets `write`
56//! read the passed-through stream, to avoid cloning data.
57//!
58//! Also note that the `append` operator's implementation would perhaps be more natural as a
59//! single-worker implementation. The purpose of sharing the work between all workers is to avoid a
60//! work imbalance where one worker is overloaded (doing both appends and the consequent persist
61//! maintenance work) while others are comparatively idle.
62//!
63//! The persist sink is written to be robust to the presence of other conflicting instances (e.g.
64//! from other replicas) writing to the same persist shard. Each of the three operators needs to be
65//! able to handle conflicting writes that unexpectedly change the contents of the output persist
66//! shard.
67//!
68//! ### Frontiers
69//!
70//! The `desired` frontier tracks the progress of the upstream dataflow, but may be rounded up to
71//! the next refresh time for dataflows that follow a refresh schedule other than "on commit".
72//!
73//! The `persist` frontier tracks the `upper` frontier of the target persist shard, with one
74//! exception: When the `persist_source` that reads back the shard is rendered, it will start
75//! reading at its `since` frontier. So if the shard's `since` is initially greater than its
76//! `upper`, the `persist` frontier too will be in advance of the shard `upper`, until the `upper`
77//! has caught up. To avoid getting confused by this edge case, the `mint` operator does not use
78//! the `persist` stream to observe the shard frontier but keeps its own `WriteHandle` instead.
79//!
80//! The `descs` frontier communicates which `lower` bounds may still be emitted in batch
81//! descriptions. All future batch descriptions will have a `lower` that is greater or equal to the
82//! current `descs` frontier.
83//!
84//! The `batches` frontier communicates for which `lower` bounds batches may still be written. All
85//! batches for descriptions with `lower`s less than the current `batches` frontier have already
86//! been written.
87//!
88//! ### Invariants
89//!
90//! The implementation upholds several invariants that can be relied upon to simplify the
91//! implementation:
92//!
93//!  1. `lower`s in minted batch descriptions are unique and strictly increasing. That is, the
94//!     `mint` operator will never mint the same `lower` twice and a minted `lower` is always
95//!     greater than any previously minted ones.
96//!  2. `upper`s in minted batch descriptions are monotonically increasing.
97//!  3. From (1) follows that there is always at most one "valid" batch description in flight in
98//!     the operator graph. "Valid" here means that the described batch can be appended to the
99//!     persist shard.
100//!
101//! The main simplification these invariants allow is that operators only need to keep track of the
102//! most recent batch description and/or `lower`. Previous batch descriptions are not valid
103//! anymore, so there is no reason to hold any state or perform any work in support of them.
104//!
105//! ### Read-only Mode
106//!
107//! The persist sink can optionally be initialized in read-only mode. In this mode it is passive
108//! and avoids any writes to persist. Activating the `read_only_rx` transitions the sink into write
109//! mode, where it commences normal operation.
110//!
111//! Read-only mode is implemented by the `mint` operator. To disable writes, the `mint` operator
112//! simply avoids minting any batch descriptions. Since both the `write` and the `append` operator
113//! require batch descriptions to write/append batches, this suppresses any persist communication.
114//! At the same time, the `write` operator still observes changes to the `desired` and `persist`
115//! collections, allowing it to keep its correction buffer up-to-date.
116
117use std::any::Any;
118use std::cell::RefCell;
119use std::pin::pin;
120use std::rc::Rc;
121use std::sync::Arc;
122
123use differential_dataflow::{AsCollection, Hashable, VecCollection};
124use futures::StreamExt;
125use mz_compute_types::dyncfgs::{
126    MV_SINK_ADVANCE_PERSIST_FRONTIERS, MV_SINK_SHARED_BATCH_MAX_BYTES,
127};
128use mz_compute_types::sinks::{ComputeSinkDesc, MaterializedViewSinkConnection};
129use mz_dyncfg::ConfigSet;
130use mz_ore::cast::CastFrom;
131use mz_persist_client::batch::{Batch, ProtoBatch};
132use mz_persist_client::cache::PersistClientCache;
133use mz_persist_client::metrics::SinkMetrics;
134use mz_persist_client::operators::shard_source::{ErrorHandler, SnapshotMode};
135use mz_persist_client::write::WriteHandle;
136use mz_persist_client::{Diagnostics, PersistClient};
137use mz_persist_types::codec_impls::UnitSchema;
138use mz_repr::{Diff, GlobalId, Row, Timestamp};
139use mz_storage_operators::persist::{SharedBatchId, SharedBatches};
140use mz_storage_types::StorageDiff;
141use mz_storage_types::controller::CollectionMetadata;
142use mz_storage_types::sources::SourceData;
143use mz_timely_util::builder_async::PressOnDropButton;
144use mz_timely_util::builder_async::{Event, OperatorBuilder};
145use mz_timely_util::probe::{Handle, ProbeNotify};
146use serde::{Deserialize, Serialize};
147use timely::PartialOrder;
148use timely::container::CapacityContainerBuilder;
149use timely::dataflow::channels::pact::{Exchange, Pipeline};
150use timely::dataflow::operators::vec::Broadcast;
151use timely::dataflow::operators::{Capability, CapabilitySet, probe};
152use timely::dataflow::{Scope, StreamVec};
153use timely::progress::Antichain;
154use tokio::sync::watch;
155use tracing::trace;
156
157use crate::compute_state::ComputeState;
158use crate::render::StartSignal;
159use crate::render::errors::DataflowErrorSer;
160use crate::render::sinks::SinkRender;
161use crate::sink::correction::{ChannelLogging, Correction, CorrectionLogger};
162use crate::sink::materialized_view_v2;
163use crate::sink::refresh::apply_refresh;
164use crate::sink::shared_batch::SharedBatchWriter;
165
166impl<'scope> SinkRender<'scope> for MaterializedViewSinkConnection<CollectionMetadata> {
167    fn render_sink(
168        &self,
169        compute_state: &mut ComputeState,
170        sink: &ComputeSinkDesc<CollectionMetadata>,
171        sink_id: GlobalId,
172        as_of: Antichain<Timestamp>,
173        start_signal: StartSignal,
174        mut ok_collection: VecCollection<'scope, Timestamp, Row, Diff>,
175        mut err_collection: VecCollection<'scope, Timestamp, DataflowErrorSer, Diff>,
176        output_probe: &Handle<Timestamp>,
177    ) -> Option<Rc<dyn Any>> {
178        // Attach probes reporting the compute frontier.
179        // The `apply_refresh` operator can round up frontiers, making it impossible to accurately
180        // track the progress of the computation, so we need to attach probes before it.
181        let probe = probe::Handle::default();
182        ok_collection = ok_collection
183            .probe_with(&probe)
184            .inner
185            .probe_notify_with(vec![output_probe.clone()])
186            .as_collection();
187        let collection_state = compute_state.expect_collection_mut(sink_id);
188        collection_state.compute_probe = Some(probe);
189
190        // If a `RefreshSchedule` was specified, round up timestamps.
191        if let Some(refresh_schedule) = &sink.refresh_schedule {
192            ok_collection = apply_refresh(ok_collection, refresh_schedule.clone());
193            err_collection = apply_refresh(err_collection, refresh_schedule.clone());
194        }
195
196        if sink.up_to != Antichain::default() {
197            unimplemented!(
198                "UP TO is not supported for persist sinks yet, and shouldn't have been accepted during parsing/planning"
199            )
200        }
201
202        let read_only_rx = collection_state.read_only_rx.clone();
203
204        let token = persist_sink(
205            sink_id,
206            &self.storage_metadata,
207            ok_collection,
208            err_collection,
209            as_of,
210            compute_state,
211            start_signal,
212            read_only_rx,
213        );
214        Some(token)
215    }
216}
217
218/// Type of the `desired` stream, split into `Ok` and `Err` streams.
219pub(super) type DesiredStreams<'s> = OkErr<
220    StreamVec<'s, Timestamp, (Row, Timestamp, Diff)>,
221    StreamVec<'s, Timestamp, (DataflowErrorSer, Timestamp, Diff)>,
222>;
223
224/// Type of the `persist` stream, split into `Ok` and `Err` streams.
225pub(super) type PersistStreams<'s> = OkErr<
226    StreamVec<'s, Timestamp, (Row, Timestamp, Diff)>,
227    StreamVec<'s, Timestamp, (DataflowErrorSer, Timestamp, Diff)>,
228>;
229
230/// Type of the `descs` stream.
231pub(super) type DescsStream<'s> = StreamVec<'s, Timestamp, BatchDescription>;
232
233/// Type of the `batches` stream.
234pub(super) type BatchesStream<'s> = StreamVec<'s, Timestamp, (BatchDescription, ProtoBatch)>;
235
236/// Type of the shared sink write frontier.
237pub(super) type SharedSinkFrontier = Rc<RefCell<Antichain<Timestamp>>>;
238
239/// Renders an MV sink writing the given desired collection into the `target` persist collection.
240pub(super) fn persist_sink<'s>(
241    sink_id: GlobalId,
242    target: &CollectionMetadata,
243    ok_collection: VecCollection<'s, Timestamp, Row, Diff>,
244    err_collection: VecCollection<'s, Timestamp, DataflowErrorSer, Diff>,
245    as_of: Antichain<Timestamp>,
246    compute_state: &mut ComputeState,
247    start_signal: StartSignal,
248    read_only_rx: watch::Receiver<bool>,
249) -> Rc<dyn Any>
250where
251{
252    if mz_compute_types::dyncfgs::ENABLE_SYNC_MV_SINK.get(&compute_state.worker_config) {
253        return materialized_view_v2::persist_sink(
254            sink_id,
255            target,
256            ok_collection,
257            err_collection,
258            as_of,
259            compute_state,
260            start_signal,
261            read_only_rx,
262        );
263    }
264
265    let scope = ok_collection.scope();
266    let desired = OkErr::new(ok_collection.inner, err_collection.inner);
267
268    // Read back the persist shard.
269    let (persist, persist_token) =
270        persist_source(scope, sink_id, target.clone(), compute_state, start_signal);
271
272    let persist_api = PersistApi {
273        persist_clients: Arc::clone(&compute_state.persist_clients),
274        persist_batches: compute_state.persist_batches.clone(),
275        collection: target.clone(),
276        shard_name: sink_id.to_string(),
277        purpose: format!("MV sink {sink_id}"),
278    };
279
280    let (desired, descs, sink_frontier, mint_token) = mint::render(
281        sink_id,
282        persist_api.clone(),
283        as_of.clone(),
284        read_only_rx.clone(),
285        desired,
286    );
287
288    let (batches, write_token) = write::render(
289        sink_id,
290        persist_api.clone(),
291        as_of,
292        desired,
293        persist,
294        descs.clone(),
295        read_only_rx,
296        Rc::clone(&compute_state.worker_config),
297    );
298
299    let append_token = append::render(sink_id, persist_api, descs, batches);
300
301    // Report sink frontier updates to the `ComputeState`.
302    let collection = compute_state.expect_collection_mut(sink_id);
303    collection.sink_write_frontier = Some(sink_frontier);
304
305    Rc::new((persist_token, mint_token, write_token, append_token))
306}
307
308/// Generic wrapper around ok/err pairs (e.g. streams, frontiers), to simplify code dealing with
309/// such pairs.
310pub(super) struct OkErr<O, E> {
311    pub(super) ok: O,
312    pub(super) err: E,
313}
314
315impl<O, E> OkErr<O, E> {
316    pub(super) fn new(ok: O, err: E) -> Self {
317        Self { ok, err }
318    }
319}
320
321impl OkErr<Antichain<Timestamp>, Antichain<Timestamp>> {
322    pub(super) fn new_frontiers() -> Self {
323        Self {
324            ok: Antichain::from_elem(Timestamp::MIN),
325            err: Antichain::from_elem(Timestamp::MIN),
326        }
327    }
328
329    /// Return the overall frontier, i.e., the minimum of `ok` and `err`.
330    pub(super) fn frontier(&self) -> &Antichain<Timestamp> {
331        if PartialOrder::less_equal(&self.ok, &self.err) {
332            &self.ok
333        } else {
334            &self.err
335        }
336    }
337}
338
339/// Advance the given `frontier` to `new`, if the latter one is greater.
340///
341/// Returns whether `frontier` was advanced.
342pub(super) fn advance(
343    frontier: &mut Antichain<Timestamp>,
344    new: timely::progress::frontier::AntichainRef<'_, Timestamp>,
345) -> bool {
346    if PartialOrder::less_than(&frontier.borrow(), &new) {
347        frontier.clear();
348        frontier.extend(new.iter().cloned());
349        true
350    } else {
351        false
352    }
353}
354
355/// A persist API specialized to a single collection.
356#[derive(Clone)]
357pub(super) struct PersistApi {
358    pub(super) persist_clients: Arc<PersistClientCache>,
359    pub(super) persist_batches: SharedBatches,
360    pub(super) collection: CollectionMetadata,
361    pub(super) shard_name: String,
362    pub(super) purpose: String,
363}
364
365impl PersistApi {
366    pub(super) async fn open_client(&self) -> PersistClient {
367        self.persist_clients
368            .open(self.collection.persist_location.clone())
369            .await
370            .unwrap_or_else(|error| panic!("error opening persist client: {error}"))
371    }
372
373    pub(super) async fn open_writer(&self) -> WriteHandle<SourceData, (), Timestamp, StorageDiff> {
374        self.open_client()
375            .await
376            .open_writer(
377                self.collection.data_shard,
378                Arc::new(self.collection.relation_desc.clone()),
379                Arc::new(UnitSchema),
380                Diagnostics {
381                    shard_name: self.shard_name.clone(),
382                    handle_purpose: self.purpose.clone(),
383                },
384            )
385            .await
386            .unwrap_or_else(|error| panic!("error opening persist writer: {error}"))
387    }
388
389    async fn open_metrics(&self) -> SinkMetrics {
390        let client = self.open_client().await;
391        client.metrics().sink.clone()
392    }
393}
394
395/// Instantiate a persist source reading back the `target` collection.
396pub(super) fn persist_source<'s>(
397    scope: Scope<'s, Timestamp>,
398    sink_id: GlobalId,
399    target: CollectionMetadata,
400    compute_state: &ComputeState,
401    start_signal: StartSignal,
402) -> (PersistStreams<'s>, Vec<PressOnDropButton>) {
403    // There is no guarantee that the sink as-of is beyond the persist shard's since. If it isn't,
404    // instantiating a `persist_source` with it would panic. So instead we leave it to
405    // `persist_source` to select an appropriate as-of. We only care about times beyond the current
406    // shard upper anyway.
407    //
408    // TODO(teskje): Ideally we would select the as-of as `join(sink_as_of, since, upper)`, to
409    // allow `persist_source` to omit as much historical detail as possible. However, we don't know
410    // the shard frontiers and we cannot get them here as that requires an `async` context. We
411    // should consider extending the `persist_source` API to allow as-of selection based on the
412    // shard's current frontiers.
413    let as_of = None;
414
415    let until = Antichain::new();
416    let map_filter_project = None;
417
418    let (ok_stream, err_stream, token) = mz_storage_operators::persist_source::persist_source::<
419        DataflowErrorSer,
420        mz_storage_operators::persist_source::RowVecBuilder<Timestamp>,
421    >(
422        scope,
423        sink_id,
424        Arc::clone(&compute_state.persist_clients),
425        &compute_state.txns_ctx,
426        target,
427        None,
428        as_of,
429        SnapshotMode::Include,
430        until,
431        map_filter_project,
432        compute_state.dataflow_max_inflight_bytes(),
433        start_signal.into_send_future(),
434        ErrorHandler::Halt("compute persist sink"),
435    );
436
437    let streams = OkErr::new(ok_stream, err_stream);
438    (streams, token)
439}
440
441/// A description for a batch of updates to be written.
442///
443/// Batch descriptions are produced by the `mint` operator and consumed by the `write` and `append`
444/// operators, where they inform which batches should be written or appended, respectively.
445///
446/// Each batch description also contains the index of its "append worker", i.e. the worker that is
447/// responsible for appending the written batches to the output shard.
448#[derive(Clone, Serialize, Deserialize)]
449pub(super) struct BatchDescription {
450    pub(super) lower: Antichain<Timestamp>,
451    pub(super) upper: Antichain<Timestamp>,
452    pub(super) append_worker: usize,
453    /// Identifies the shared batch all workers building this description contribute to.
454    pub(super) shared_id: SharedBatchId,
455}
456
457impl BatchDescription {
458    pub(super) fn new(
459        lower: Antichain<Timestamp>,
460        upper: Antichain<Timestamp>,
461        append_worker: usize,
462    ) -> Self {
463        assert!(PartialOrder::less_than(&lower, &upper));
464        Self {
465            lower,
466            upper,
467            append_worker,
468            shared_id: SharedBatchId::new(),
469        }
470    }
471}
472
473impl std::fmt::Debug for BatchDescription {
474    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
475        write!(
476            f,
477            "({:?}, {:?})@{} {:?}",
478            self.lower.elements(),
479            self.upper.elements(),
480            self.append_worker,
481            self.shared_id,
482        )
483    }
484}
485
486/// Construct a name for the given sub-operator.
487pub(super) fn operator_name(sink_id: GlobalId, sub_operator: &str) -> String {
488    format!("mv_sink({sink_id})::{sub_operator}")
489}
490
491/// Implementation of the `mint` operator.
492mod mint {
493    use super::*;
494
495    /// Render the `mint` operator.
496    ///
497    /// The parameters passed in are:
498    ///  * `sink_id`: The `GlobalId` of the sink export.
499    ///  * `persist_api`: An object providing access to the output persist shard.
500    ///  * `as_of`: The first time for which the sink may produce output.
501    ///  * `read_only_tx`: A receiver that reports the sink is in read-only mode.
502    ///  * `desired`: The ok/err streams that should be sinked to persist.
503    pub fn render<'s>(
504        sink_id: GlobalId,
505        persist_api: PersistApi,
506        as_of: Antichain<Timestamp>,
507        mut read_only_rx: watch::Receiver<bool>,
508        desired: DesiredStreams<'s>,
509    ) -> (
510        DesiredStreams<'s>,
511        DescsStream<'s>,
512        SharedSinkFrontier,
513        PressOnDropButton,
514    ) {
515        let scope = desired.ok.scope();
516        let worker_id = scope.index();
517        let worker_count = scope.peers();
518
519        // Determine the active worker for the mint operator.
520        let active_worker_id = usize::cast_from(sink_id.hashed()) % scope.peers();
521
522        let sink_frontier = Rc::new(RefCell::new(Antichain::from_elem(Timestamp::MIN)));
523        let shared_frontier = Rc::clone(&sink_frontier);
524
525        let name = operator_name(sink_id, "mint");
526        let mut op = OperatorBuilder::new(name, scope);
527
528        let (ok_output, ok_stream) = op.new_output::<CapacityContainerBuilder<_>>();
529        let (err_output, err_stream) = op.new_output::<CapacityContainerBuilder<_>>();
530        let desired_outputs = OkErr::new(ok_output, err_output);
531        let desired_output_streams = OkErr::new(ok_stream, err_stream);
532
533        let (desc_output, desc_output_stream) = op.new_output::<CapacityContainerBuilder<_>>();
534
535        let mut desired_inputs = OkErr {
536            ok: op.new_input_for(desired.ok, Pipeline, &desired_outputs.ok),
537            err: op.new_input_for(desired.err, Pipeline, &desired_outputs.err),
538        };
539
540        let button = op.build(move |capabilities| async move {
541            // Passing through the `desired` streams only requires data capabilities, so we can
542            // immediately drop their initial capabilities here.
543            let [_, _, desc_cap]: [_; 3] =
544                capabilities.try_into().expect("one capability per output");
545
546            // Non-active workers just pass the `desired` and `persist` data through.
547            if worker_id != active_worker_id {
548                drop(desc_cap);
549                shared_frontier.borrow_mut().clear();
550
551                loop {
552                    tokio::select! {
553                        Some(event) = desired_inputs.ok.next() => {
554                            if let Event::Data(cap, mut data) = event {
555                                desired_outputs.ok.give_container(&cap, &mut data);
556                            }
557                        }
558                        Some(event) = desired_inputs.err.next() => {
559                            if let Event::Data(cap, mut data) = event {
560                                desired_outputs.err.give_container(&cap, &mut data);
561                            }
562                        }
563                        // All inputs are exhausted, so we can shut down.
564                        else => return,
565                    }
566                }
567            }
568
569            let mut cap_set = CapabilitySet::from_elem(desc_cap);
570
571            let read_only = *read_only_rx.borrow_and_update();
572            let mut state = State::new(sink_id, worker_count, as_of, read_only);
573
574            // Create a stream that reports advancements of the target shard's frontier and updates
575            // the shared sink frontier.
576            //
577            // We collect the persist frontier from a write handle directly, rather than inspecting
578            // the `persist` stream, because the latter has two annoying glitches:
579            //  (a) It starts at the shard's read frontier, not its write frontier.
580            //  (b) It can lag behind if there are spikes in ingested data.
581            //
582            // The decoupling from the `persist` stream is load-bearing: that stream can fall
583            // arbitrarily behind the shard upper during snapshot replay or write spikes. Using it
584            // would delay both (1) the controller-visible sink frontier (`shared_frontier`),
585            // which previously caused a CrossJoin feature-bench regression where the controller
586            // held a finished MV dataflow open waiting for the empty frontier, and (2) the
587            // `state.persist_frontier` that gates batch-description minting, stalling mint until
588            // the read-back stream catches up. The goal isn't tick-granular descriptions per
589            // se — it's avoiding the stream-induced stall. See 5eab5ff896 for the original
590            // regression and rationale.
591            let mut persist_frontiers = pin!(async_stream::stream! {
592                let mut writer = persist_api.open_writer().await;
593                let mut frontier = Antichain::from_elem(Timestamp::MIN);
594                while !frontier.is_empty() {
595                    writer.wait_for_upper_past(&frontier).await;
596                    frontier = writer.upper().clone();
597                    shared_frontier.borrow_mut().clone_from(&frontier);
598                    yield frontier.clone();
599                }
600            });
601
602            loop {
603                // Read from the inputs, pass through all data to the respective outputs, and keep
604                // track of the input frontiers. When a frontier advances we might have to mint a
605                // new batch description.
606                let maybe_desc = tokio::select! {
607                    Some(event) = desired_inputs.ok.next() => {
608                        match event {
609                            Event::Data(cap, mut data) => {
610                                desired_outputs.ok.give_container(&cap, &mut data);
611                                None
612                            }
613                            Event::Progress(frontier) => {
614                                state.advance_desired_ok_frontier(frontier);
615                                state.maybe_mint_batch_description()
616                            }
617                        }
618                    }
619                    Some(event) = desired_inputs.err.next() => {
620                        match event {
621                            Event::Data(cap, mut data) => {
622                                desired_outputs.err.give_container(&cap, &mut data);
623                                None
624                            }
625                            Event::Progress(frontier) => {
626                                state.advance_desired_err_frontier(frontier);
627                                state.maybe_mint_batch_description()
628                            }
629                        }
630                    }
631                    Some(frontier) = persist_frontiers.next() => {
632                        state.advance_persist_frontier(frontier);
633                        state.maybe_mint_batch_description()
634                    }
635                    Ok(()) = read_only_rx.changed(), if read_only => {
636                        state.allow_writes();
637                        state.maybe_mint_batch_description()
638                    }
639                    // All inputs are exhausted, so we can shut down.
640                    else => return,
641                };
642
643                if let Some(desc) = maybe_desc {
644                    let lower_ts = *desc.lower.as_option().expect("not empty");
645                    let cap = cap_set.delayed(&lower_ts);
646                    desc_output.give(&cap, desc);
647
648                    // We only emit strictly increasing `lower`s, so we can let our output frontier
649                    // advance beyond the current `lower`.
650                    cap_set.downgrade([lower_ts.step_forward()]);
651                } else {
652                    // The next emitted `lower` will be at least the `persist` frontier, so we can
653                    // advance our output frontier as far.
654                    let _ = cap_set.try_downgrade(state.persist_frontier.iter());
655                }
656            }
657        });
658
659        (
660            desired_output_streams,
661            desc_output_stream,
662            sink_frontier,
663            button.press_on_drop(),
664        )
665    }
666
667    /// State maintained by the `mint` operator.
668    struct State {
669        sink_id: GlobalId,
670        /// The number of workers in the Timely cluster.
671        worker_count: usize,
672        /// The frontiers of the `desired` inputs.
673        desired_frontiers: OkErr<Antichain<Timestamp>, Antichain<Timestamp>>,
674        /// The frontier of the target persist shard.
675        persist_frontier: Antichain<Timestamp>,
676        /// The append worker for the next batch description, chosen in round-robin fashion.
677        next_append_worker: usize,
678        /// The last `lower` we have emitted in a batch description, if any. Whenever the
679        /// `persist_frontier` moves beyond this frontier, we need to mint a new description.
680        last_lower: Option<Antichain<Timestamp>>,
681        /// Whether we are operating in read-only mode.
682        ///
683        /// In read-only mode, minting of batch descriptions is disabled.
684        read_only: bool,
685    }
686
687    impl State {
688        fn new(
689            sink_id: GlobalId,
690            worker_count: usize,
691            as_of: Antichain<Timestamp>,
692            read_only: bool,
693        ) -> Self {
694            // Initializing `persist_frontier` to the `as_of` ensures that the first minted batch
695            // description will have a `lower` of `as_of` or beyond, and thus that we don't spend
696            // work needlessly writing batches at previous times.
697            let persist_frontier = as_of;
698
699            Self {
700                sink_id,
701                worker_count,
702                desired_frontiers: OkErr::new_frontiers(),
703                persist_frontier,
704                next_append_worker: 0,
705                last_lower: None,
706                read_only,
707            }
708        }
709
710        fn trace<S: AsRef<str>>(&self, message: S) {
711            let message = message.as_ref();
712            trace!(
713                sink_id = %self.sink_id,
714                desired_frontier = ?self.desired_frontiers.frontier().elements(),
715                persist_frontier = ?self.persist_frontier.elements(),
716                last_lower = ?self.last_lower.as_ref().map(|f| f.elements()),
717                message,
718            );
719        }
720
721        fn advance_desired_ok_frontier(&mut self, frontier: Antichain<Timestamp>) {
722            if advance(&mut self.desired_frontiers.ok, frontier.borrow()) {
723                self.trace("advanced `desired` ok frontier");
724            }
725        }
726
727        fn advance_desired_err_frontier(&mut self, frontier: Antichain<Timestamp>) {
728            if advance(&mut self.desired_frontiers.err, frontier.borrow()) {
729                self.trace("advanced `desired` err frontier");
730            }
731        }
732
733        fn advance_persist_frontier(&mut self, frontier: Antichain<Timestamp>) {
734            if advance(&mut self.persist_frontier, frontier.borrow()) {
735                self.trace("advanced `persist` frontier");
736            }
737        }
738
739        fn allow_writes(&mut self) {
740            if self.read_only {
741                self.read_only = false;
742                self.trace("disabled read-only mode");
743            }
744        }
745
746        fn maybe_mint_batch_description(&mut self) -> Option<BatchDescription> {
747            let desired_frontier = self.desired_frontiers.frontier();
748            let persist_frontier = &self.persist_frontier;
749
750            // We only mint new batch descriptions when:
751            //  1. We are _not_ in read-only mode.
752            //  2. The `desired` frontier is ahead of the `persist` frontier.
753            //  3. The `persist` frontier advanced since we last emitted a batch description.
754            let desired_ahead = PartialOrder::less_than(persist_frontier, desired_frontier);
755            let persist_advanced = self.last_lower.as_ref().map_or(true, |lower| {
756                PartialOrder::less_than(lower, persist_frontier)
757            });
758
759            if self.read_only || !desired_ahead || !persist_advanced {
760                return None;
761            }
762
763            let lower = persist_frontier.clone();
764            let upper = desired_frontier.clone();
765            let append_worker = self.next_append_worker;
766            let desc = BatchDescription::new(lower, upper, append_worker);
767
768            self.next_append_worker = (append_worker + 1) % self.worker_count;
769            self.last_lower = Some(desc.lower.clone());
770
771            self.trace(format!("minted batch description: {desc:?}"));
772            Some(desc)
773        }
774    }
775}
776
777/// Implementation of the `write` operator.
778mod write {
779    use super::*;
780
781    /// Render the `write` operator.
782    ///
783    /// The parameters passed in are:
784    ///  * `sink_id`: The `GlobalId` of the sink export.
785    ///  * `persist_api`: An object providing access to the output persist shard.
786    ///  * `as_of`: The first time for which the sink may produce output.
787    ///  * `desired`: The ok/err streams that should be sinked to persist.
788    ///  * `persist`: The ok/err streams read back from the output persist shard.
789    ///  * `descs`: The stream of batch descriptions produced by the `mint` operator.
790    pub fn render<'s>(
791        sink_id: GlobalId,
792        persist_api: PersistApi,
793        as_of: Antichain<Timestamp>,
794        desired: DesiredStreams<'s>,
795        persist: PersistStreams<'s>,
796        descs: DescsStream<'s>,
797        mut read_only_rx: watch::Receiver<bool>,
798        worker_config: Rc<ConfigSet>,
799    ) -> (BatchesStream<'s>, PressOnDropButton) {
800        let scope = desired.ok.scope();
801        let worker_id = scope.index();
802
803        let name = operator_name(sink_id, "write");
804        let mut op = OperatorBuilder::new(name, scope.clone());
805
806        let mut channel_logging = None;
807        let mut correction_logger = None;
808        if let (Some(compute_logger), Some(differential_logger)) = (
809            scope.worker().logger_for("materialize/compute"),
810            scope.worker().logger_for("differential/arrange"),
811        ) {
812            let operator_info = op.operator_info();
813            let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
814            channel_logging = Some(ChannelLogging::new(tx));
815            correction_logger = Some(CorrectionLogger::new(
816                compute_logger,
817                differential_logger.into(),
818                operator_info.global_id,
819                operator_info.address.to_vec(),
820                rx,
821            ));
822        }
823
824        let (batches_output, batches_output_stream) =
825            op.new_output::<CapacityContainerBuilder<_>>();
826
827        // It is important that we exchange the `desired` and `persist` data the same way, so
828        // updates that cancel each other out end up on the same worker.
829        let exchange_ok = |(d, _, _): &(Row, Timestamp, Diff)| d.hashed();
830        let exchange_err = |(d, _, _): &(DataflowErrorSer, Timestamp, Diff)| d.hashed();
831
832        let mut desired_inputs = OkErr::new(
833            op.new_disconnected_input(desired.ok, Exchange::new(exchange_ok)),
834            op.new_disconnected_input(desired.err, Exchange::new(exchange_err)),
835        );
836        let mut persist_inputs = OkErr::new(
837            op.new_disconnected_input(persist.ok, Exchange::new(exchange_ok)),
838            op.new_disconnected_input(persist.err, Exchange::new(exchange_err)),
839        );
840        let mut descs_input = op.new_input_for(descs.broadcast(), Pipeline, &batches_output);
841
842        let button = op.build(move |capabilities| async move {
843            // We will use the data capabilities from the `descs` input to produce output, so no
844            // need to hold onto the initial capabilities.
845            drop(capabilities);
846
847            let writer = persist_api.open_writer().await;
848            let shared_batch_writer = SharedBatchWriter::new(
849                &persist_api,
850                &writer,
851                MV_SINK_SHARED_BATCH_MAX_BYTES.handle(&worker_config),
852            )
853            .await;
854            // The write operator needs the writer only for its schema id.
855            drop(writer);
856            let sink_metrics = persist_api.open_metrics().await;
857            let read_only = *read_only_rx.borrow_and_update();
858            let mut state = State::new(
859                sink_id,
860                worker_id,
861                shared_batch_writer,
862                sink_metrics,
863                channel_logging,
864                as_of,
865                read_only,
866                &worker_config,
867            );
868            let mut correction_logger = correction_logger;
869
870            loop {
871                // Drain correction logging events from the channel.
872                if let Some(logger) = &mut correction_logger {
873                    logger.apply_events();
874                }
875
876                // Read from the inputs, extract `desired` updates as positive contributions to
877                // `correction` and `persist` updates as negative contributions. If either the
878                // `desired` or `persist` frontier advances, or if we receive a new batch description,
879                // we might have to write a new batch.
880                let maybe_batch = tokio::select! {
881                    Some(event) = desired_inputs.ok.next() => {
882                        match event {
883                            Event::Data(_cap, mut data) => {
884                                state.corrections.ok.insert(&mut data);
885                                None
886                            }
887                            Event::Progress(frontier) => {
888                                state.advance_desired_ok_frontier(frontier);
889                                state.maybe_write_batch().await
890                            }
891                        }
892                    }
893                    Some(event) = desired_inputs.err.next() => {
894                        match event {
895                            Event::Data(_cap, mut data) => {
896                                state.corrections.err.insert(&mut data);
897                                None
898                            }
899                            Event::Progress(frontier) => {
900                                state.advance_desired_err_frontier(frontier);
901                                state.maybe_write_batch().await
902                            }
903                        }
904                    }
905                    Some(event) = persist_inputs.ok.next() => {
906                        match event {
907                            Event::Data(_cap, mut data) => {
908                                state.corrections.ok.insert_negated(&mut data);
909                                None
910                            }
911                            Event::Progress(frontier) => {
912                                state.advance_persist_ok_frontier(frontier);
913                                state.maybe_write_batch().await
914                            }
915                        }
916                    }
917                    Some(event) = persist_inputs.err.next() => {
918                        match event {
919                            Event::Data(_cap, mut data) => {
920                                state.corrections.err.insert_negated(&mut data);
921                                None
922                            }
923                            Event::Progress(frontier) => {
924                                state.advance_persist_err_frontier(frontier);
925                                state.maybe_write_batch().await
926                            }
927                        }
928                    }
929                    Some(event) = descs_input.next() => {
930                        match event {
931                            Event::Data(cap, data) => {
932                                for desc in data {
933                                    state.absorb_batch_description(desc, cap.clone());
934                                }
935                                state.maybe_write_batch().await
936                            }
937                            Event::Progress(_frontier) => None,
938                        }
939                    }
940                    // Track read-only mode so the forced consolidation stops re-arming once
941                    // writes are allowed and the batch-write path takes over.
942                    Ok(()) = read_only_rx.changed(), if state.read_only => {
943                        if !*read_only_rx.borrow_and_update() {
944                            state.set_read_only(false);
945                        }
946                        None
947                    }
948                    // All inputs are exhausted, so we can shut down.
949                    else => return,
950                };
951
952                if let Some((desc, batches, cap)) = maybe_batch {
953                    for batch in batches {
954                        batches_output.give(&cap, (desc.clone(), batch));
955                    }
956                }
957            }
958        });
959
960        (batches_output_stream, button.press_on_drop())
961    }
962
963    /// State maintained by the `write` operator.
964    struct State {
965        sink_id: GlobalId,
966        worker_id: usize,
967        /// Coalesces batches across the workers in this process.
968        shared_batch_writer: SharedBatchWriter,
969        /// Contains `desired - persist`, reflecting the updates we would like to commit to
970        /// `persist` in order to "correct" it to track `desired`. This collection is only modified
971        /// by updates received from either the `desired` or `persist` inputs.
972        corrections: OkErr<Correction<Row>, Correction<DataflowErrorSer>>,
973        /// The frontiers of the `desired` inputs.
974        desired_frontiers: OkErr<Antichain<Timestamp>, Antichain<Timestamp>>,
975        /// The frontiers of the `persist` inputs.
976        ///
977        /// Note that this is _not_ the same as the write frontier of the output persist shard! It
978        /// usually is, but during snapshot processing, these frontiers will start at the shard's
979        /// read frontier, so they can be beyond its write frontier. This is important as it means
980        /// we must not discard batch descriptions based on these persist frontiers: A batch
981        /// description might still be valid even if its `lower` is before the persist frontiers we
982        /// observe.
983        persist_frontiers: OkErr<Antichain<Timestamp>, Antichain<Timestamp>>,
984        /// The current valid batch description and associated output capability, if any.
985        batch_description: Option<(BatchDescription, Capability<Timestamp>)>,
986        /// A request to force a consolidation of `corrections` once both `desired_frontiers` and
987        /// `persist_frontiers` become greater than the given frontier.
988        ///
989        /// Normally we force a consolidation whenever we write a batch, but there are periods
990        /// (like read-only mode) when that doesn't happen, and we need to manually force
991        /// consolidation instead. In read-only mode this is re-armed after each firing so it
992        /// keeps sweeping forward as the frontiers advance (see `maybe_force_consolidation`).
993        force_consolidation_after: Option<Antichain<Timestamp>>,
994        /// Whether the sink is in read-only mode. While read-only the `write` operator mints no
995        /// batches, so the batch-write path never sweeps `consolidate_before(upper)` forward; the
996        /// forced consolidation stands in for it and is re-armed as long as this holds.
997        read_only: bool,
998    }
999
1000    impl State {
1001        fn new(
1002            sink_id: GlobalId,
1003            worker_id: usize,
1004            shared_batch_writer: SharedBatchWriter,
1005            metrics: SinkMetrics,
1006            logging: Option<ChannelLogging>,
1007            as_of: Antichain<Timestamp>,
1008            read_only: bool,
1009            worker_config: &ConfigSet,
1010        ) -> Self {
1011            let worker_metrics = metrics.for_worker(worker_id);
1012
1013            // Force a consolidation of `corrections` after the snapshot updates have been fully
1014            // processed, to ensure we get rid of those as quickly as possible.
1015            let force_consolidation_after = Some(as_of.clone());
1016
1017            let mut state = Self {
1018                sink_id,
1019                worker_id,
1020                shared_batch_writer,
1021                corrections: OkErr::new(
1022                    Correction::new(
1023                        metrics.clone(),
1024                        worker_metrics.clone(),
1025                        logging.clone(),
1026                        worker_config,
1027                    ),
1028                    Correction::new(metrics, worker_metrics, logging, worker_config),
1029                ),
1030                desired_frontiers: OkErr::new_frontiers(),
1031                persist_frontiers: OkErr::new_frontiers(),
1032                batch_description: None,
1033                force_consolidation_after,
1034                read_only,
1035            };
1036
1037            // Immediately advance the persist frontier tracking to the `as_of`.
1038            // This is important to ensure the persist sink doesn't get stuck if the output shard's
1039            // initial frontier is less than the `as_of`. The `mint` operator first emits a batch
1040            // description with `lower = as_of`, and the `write` operator only emits a batch when
1041            // its observed persist frontier is >= the batch description's `lower`, which (assuming
1042            // no other writers) would be never if we didn't advance the observed persist frontier
1043            // to the `as_of`.
1044            if MV_SINK_ADVANCE_PERSIST_FRONTIERS.get(worker_config) {
1045                state.advance_persist_ok_frontier(as_of.clone());
1046                state.advance_persist_err_frontier(as_of);
1047            }
1048
1049            state
1050        }
1051
1052        fn trace<S: AsRef<str>>(&self, message: S) {
1053            let message = message.as_ref();
1054            trace!(
1055                sink_id = %self.sink_id,
1056                worker = %self.worker_id,
1057                desired_frontier = ?self.desired_frontiers.frontier().elements(),
1058                persist_frontier = ?self.persist_frontiers.frontier().elements(),
1059                batch_description = ?self.batch_description.as_ref().map(|(d, _)| d),
1060                message,
1061            );
1062        }
1063
1064        fn advance_desired_ok_frontier(&mut self, frontier: Antichain<Timestamp>) {
1065            if advance(&mut self.desired_frontiers.ok, frontier.borrow()) {
1066                self.apply_desired_frontier_advancement();
1067                self.trace("advanced `desired` ok frontier");
1068            }
1069        }
1070
1071        fn advance_desired_err_frontier(&mut self, frontier: Antichain<Timestamp>) {
1072            if advance(&mut self.desired_frontiers.err, frontier.borrow()) {
1073                self.apply_desired_frontier_advancement();
1074                self.trace("advanced `desired` err frontier");
1075            }
1076        }
1077
1078        fn advance_persist_ok_frontier(&mut self, frontier: Antichain<Timestamp>) {
1079            if advance(&mut self.persist_frontiers.ok, frontier.borrow()) {
1080                self.apply_persist_frontier_advancement();
1081                self.trace("advanced `persist` ok frontier");
1082            }
1083        }
1084
1085        fn advance_persist_err_frontier(&mut self, frontier: Antichain<Timestamp>) {
1086            if advance(&mut self.persist_frontiers.err, frontier.borrow()) {
1087                self.apply_persist_frontier_advancement();
1088                self.trace("advanced `persist` err frontier");
1089            }
1090        }
1091
1092        /// Apply the effects of a previous `desired` frontier advancement.
1093        fn apply_desired_frontier_advancement(&mut self) {
1094            self.maybe_force_consolidation();
1095        }
1096
1097        /// Apply the effects of a previous `persist` frontier advancement.
1098        fn apply_persist_frontier_advancement(&mut self) {
1099            let frontier = self.persist_frontiers.frontier();
1100
1101            // We will only emit times at or after the `persist` frontier, so now is a good time to
1102            // advance the times of stashed updates.
1103            self.corrections.ok.advance_since(frontier.clone());
1104            self.corrections.err.advance_since(frontier.clone());
1105
1106            self.maybe_force_consolidation();
1107        }
1108
1109        /// If the current consolidation request has become applicable, apply it.
1110        fn maybe_force_consolidation(&mut self) {
1111            let Some(request) = &self.force_consolidation_after else {
1112                return;
1113            };
1114
1115            let desired_frontier = self.desired_frontiers.frontier();
1116            let persist_frontier = self.persist_frontiers.frontier();
1117            if PartialOrder::less_than(request, desired_frontier)
1118                && PartialOrder::less_than(request, persist_frontier)
1119            {
1120                self.trace("forcing correction consolidation");
1121                self.corrections.ok.consolidate_at_since();
1122                self.corrections.err.consolidate_at_since();
1123
1124                // In read-only mode the sink mints no batches, so the batch-write path never
1125                // sweeps `consolidate_before(upper)` forward and a single forced consolidation
1126                // only covers the band below the `since` at the moment it fires, leaving the
1127                // forward region uncancelled for the whole read-only window (CLU-131). Re-arm at
1128                // the current persist frontier so the forced consolidation keeps sweeping forward
1129                // as `desired`/`persist` advance. Once writes are allowed, the batch-write path
1130                // takes over and we disarm.
1131                self.force_consolidation_after = if self.read_only {
1132                    Some(persist_frontier.to_owned())
1133                } else {
1134                    None
1135                };
1136            }
1137        }
1138
1139        /// Update read-only mode. While read-only the forced consolidation re-arms itself; once
1140        /// writes are allowed the still-armed request fires one final time and then disarms, after
1141        /// which the batch-write path drives consolidation.
1142        fn set_read_only(&mut self, read_only: bool) {
1143            self.read_only = read_only;
1144        }
1145
1146        fn absorb_batch_description(&mut self, desc: BatchDescription, cap: Capability<Timestamp>) {
1147            // The incoming batch description is outdated if we already have a batch description
1148            // with a greater `lower`.
1149            //
1150            // Note that we cannot assume a description is outdated based on the comparison of its
1151            // `lower` with the `persist_frontier`. The persist frontier observed by the `write`
1152            // operator is initialized with the shard's read frontier, so it can be greater than
1153            // the shard's write frontier.
1154            if let Some((prev, _)) = &self.batch_description {
1155                if PartialOrder::less_than(&desc.lower, &prev.lower) {
1156                    self.trace(format!("skipping outdated batch description: {desc:?}"));
1157                    return;
1158                }
1159            }
1160
1161            self.batch_description = Some((desc, cap));
1162            self.trace("set batch description");
1163        }
1164
1165        async fn maybe_write_batch(
1166            &mut self,
1167        ) -> Option<(BatchDescription, Vec<ProtoBatch>, Capability<Timestamp>)> {
1168            let (desc, _cap) = self.batch_description.as_ref()?;
1169
1170            // We can write a new batch if we have seen all `persist` updates before `lower` and
1171            // all `desired` updates up to `upper`.
1172            let persist_complete =
1173                PartialOrder::less_equal(&desc.lower, self.persist_frontiers.frontier());
1174            let desired_complete =
1175                PartialOrder::less_equal(&desc.upper, self.desired_frontiers.frontier());
1176            if !persist_complete || !desired_complete {
1177                return None;
1178            }
1179
1180            let (desc, cap) = self.batch_description.take()?;
1181
1182            let ok_updates = self.corrections.ok.updates_before(&desc.upper);
1183            let err_updates = self.corrections.err.updates_before(&desc.upper);
1184
1185            let oks = ok_updates.map(|(d, t, r)| ((SourceData(Ok(d)), ()), t, r.into_inner()));
1186            let errs = err_updates
1187                .map(|(d, t, r)| ((SourceData(Err(d.deserialize())), ()), t, r.into_inner()));
1188            let mut updates = oks.chain(errs).peekable();
1189
1190            // Don't write empty batches.
1191            if updates.peek().is_none() {
1192                drop(updates);
1193                self.trace("skipping empty batch");
1194                return None;
1195            }
1196
1197            let mut write = self.shared_batch_writer.start(&desc);
1198            for (kv, t, d) in updates {
1199                write.push(&kv, t, d).await;
1200            }
1201            let batches = write.finish().await;
1202            if batches.is_empty() {
1203                self.trace("contributed to a shared batch");
1204                return None;
1205            }
1206
1207            self.trace(format!("wrote {} batches", batches.len()));
1208            Some((desc, batches, cap))
1209        }
1210    }
1211}
1212
1213/// Implementation of the `append` operator.
1214mod append {
1215    use super::*;
1216
1217    /// Render the `append` operator.
1218    ///
1219    /// The parameters passed in are:
1220    ///  * `sink_id`: The `GlobalId` of the sink export.
1221    ///  * `persist_api`: An object providing access to the output persist shard.
1222    ///  * `descs`: The stream of batch descriptions produced by the `mint` operator.
1223    ///  * `batches`: The stream of written batches produced by the `write` operator.
1224    pub fn render<'s>(
1225        sink_id: GlobalId,
1226        persist_api: PersistApi,
1227        descs: DescsStream<'s>,
1228        batches: BatchesStream<'s>,
1229    ) -> PressOnDropButton {
1230        let scope = descs.scope();
1231        let worker_id = scope.index();
1232
1233        let name = operator_name(sink_id, "append");
1234        let mut op = OperatorBuilder::new(name, scope);
1235
1236        // Broadcast batch descriptions to all workers, regardless of whether or not they are
1237        // responsible for the append, to give them a chance to clean up any outdated state they
1238        // might still hold.
1239        let mut descs_input = op.new_disconnected_input(descs.broadcast(), Pipeline);
1240        let mut batches_input = op.new_disconnected_input(
1241            batches,
1242            Exchange::new(move |(desc, _): &(BatchDescription, _)| {
1243                u64::cast_from(desc.append_worker)
1244            }),
1245        );
1246
1247        let button = op.build(move |_capabilities| async move {
1248            let writer = persist_api.open_writer().await;
1249            let mut state = State::new(sink_id, worker_id, writer);
1250
1251            loop {
1252                // Read from the inputs, absorb batch descriptions and batches. If the `batches`
1253                // frontier advances, or if we receive a new batch description, we might have to
1254                // append a new batch.
1255                tokio::select! {
1256                    Some(event) = descs_input.next() => {
1257                        if let Event::Data(_cap, data) = event {
1258                            for desc in data {
1259                                state.absorb_batch_description(desc).await;
1260                                state.maybe_append_batches().await;
1261                            }
1262                        }
1263                    }
1264                    Some(event) = batches_input.next() => {
1265                        match event {
1266                            Event::Data(_cap, data) => {
1267                                // The batch description is only used for routing and we ignore it
1268                                // here since we already get one from `descs_input`.
1269                                for (_desc, batch) in data {
1270                                    state.absorb_batch(batch).await;
1271                                }
1272                            }
1273                            Event::Progress(frontier) => {
1274                                state.advance_batches_frontier(frontier);
1275                                state.maybe_append_batches().await;
1276                            }
1277                        }
1278                    }
1279                    // All inputs are exhausted, so we can shut down.
1280                    else => return,
1281                }
1282            }
1283        });
1284
1285        button.press_on_drop()
1286    }
1287
1288    /// State maintained by the `append` operator.
1289    struct State {
1290        sink_id: GlobalId,
1291        worker_id: usize,
1292        persist_writer: WriteHandle<SourceData, (), Timestamp, StorageDiff>,
1293        /// The current input frontier of `batches`.
1294        batches_frontier: Antichain<Timestamp>,
1295        /// The greatest observed `lower` from both `descs` and `batches`.
1296        lower: Antichain<Timestamp>,
1297        /// The batch description for `lower`, if any.
1298        batch_description: Option<BatchDescription>,
1299        /// Batches received for `lower`.
1300        batches: Vec<Batch<SourceData, (), Timestamp, StorageDiff>>,
1301    }
1302
1303    impl State {
1304        fn new(
1305            sink_id: GlobalId,
1306            worker_id: usize,
1307            persist_writer: WriteHandle<SourceData, (), Timestamp, StorageDiff>,
1308        ) -> Self {
1309            Self {
1310                sink_id,
1311                worker_id,
1312                persist_writer,
1313                batches_frontier: Antichain::from_elem(Timestamp::MIN),
1314                lower: Antichain::from_elem(Timestamp::MIN),
1315                batch_description: None,
1316                batches: Default::default(),
1317            }
1318        }
1319
1320        fn trace<S: AsRef<str>>(&self, message: S) {
1321            let message = message.as_ref();
1322            trace!(
1323                sink_id = %self.sink_id,
1324                worker = %self.worker_id,
1325                batches_frontier = ?self.batches_frontier.elements(),
1326                lower = ?self.lower.elements(),
1327                batch_description = ?self.batch_description,
1328                message,
1329            );
1330        }
1331
1332        fn advance_batches_frontier(&mut self, frontier: Antichain<Timestamp>) {
1333            if advance(&mut self.batches_frontier, frontier.borrow()) {
1334                self.trace("advanced `batches` frontier");
1335            }
1336        }
1337
1338        /// Advance the current `lower`.
1339        ///
1340        /// Discards all currently stashed batches and batch descriptions, assuming that they are
1341        /// now invalid.
1342        async fn advance_lower(&mut self, frontier: Antichain<Timestamp>) {
1343            assert!(PartialOrder::less_than(&self.lower, &frontier));
1344
1345            self.lower = frontier;
1346            self.batch_description = None;
1347
1348            // Remove stashed batches, cleaning up those we didn't append.
1349            for batch in self.batches.drain(..) {
1350                batch.delete().await;
1351            }
1352
1353            self.trace("advanced `lower`");
1354        }
1355
1356        /// Absorb the given batch description into the state, provided it is not outdated.
1357        async fn absorb_batch_description(&mut self, desc: BatchDescription) {
1358            if PartialOrder::less_than(&self.lower, &desc.lower) {
1359                self.advance_lower(desc.lower.clone()).await;
1360            } else if &self.lower != &desc.lower {
1361                self.trace(format!("skipping outdated batch description: {desc:?}"));
1362                return;
1363            }
1364
1365            if desc.append_worker == self.worker_id {
1366                self.batch_description = Some(desc);
1367                self.trace("set batch description");
1368            }
1369        }
1370
1371        /// Absorb the given batch into the state, provided it is not outdated.
1372        async fn absorb_batch(&mut self, batch: ProtoBatch) {
1373            let batch = self.persist_writer.batch_from_transmittable_batch(batch);
1374            if PartialOrder::less_than(&self.lower, batch.lower()) {
1375                self.advance_lower(batch.lower().clone()).await;
1376            } else if &self.lower != batch.lower() {
1377                self.trace(format!(
1378                    "skipping outdated batch: ({:?}, {:?})",
1379                    batch.lower().elements(),
1380                    batch.upper().elements(),
1381                ));
1382
1383                // Ensure the batch's data gets properly cleaned up before dropping it.
1384                batch.delete().await;
1385                return;
1386            }
1387
1388            self.batches.push(batch);
1389            self.trace("absorbed a batch");
1390        }
1391
1392        async fn maybe_append_batches(&mut self) {
1393            let batches_complete = PartialOrder::less_than(&self.lower, &self.batches_frontier);
1394            if !batches_complete {
1395                return;
1396            }
1397
1398            let Some(desc) = self.batch_description.take() else {
1399                return;
1400            };
1401
1402            let new_lower = match self.append_batches(desc).await {
1403                Ok(shard_upper) => {
1404                    self.trace("appended a batch");
1405                    shard_upper
1406                }
1407                Err(shard_upper) => {
1408                    // Failing the append is expected in the presence of concurrent replicas. There
1409                    // is nothing special to do here: The self-correcting feedback mechanism
1410                    // ensures that we observe the concurrent changes, compute their consequences,
1411                    // and append them at a future time.
1412                    self.trace(format!(
1413                        "append failed due to `lower` mismatch: {:?}",
1414                        shard_upper.elements(),
1415                    ));
1416                    shard_upper
1417                }
1418            };
1419
1420            self.advance_lower(new_lower).await;
1421        }
1422
1423        /// Append the current `batches` to the output shard.
1424        ///
1425        /// Returns whether the append was successful or not, and the current shard upper in either
1426        /// case.
1427        ///
1428        /// This method advances the shard upper to the batch `lower` if necessary. This is the
1429        /// mechanism that brings the shard upper to the sink as-of when appending the initial
1430        /// batch.
1431        ///
1432        /// An alternative mechanism for bringing the shard upper to the sink as-of would be making
1433        /// a single append at operator startup. The reason we are doing it here instead is that it
1434        /// simplifies the implementation of read-only mode. In read-only mode we have to defer any
1435        /// persist writes, including the initial upper bump. Having only a single place that
1436        /// performs writes makes it easy to ensure we are doing that correctly.
1437        async fn append_batches(
1438            &mut self,
1439            desc: BatchDescription,
1440        ) -> Result<Antichain<Timestamp>, Antichain<Timestamp>> {
1441            let (lower, upper) = (desc.lower, desc.upper);
1442            let mut to_append: Vec<_> = self.batches.iter_mut().collect();
1443
1444            loop {
1445                let result = self
1446                    .persist_writer
1447                    .compare_and_append_batch(&mut to_append, lower.clone(), upper.clone(), true)
1448                    .await
1449                    .expect("valid usage");
1450
1451                match result {
1452                    Ok(()) => return Ok(upper),
1453                    Err(mismatch) if PartialOrder::less_than(&mismatch.current, &lower) => {
1454                        advance_shard_upper(&mut self.persist_writer, lower.clone()).await;
1455
1456                        // At this point the shard's since and upper are likely the same, a state
1457                        // that is likely to hit edge-cases in logic reasoning about frontiers.
1458                        fail::fail_point!("mv_advanced_upper");
1459                    }
1460                    Err(mismatch) => return Err(mismatch.current),
1461                }
1462            }
1463        }
1464    }
1465
1466    /// Advance the frontier of the given writer's shard to at least the given `upper`.
1467    async fn advance_shard_upper(
1468        persist_writer: &mut WriteHandle<SourceData, (), Timestamp, StorageDiff>,
1469        upper: Antichain<Timestamp>,
1470    ) {
1471        let empty_updates: &[((SourceData, ()), Timestamp, StorageDiff)] = &[];
1472        let lower = Antichain::from_elem(Timestamp::MIN);
1473        persist_writer
1474            .append(empty_updates, lower, upper)
1475            .await
1476            .expect("valid usage")
1477            .expect("should always succeed");
1478    }
1479}