Skip to main content

mz_storage/render/
persist_sink.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//! Render an operator that persists a source collection.
11//!
12//! ## Implementation
13//!
14//! This module defines the `persist_sink` operator, that writes
15//! a collection produced by source rendering into a persist shard.
16//!
17//! It attempts to use all workers to write data to persist, and uses
18//! single-instance workers to coordinate work. The below diagram
19//! is an overview how it it shaped. There is more information
20//! in the doc comments of the top-level functions of this module.
21//!
22//!```text
23//!
24//!                                       ,------------.
25//!                                       | source     |
26//!                                       | collection |
27//!                                       +---+--------+
28//!                                       /   |
29//!                                      /    |
30//!                                     /     |
31//!                                    /      |
32//!                                   /       |
33//!                                  /        |
34//!                                 /         |
35//!                                /          |
36//!                               /     ,-+-----------------------.
37//!                              /      | mint_batch_descriptions |
38//!                             /       | one arbitrary worker    |
39//!                            |        +-,--,--------+----+------+
40//!                           ,----------´.-´         |     \
41//!                       _.-´ |       .-´            |      \
42//!                   _.-´     |    .-´               |       \
43//!                .-´  .------+----|-------+---------|--------\-----.
44//!               /    /            |       |         |         \     \
45//!        ,--------------.   ,-----------------.     |     ,-----------------.
46//!        | write_batches|   |  write_batches  |     |     |  write_batches  |
47//!        | worker 0     |   | worker 1        |     |     | worker N        |
48//!        +-----+--------+   +-+---------------+     |     +--+--------------+
49//!               \              \                    |        /
50//!                `-.            `,                  |       /
51//!                   `-._          `-.               |      /
52//!                       `-._         `-.            |     /
53//!                           `---------. `-.         |    /
54//!                                     +`---`---+-------------,
55//!                                     | append_batches       |
56//!                                     | one arbitrary worker |
57//!                                     +------+---------------+
58//!```
59//!
60//! `mint_batch_descriptions` also takes the ingestion's remap upper as a disconnected input, and
61//! broadcasts a committed ceiling to the writers, which `append_batches` does not see. Both edges
62//! are left out of the diagram to avoid clutter.
63//!
64//! ## Similarities with `mz_compute::sink::persist_sink`
65//!
66//! This module has many similarities with the compute version of
67//! the same concept, and in fact, is entirely derived from it.
68//!
69//! Compute requires that its `persist_sink` is _self-correcting_;
70//! that is, it corrects what the collection in persist
71//! accumulates to if the collection has values changed at
72//! previous timestamps. It does this by continually comparing
73//! the input stream with the collection as read back from persist.
74//!
75//! Source collections, while definite, cannot be reliably by
76//! re-produced once written down, which means compute's
77//! `persist_sink`'s self-correction mechanism would need to be
78//! skipped on operator startup, and would cause unnecessary read
79//! load on persist.
80//!
81//! Additionally, persisting sources requires we use bounded
82//! amounts of memory, even if a single timestamp represents
83//! a huge amount of data. This is not (currently) possible
84//! to guarantee while also performing self-correction.
85//!
86//! Because of this, we have ripped out the self-correction
87//! mechanism, and aggressively simplified the sub-operators.
88//! Some, particularly `append_batches` could be merged with
89//! the compute version, but that requires some amount of
90//! onerous refactoring that we have chosen to skip for now.
91//!
92// TODO(guswynn): merge at least the `append_batches` operator`
93
94use std::cmp::Ordering;
95use std::collections::{BTreeMap, VecDeque};
96use std::fmt::Debug;
97use std::ops::AddAssign;
98use std::rc::Rc;
99use std::sync::Arc;
100use std::time::Duration;
101
102use differential_dataflow::difference::Monoid;
103use differential_dataflow::lattice::Lattice;
104use differential_dataflow::{AsCollection, Hashable, VecCollection};
105use futures::{StreamExt, future};
106use itertools::Itertools;
107use mz_ore::cast::CastFrom;
108use mz_ore::collections::HashMap;
109use mz_persist_client::Diagnostics;
110use mz_persist_client::batch::{Batch, BatchBuilder, ProtoBatch};
111use mz_persist_client::cache::PersistClientCache;
112use mz_persist_client::error::UpperMismatch;
113use mz_persist_types::codec_impls::UnitSchema;
114use mz_persist_types::{Codec, Codec64};
115use mz_repr::{Diff, GlobalId, Row};
116use mz_storage_types::controller::CollectionMetadata;
117use mz_storage_types::errors::DataflowError;
118use mz_storage_types::sources::SourceData;
119use mz_storage_types::{StorageDiff, dyncfgs};
120use mz_timely_util::builder_async::{
121    Event, OperatorBuilder as AsyncOperatorBuilder, PressOnDropButton, sole_capability,
122};
123use serde::{Deserialize, Serialize};
124use timely::PartialOrder;
125use timely::container::CapacityContainerBuilder;
126use timely::dataflow::channels::pact::{Exchange, Pipeline};
127use timely::dataflow::operators::vec::Broadcast;
128use timely::dataflow::operators::{Capability, CapabilitySet, Inspect};
129use timely::dataflow::{Scope, Stream, StreamVec};
130use timely::progress::{Antichain, Timestamp};
131use tokio::sync::Semaphore;
132use tracing::{info, trace};
133
134use crate::metrics::source::SourcePersistSinkMetrics;
135use crate::statistics::SourceStatistics;
136use crate::storage_state::StorageState;
137
138/// Metrics about batches.
139#[derive(Clone, Debug, Default, Deserialize, Serialize)]
140struct BatchMetrics {
141    inserts: u64,
142    retractions: u64,
143    error_inserts: u64,
144    error_retractions: u64,
145}
146
147impl AddAssign<&BatchMetrics> for BatchMetrics {
148    fn add_assign(&mut self, rhs: &BatchMetrics) {
149        let BatchMetrics {
150            inserts: self_inserts,
151            retractions: self_retractions,
152            error_inserts: self_error_inserts,
153            error_retractions: self_error_retractions,
154        } = self;
155        let BatchMetrics {
156            inserts: rhs_inserts,
157            retractions: rhs_retractions,
158            error_inserts: rhs_error_inserts,
159            error_retractions: rhs_error_retractions,
160        } = rhs;
161        *self_inserts += rhs_inserts;
162        *self_retractions += rhs_retractions;
163        *self_error_inserts += rhs_error_inserts;
164        *self_error_retractions += rhs_error_retractions;
165    }
166}
167
168/// Manages batches and metrics.
169struct BatchBuilderAndMetadata<K, V, T, D>
170where
171    K: Codec,
172    V: Codec,
173    T: Timestamp + Lattice + Codec64,
174{
175    builder: BatchBuilder<K, V, T, D>,
176    /// Largest update timestamp staged so far, `None` while empty.
177    ///
178    /// `append_batches` needs this to decide, after an `UpperMismatch`, whether a batch lies
179    /// entirely below a raised append lower. A batch completely below the append lower is
180    /// deleted. A batch whose data straddles an append lower has its bounds adjusted instead.
181    data_max_ts: Option<T>,
182    metrics: BatchMetrics,
183}
184
185impl<K, V, T, D> BatchBuilderAndMetadata<K, V, T, D>
186where
187    K: Codec + Debug,
188    V: Codec + Debug,
189    T: Timestamp + Lattice + Codec64,
190    D: Monoid + Codec64,
191{
192    /// Creates a new batch. Updates at any timestamp at or beyond the builder's lower may be
193    /// added, in any order.
194    fn new(builder: BatchBuilder<K, V, T, D>) -> Self {
195        BatchBuilderAndMetadata {
196            builder,
197            data_max_ts: None,
198            metrics: Default::default(),
199        }
200    }
201
202    /// Adds an update to the batch.
203    async fn add(&mut self, k: &K, v: &V, t: &T, d: &D) {
204        self.data_max_ts = Some(match self.data_max_ts.take() {
205            Some(max) => max.join(t),
206            None => t.clone(),
207        });
208
209        self.builder.add(k, v, t, d).await.expect("invalid usage");
210    }
211
212    /// Finishes the batch, registering it under `lower` and `upper`.
213    ///
214    /// Panics if no update was ever added, since an empty batch has no largest timestamp. Callers
215    /// open a builder on the first update rather than up front, so reaching this is a bug.
216    async fn finish(self, lower: Antichain<T>, upper: Antichain<T>) -> HollowBatchAndMetadata<T> {
217        let data_max_ts = self.data_max_ts.expect("finishing an empty builder");
218        // `BatchBuilder::finish` rejects an update at or beyond `upper`, so a builder that was
219        // handed updates outside the description it is being finished under fails here rather
220        // than producing a batch whose parts reach past their registered bounds.
221        let batch = self
222            .builder
223            .finish(upper.clone())
224            .await
225            .expect("invalid usage");
226        HollowBatchAndMetadata {
227            lower,
228            upper,
229            data_max_ts,
230            batch: batch.into_transmittable_batch(),
231            metrics: self.metrics,
232        }
233    }
234}
235
236/// A batch or data + metrics moved from `write_batches` to `append_batches`.
237#[derive(Clone, Debug, Deserialize, Serialize)]
238#[serde(bound(
239    serialize = "T: Timestamp + Codec64",
240    deserialize = "T: Timestamp + Codec64"
241))]
242struct HollowBatchAndMetadata<T> {
243    lower: Antichain<T>,
244    upper: Antichain<T>,
245    data_max_ts: T,
246    batch: ProtoBatch,
247    metrics: BatchMetrics,
248}
249
250/// Holds finished batches for `append_batches`.
251#[derive(Debug, Default)]
252struct BatchSet {
253    finished: Vec<FinishedBatch>,
254    batch_metrics: BatchMetrics,
255}
256
257#[derive(Debug)]
258struct FinishedBatch {
259    batch: Batch<SourceData, (), mz_repr::Timestamp, StorageDiff>,
260    data_max_ts: mz_repr::Timestamp,
261}
262
263/// The batch builder the source sink writes with.
264type SourceBatchBuilder = BatchBuilderAndMetadata<SourceData, (), mz_repr::Timestamp, StorageDiff>;
265
266/// How far past the remap upper the minter commits its ceiling while a snapshot pins the
267/// frontier, in milliseconds, the unit of `mz_repr::Timestamp`. `None` turns committing ahead off.
268///
269/// `timestamp_interval` is the floor to avoid having data outrun the ceiling. Remap emits a new
270/// binding and downgrades its capability together. Rows emitted by reclock under the new binding
271/// are racing the new ceiling to the batch writers. At one timestamp interval the previous ceiling
272/// already covers the current timestamp's rows. Rows still outrun it when consecutive probes are
273/// further apart than the lookahead, which costs a batch rather than correctness.
274fn description_lookahead(lookahead: Duration, timestamp_interval: Duration) -> Option<u64> {
275    // Saturates rather than panics, since the lookahead is operator-supplied. A lookahead that
276    // large overflows the ceiling in `next_mint`, which then commits nothing.
277    let to_millis = |d: Duration| u64::try_from(d.as_millis()).unwrap_or(u64::MAX);
278    // The zero test is on the `Duration`, matching the sources that read the same dyncfg to
279    // enable concurrent replication. Testing after truncating to millis would turn a sub-ms
280    // lookahead into concurrent replication without grouping. The floor of one millisecond keeps
281    // a nonzero setting nonzero when the timestamp interval is itself sub-ms.
282    (!lookahead.is_zero()).then(|| {
283        to_millis(lookahead)
284            .max(to_millis(timestamp_interval))
285            .max(1)
286    })
287}
288
289/// Adds one update to `builder`, keeping the batch metrics in step.
290async fn stage_update(
291    builder: &mut SourceBatchBuilder,
292    row: Result<Row, DataflowError>,
293    ts: mz_repr::Timestamp,
294    diff: Diff,
295) {
296    let is_value = row.is_ok();
297
298    builder
299        .add(&SourceData(row), &(), &ts, &diff.into_inner())
300        .await;
301
302    // Note that we assume `diff` is either +1 or -1 here, being anything else is a logic bug we
303    // can't handle at the metric layer. We also assume this addition doesn't overflow.
304    match (is_value, diff.is_positive()) {
305        (true, true) => builder.metrics.inserts += diff.unsigned_abs(),
306        (true, false) => builder.metrics.retractions += diff.unsigned_abs(),
307        (false, true) => builder.metrics.error_inserts += diff.unsigned_abs(),
308        (false, false) => builder.metrics.error_retractions += diff.unsigned_abs(),
309    }
310}
311
312/// Continuously writes the `desired_stream` into persist
313/// This is done via a multi-stage operator graph:
314///
315/// 1. `mint_batch_descriptions` emits new batch descriptions whenever the
316///    frontier of `desired_collection` advances. A batch description is
317///    a pair of `(lower, upper)` that tells write operators
318///    which updates to write and in the end tells the append operator
319///    what frontiers to use when calling `append`/`compare_and_append`.
320///    This is a single-worker operator.
321/// 2. `write_batches` writes the `desired_collection` to persist as
322///    batches and sends those batches along.
323///    This does not yet append the batches to the persist shard, the update are
324///    only uploaded/prepared to be appended to a shard. Also: we only write
325///    updates for batch descriptions that we learned about from
326///    `mint_batch_descriptions`.
327/// 3. `append_batches` takes as input the minted batch descriptions and written
328///    batches. Whenever the frontiers sufficiently advance, we take a batch
329///    description and all the batches that belong to it and append it to the
330///    persist shard.
331///
332/// This operator assumes that the `desired_collection` comes pre-sharded.
333///
334/// Note that `mint_batch_descriptions` inspects the frontier of
335/// `desired_collection`, and passes the data through to `write_batches`.
336/// This is done to avoid a clone of the underlying data so that both
337/// operators can have the collection as input.
338///
339/// `snapshot_time` is the as_of, the time the collection's snapshot lands at, given only for an
340/// export that snapshots in this incarnation. The frontier sits there for the length of the
341/// snapshot, which is what gives the sink something to group. `remap_upper` carries the
342/// ingestion's remap upper as its frontier, which paces the grouping. See
343/// [`mint_batch_descriptions`].
344pub(crate) fn render<'scope>(
345    scope: Scope<'scope, mz_repr::Timestamp>,
346    collection_id: GlobalId,
347    target: CollectionMetadata,
348    desired_collection: VecCollection<'scope, mz_repr::Timestamp, Result<Row, DataflowError>, Diff>,
349    storage_state: &StorageState,
350    metrics: SourcePersistSinkMetrics,
351    busy_signal: Arc<Semaphore>,
352    snapshot_time: Option<Antichain<mz_repr::Timestamp>>,
353    timestamp_interval: Duration,
354    remap_upper: StreamVec<'scope, mz_repr::Timestamp, ()>,
355) -> (
356    StreamVec<'scope, mz_repr::Timestamp, ()>,
357    StreamVec<'scope, mz_repr::Timestamp, Rc<anyhow::Error>>,
358    Vec<PressOnDropButton>,
359) {
360    let persist_clients = Arc::clone(&storage_state.persist_clients);
361
362    let operator_name = format!("persist_sink({})", collection_id);
363
364    let config_set = storage_state.storage_configuration.config_set();
365    let lookahead = description_lookahead(
366        dyncfgs::STORAGE_PERSIST_SINK_DESCRIPTION_LOOKAHEAD.get(config_set),
367        timestamp_interval,
368    );
369    // Once per dataflow build, so a hydration's logs show whether this sink groups a snapshot's
370    // updates into one batch or writes one batch per timestamp.
371    info!(
372        %collection_id,
373        lookahead_ms = lookahead.unwrap_or(0),
374        "persist_sink description lookahead"
375    );
376
377    let (batch_descriptions, commitments, passthrough_desired_stream, mint_token) =
378        mint_batch_descriptions(
379            scope,
380            collection_id,
381            &operator_name,
382            &target,
383            desired_collection,
384            remap_upper,
385            Arc::clone(&persist_clients),
386            lookahead,
387            snapshot_time,
388        );
389
390    let source_statistics = storage_state
391        .aggregated_statistics
392        .get_source(&collection_id)
393        .expect("statistics initialized")
394        .clone();
395
396    let (written_batches, write_token) = write_batches(
397        scope,
398        collection_id.clone(),
399        &operator_name,
400        &target,
401        batch_descriptions.clone(),
402        commitments,
403        passthrough_desired_stream.as_collection(),
404        Arc::clone(&persist_clients),
405        source_statistics,
406        Arc::clone(&busy_signal),
407    );
408
409    let (upper_stream, append_errors, append_token) = append_batches(
410        scope,
411        collection_id.clone(),
412        operator_name,
413        &target,
414        batch_descriptions,
415        written_batches,
416        persist_clients,
417        storage_state,
418        metrics,
419        Arc::clone(&busy_signal),
420    );
421
422    (
423        upper_stream,
424        append_errors,
425        vec![mint_token, write_token, append_token],
426    )
427}
428
429/// Whenever the frontier advances, this mints a new batch description (lower
430/// and upper) that writers should use for writing the next set of batches to
431/// persist.
432///
433/// With a `lookahead`, and while the frontier sits at `snapshot_time`, it also commits to a
434/// ceiling that far past the data and broadcasts it on the second output. A
435/// ceiling is not a description: it gives the writers a bound to group updates under before the
436/// frontier certifies anything, and it binds this operator, which mints nothing below an
437/// outstanding ceiling. So the whole snapshot and the catch-up behind it become one description,
438/// emitted when the frontier reaches the ceiling. See [`next_mint`].
439///
440/// Only one of the workers does this, meaning there will only be one
441/// description in the stream, even in case of multiple timely workers. Use
442/// `broadcast()` to, ahem, broadcast, the one description to all downstream
443/// write operators/workers.
444///
445/// `remap_upper` carries the ingestion's remap upper as its frontier, which paces the
446/// commitments. Reclocking stamps every update below that upper before the update exists, so a
447/// ceiling ahead of it leads every row that can still arrive, however long this export's own data
448/// has been quiet.
449fn mint_batch_descriptions<'scope>(
450    scope: Scope<'scope, mz_repr::Timestamp>,
451    collection_id: GlobalId,
452    operator_name: &str,
453    target: &CollectionMetadata,
454    desired_collection: VecCollection<'scope, mz_repr::Timestamp, Result<Row, DataflowError>, Diff>,
455    remap_upper: StreamVec<'scope, mz_repr::Timestamp, ()>,
456    persist_clients: Arc<PersistClientCache>,
457    lookahead: Option<u64>,
458    snapshot_time: Option<Antichain<mz_repr::Timestamp>>,
459) -> (
460    StreamVec<
461        'scope,
462        mz_repr::Timestamp,
463        (Antichain<mz_repr::Timestamp>, Antichain<mz_repr::Timestamp>),
464    >,
465    StreamVec<'scope, mz_repr::Timestamp, Commitment>,
466    StreamVec<'scope, mz_repr::Timestamp, (Result<Row, DataflowError>, mz_repr::Timestamp, Diff)>,
467    PressOnDropButton,
468) {
469    let persist_location = target.persist_location.clone();
470    let shard_id = target.data_shard;
471    let target_relation_desc = target.relation_desc.clone();
472
473    // Only one worker is responsible for determining batch descriptions. All
474    // workers must write batches with the same description, to ensure that they
475    // can be combined into one batch that gets appended to Consensus state.
476    let hashed_id = collection_id.hashed();
477    let active_worker = usize::cast_from(hashed_id) % scope.peers() == scope.index();
478
479    // Only the "active" operator will mint batches. All other workers have an
480    // empty frontier. It's necessary to insert all of these into
481    // `compute_state.sink_write_frontier` below so we properly clear out
482    // default frontiers of non-active workers.
483
484    let mut mint_op = AsyncOperatorBuilder::new(
485        format!("{} mint_batch_descriptions", operator_name),
486        scope.clone(),
487    );
488
489    let (output, output_stream) = mint_op.new_output::<CapacityContainerBuilder<Vec<_>>>();
490    let (ceiling_output, ceiling_output_stream) =
491        mint_op.new_output::<CapacityContainerBuilder<Vec<_>>>();
492    let (data_output, data_output_stream) =
493        mint_op.new_output::<CapacityContainerBuilder<Vec<_>>>();
494
495    // The description, ceiling and data-passthrough outputs are all driven by this input, so
496    // they use a standard input connection.
497    let mut desired_input = mint_op.new_input_for_many(
498        desired_collection.inner,
499        Pipeline,
500        [&output, &ceiling_output, &data_output],
501    );
502
503    // The remap upper only influences the bounds of descriptions, it doesn't drive the outputs, so
504    // this input is disconnected. The stream is already broadcast, so every worker holds its
505    // frontier and the active one reads it in place.
506    let mut remap_input = mint_op.new_disconnected_input(remap_upper, Pipeline);
507
508    let shutdown_button = mint_op.build(move |capabilities| async move {
509        // Non-active workers should just pass the data through.
510        if !active_worker {
511            // The description and ceiling outputs are entirely driven by the active worker, so we
512            // drop their capabilities here. The data-passthrough output just uses the data
513            // capabilities.
514            drop(capabilities);
515            while let Some(event) = desired_input.next().await {
516                match event {
517                    Event::Data([_output_cap, _ceiling_cap, data_output_cap], mut data) => {
518                        data_output.give_container(&data_output_cap, &mut data);
519                    }
520                    Event::Progress(_) => {}
521                }
522            }
523            return;
524        }
525        // The data-passthrough output uses the data capabilities, so we drop its capability here.
526        let [desc_cap, ceiling_cap, _]: [_; 3] =
527            capabilities.try_into().expect("one capability per output");
528        let mut cap_set = CapabilitySet::from_elem(desc_cap);
529        let mut ceiling_cap_set = CapabilitySet::from_elem(ceiling_cap);
530
531        // Initialize this operators's `upper` to the `upper` of the persist shard we are writing
532        // to. Data from the source not beyond this time will be dropped, as it has already
533        // been persisted.
534        // In the future, sources will avoid passing through data not beyond this upper
535        let mut current_upper = {
536            // TODO(aljoscha): We need to figure out what to do with error
537            // results from these calls.
538            let persist_client = persist_clients
539                .open(persist_location)
540                .await
541                .expect("could not open persist client");
542
543            let mut write = persist_client
544                .open_writer::<SourceData, (), mz_repr::Timestamp, StorageDiff>(
545                    shard_id,
546                    Arc::new(target_relation_desc),
547                    Arc::new(UnitSchema),
548                    Diagnostics {
549                        shard_name: collection_id.to_string(),
550                        handle_purpose: format!(
551                            "storage::persist_sink::mint_batch_descriptions {}",
552                            collection_id
553                        ),
554                    },
555                )
556                .await
557                .expect("could not open persist shard");
558
559            // TODO: this sink currently cannot tolerate a stale upper... which is bad because the
560            // upper can become stale as soon as it is read. (For example, if another concurrent
561            // instance of the sink has updated it.) Fetching a recent upper helps to mitigate this,
562            // but ideally we would just skip ahead if we discover that our upper is stale.
563            let upper = write.fetch_recent_upper().await.clone();
564            // explicitly expire the once-used write handle.
565            write.expire().await;
566            upper
567        };
568
569        // The current input frontier.
570        let mut desired_frontier = Antichain::from_elem(mz_repr::Timestamp::minimum());
571
572        // The ingestion's remap upper, which drives the ceiling.
573        let mut remap_upper = Antichain::from_elem(mz_repr::Timestamp::minimum());
574
575        // The outstanding ceiling, if any. While one is held nothing below it is minted, which is
576        // what makes it binding.
577        let mut committed: Option<mz_repr::Timestamp> = None;
578
579        // The wait sits at the bottom of the loop, so a mint pass runs once before the first
580        // event as well as after every event. A fresh export's snapshot lands at the minimum and
581        // its frontier never moves during it, so no event announces the pin. Its first ceiling
582        // has to go out ahead of the snapshot's rows rather than alongside them.
583        loop {
584            // `desired_frontier` starts at the minimum. A fresh export's as_of is the minimum too,
585            // so it reads as snapshotting from the first activation. A later as_of waits for the
586            // progress statement that moves the frontier there. A row can arrive ahead of that
587            // statement, and committing on it would anchor the ceiling at the shard upper instead
588            // of the snapshot's time.
589            let snapshot_in_progress = snapshot_time
590                .as_ref()
591                .is_some_and(|time| *time == desired_frontier);
592
593            while let Some(mint) = next_mint(
594                &current_upper,
595                &desired_frontier,
596                &remap_upper,
597                committed,
598                lookahead.filter(|_| snapshot_in_progress),
599            ) {
600                let lower = current_upper
601                    .as_option()
602                    .copied()
603                    .expect("a non-empty current upper, or nothing is minted");
604
605                let upper = match mint {
606                    Mint::Ceiling(ceiling) => {
607                        let cap = ceiling_cap_set
608                            .try_delayed(&lower)
609                            .expect("ceiling capability holds the current upper");
610                        trace!(
611                            "persist_sink {collection_id}/{shard_id}: \
612                                committing ceiling: {:?}",
613                            ceiling
614                        );
615                        ceiling_output.give(&cap, Commitment { lower, ceiling });
616                        committed = Some(ceiling);
617                        continue;
618                    }
619                    Mint::Description(upper) => upper,
620                };
621
622                let batch_description = (current_upper.to_owned(), upper.to_owned());
623
624                let cap = cap_set
625                    .try_delayed(&lower)
626                    .ok_or_else(|| {
627                        format!(
628                            "minter cannot delay {:?} to {:?}. \
629                                Likely because we already emitted a \
630                                batch description and delayed.",
631                            cap_set, lower
632                        )
633                    })
634                    .unwrap();
635
636                trace!(
637                    "persist_sink {collection_id}/{shard_id}: \
638                        new batch_description: {:?}",
639                    batch_description
640                );
641
642                output.give(&cap, batch_description);
643
644                // We downgrade our capability to the batch
645                // description upper, as there will never be
646                // any overlapping descriptions.
647                trace!(
648                    "persist_sink {collection_id}/{shard_id}: \
649                        downgrading to {:?}",
650                    upper
651                );
652                cap_set.downgrade(upper.iter());
653                ceiling_cap_set.downgrade(upper.iter());
654
655                // The description that retires a ceiling covers everything the writers grouped
656                // under it, so the next one starts fresh.
657                committed = None;
658                current_upper = upper;
659            }
660
661            tokio::select! {
662                event = desired_input.next() => match event {
663                    Some(Event::Data([_output_cap, _ceiling_cap, data_output_cap], mut data)) => {
664                        // Just passthrough the data.
665                        data_output.give_container(&data_output_cap, &mut data);
666                    }
667                    Some(Event::Progress(frontier)) => desired_frontier = frontier,
668                    // Input is exhausted, so we can shut down.
669                    None => return,
670                },
671                // During the snapshot, the frontier is pinned and the remap upper drives the
672                // ceiling.
673                _ = remap_input.ready() => {},
674            }
675
676            // The stream carries no data, only its frontier.
677            while let Some(event) = remap_input.next_sync() {
678                if let Event::Progress(frontier) = event {
679                    remap_upper = frontier;
680                }
681            }
682        }
683    });
684
685    (
686        output_stream,
687        ceiling_output_stream,
688        data_output_stream,
689        shutdown_button.press_on_drop(),
690    )
691}
692
693/// A bound the minter commits ahead of the frontier, broadcast to the writers.
694///
695/// `lower` is the lower of the description that will end at or past `ceiling`, which the writers
696/// need before that description exists, since a builder declares its lower when it opens.
697#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize)]
698struct Commitment {
699    lower: mz_repr::Timestamp,
700    ceiling: mz_repr::Timestamp,
701}
702
703/// What [`mint_batch_descriptions`] does with the frontier and the data it has seen.
704#[derive(Debug, Clone, PartialEq, Eq)]
705enum Mint {
706    /// Emit a description ending here, the unit `append_batches` appends.
707    Description(Antichain<mz_repr::Timestamp>),
708    /// Commit to this ceiling and broadcast it to the writers. Not a description: it gives them a
709    /// bound to group under, and binds the minter to mint nothing below it.
710    Ceiling(mz_repr::Timestamp),
711}
712
713/// The next thing the minter emits, or `None` when there is nothing to do.
714///
715/// A description is derived from the frontier, which certifies that everything below it has
716/// arrived. An outstanding `committed` ceiling suppresses that until the frontier reaches the
717/// ceiling, which is what collapses a whole snapshot and the catch-up behind it into one
718/// description.
719///
720/// * `current_upper`:
721///     The upper of the last minted description, initialized from the shard.
722///     Any data before this has already been processed.
723/// * `desired_frontier`:
724///     Data input frontier.
725///     This frontier advances when the reclocked data stream advances.
726/// * `remap_upper`:
727///     The frontier of the remap shard, which advances on each successful probe.
728/// * `committed` (Optional):
729///     This is the current value of the ceiling. If set, the upper of a minted description will be
730///     at or beyond this.
731/// * `lookahead` (Optional):
732///     If present, `desired_frontier` is pinned to `as_of` (e.g. due to snapshot), so we advance
733///     `committed` as `remap_upper` + `lookahead` to consolidate in-flight data into one shared
734///     batch.
735fn next_mint(
736    current_upper: &Antichain<mz_repr::Timestamp>,
737    desired_frontier: &Antichain<mz_repr::Timestamp>,
738    remap_upper: &Antichain<mz_repr::Timestamp>,
739    committed: Option<mz_repr::Timestamp>,
740    lookahead: Option<u64>,
741) -> Option<Mint> {
742    let frontier_reached = |ts: mz_repr::Timestamp| {
743        PartialOrder::less_equal(&Antichain::from_elem(ts), desired_frontier)
744    };
745    if committed.is_none_or(frontier_reached)
746        && PartialOrder::less_than(current_upper, desired_frontier)
747    {
748        return Some(Mint::Description(desired_frontier.clone()));
749    }
750
751    // A bunch of checks that will return None (meaning there's no ceiling).
752    // If feature is not enabled or not snapshotting...
753    let lookahead = lookahead?;
754
755    // If the data shard has advanced to empty frontier...
756    let lower = *current_upper.as_option()?;
757
758    // If the remap shard has advanced to the empty frontier, or the add overflows because of
759    // a large lookahead value (see `description_lookahead`).
760    let ceiling = remap_upper.as_option()?.checked_add(lookahead)?;
761
762    // lower can be beyond remap_upper + lookahead any time the current_upper jumps ahead of
763    // of the remap_upper the operator is tracking.
764    (lower < ceiling && committed.is_none_or(|c| c < ceiling)).then_some(Mint::Ceiling(ceiling))
765}
766
767/// Writes `desired_collection` to persist, but only for updates
768/// that fall into batch a description that we get via `batch_descriptions`.
769/// This forwards a `HollowBatch` (with additional metadata)
770/// for any batch of updates that was written.
771///
772/// Every update below the ceiling `mint_batch_descriptions` commits goes into one open builder,
773/// whatever its timestamp, so a pinned frontier costs one batch rather than one per timestamp. An
774/// update that outruns the ceiling, or arrives while none is committed, has no bound to be grouped
775/// under and writes a batch of its own timestamp, which is what the sink does for every update when
776/// nothing is ever committed ahead of the frontier.
777///
778/// This operator assumes that the `desired_collection` comes pre-sharded.
779///
780/// This also and updates various metrics.
781fn write_batches<'scope>(
782    scope: Scope<'scope, mz_repr::Timestamp>,
783    collection_id: GlobalId,
784    operator_name: &str,
785    target: &CollectionMetadata,
786    batch_descriptions: Stream<
787        'scope,
788        mz_repr::Timestamp,
789        Vec<(Antichain<mz_repr::Timestamp>, Antichain<mz_repr::Timestamp>)>,
790    >,
791    commitments: StreamVec<'scope, mz_repr::Timestamp, Commitment>,
792    desired_collection: VecCollection<'scope, mz_repr::Timestamp, Result<Row, DataflowError>, Diff>,
793    persist_clients: Arc<PersistClientCache>,
794    source_statistics: SourceStatistics,
795    busy_signal: Arc<Semaphore>,
796) -> (
797    StreamVec<'scope, mz_repr::Timestamp, HollowBatchAndMetadata<mz_repr::Timestamp>>,
798    PressOnDropButton,
799) {
800    let worker_index = scope.index();
801
802    let persist_location = target.persist_location.clone();
803    let shard_id = target.data_shard;
804    let target_relation_desc = target.relation_desc.clone();
805
806    let mut write_op =
807        AsyncOperatorBuilder::new(format!("{} write_batches", operator_name), scope.clone());
808
809    let (output, output_stream) = write_op.new_output::<CapacityContainerBuilder<Vec<_>>>();
810
811    let mut descriptions_input =
812        write_op.new_input_for(batch_descriptions.broadcast(), Pipeline, &output);
813    // Commitments only route updates into builders, so this input is disconnected: holding the
814    // output back on it would stall every batch behind the minter's own progress.
815    let mut commitments_input = write_op.new_disconnected_input(commitments.broadcast(), Pipeline);
816    let mut desired_input = write_op.new_disconnected_input(desired_collection.inner, Pipeline);
817
818    // This operator accepts the current and desired update streams for a `persist` shard.
819    // It attempts to write out updates, starting from the current's upper frontier, that
820    // will cause the changes of desired to be committed to persist, _but only those also past the
821    // upper_.
822
823    let shutdown_button = write_op.build(move |_capabilities| async move {
824        // Builders for timestamps no ceiling covers, keyed by timestamp.
825        //
826        // A batch builder cannot be split, so an update may only join updates at other timestamps
827        // once a bound they all fall below is known. Until then a timestamp gets a builder to
828        // itself, which is safe without knowing the descriptions because a description covers a
829        // timestamp entirely or not at all. It is finished only once its description is ready,
830        // since rows at its timestamp can still be staged after that description is processed.
831        let mut uncovered_builders: BTreeMap<mz_repr::Timestamp, SourceBatchBuilder> =
832            BTreeMap::new();
833
834        // The outstanding commitment: the bounds the open builder takes updates in, and the lower
835        // it declares. The ceiling is raised to the description's own upper once that arrives, so
836        // rows landing between arrival and readiness still join. `None` until the first commitment,
837        // which leaves every timestamp writing its own batch.
838        let mut commitment: Option<Commitment> = None;
839
840        // The one open builder, holding every update inside the commitment regardless of
841        // timestamp. It reaches `persist_blob_target_size` and spills its parts to blob, so what it
842        // holds resident is one unflushed part however long the frontier stays pinned. Only ever
843        // one, because the minter mints nothing below an outstanding ceiling, so there is exactly
844        // one description in flight for it to be finished under.
845        let mut open_builder: Option<SourceBatchBuilder> = None;
846
847        // Contains descriptions of batches for which we know that we can
848        // write data. We got these from the "centralized" operator that
849        // determines batch descriptions for all writers.
850        //
851        // `Antichain` does not implement `Ord`, so we cannot use a `BTreeMap`. We need to search
852        // through the map, so we cannot use the `mz_ore` wrapper either.
853        #[allow(clippy::disallowed_types)]
854        let mut in_flight_batches = std::collections::HashMap::<
855            (Antichain<mz_repr::Timestamp>, Antichain<mz_repr::Timestamp>),
856            Capability<mz_repr::Timestamp>,
857        >::new();
858
859        // TODO(aljoscha): We need to figure out what to do with error results from these calls.
860        let persist_client = persist_clients
861            .open(persist_location)
862            .await
863            .expect("could not open persist client");
864
865        let write = persist_client
866            .open_writer::<SourceData, (), mz_repr::Timestamp, StorageDiff>(
867                shard_id,
868                Arc::new(target_relation_desc),
869                Arc::new(UnitSchema),
870                Diagnostics {
871                    shard_name: collection_id.to_string(),
872                    handle_purpose: format!(
873                        "storage::persist_sink::write_batches {}",
874                        collection_id
875                    ),
876                },
877            )
878            .await
879            .expect("could not open persist shard");
880
881        // The current input frontiers.
882        let mut batch_descriptions_frontier = Antichain::from_elem(Timestamp::minimum());
883        let mut desired_frontier = Antichain::from_elem(Timestamp::minimum());
884
885        // The frontiers of the inputs we have processed, used to avoid redoing work
886        let mut processed_desired_frontier = Antichain::from_elem(Timestamp::minimum());
887        let mut processed_descriptions_frontier = Antichain::from_elem(Timestamp::minimum());
888
889        // A "safe" choice for the lower of new batches we are creating.
890        let mut operator_batch_lower = Antichain::from_elem(Timestamp::minimum());
891
892        while !(batch_descriptions_frontier.is_empty() && desired_frontier.is_empty()) {
893            // Wait for either inputs to become ready
894            tokio::select! {
895                _ = descriptions_input.ready() => {},
896                _ = commitments_input.ready() => {},
897                _ = desired_input.ready() => {},
898            }
899
900            // Collect ready work from all three inputs. Commitments and descriptions are processed
901            // before the data of the same round so that an update whose bound arrived alongside it
902            // can go straight into the open builder instead of writing a batch of its own.
903            let ready_commitments =
904                std::iter::from_fn(|| commitments_input.next_sync()).collect_vec();
905            let ready_descriptions =
906                std::iter::from_fn(|| descriptions_input.next_sync()).collect_vec();
907            let ready_events = std::iter::from_fn(|| desired_input.next_sync()).collect_vec();
908
909            // We now start the async work for the input we received. Until we finish the dataflow
910            // should be marked as busy.
911            let permit = busy_signal.acquire().await;
912
913            for event in ready_commitments {
914                let Event::Data(_cap, data) = event else {
915                    continue;
916                };
917                for next in data {
918                    // The minter only commits while the source is snapshotting, which pins the
919                    // frontier. The description that retires the commitment moves the frontier for
920                    // good. Every commitment during that snapshot shares one lower.
921                    let held = commitment.get_or_insert(next);
922                    assert_eq!(
923                        held.lower,
924                        next.lower,
925                        "persist_sink {collection_id}/{shard_id}: commitment {next:?} while {held:?} is outstanding"
926                    );
927                    held.ceiling = held.ceiling.max(next.ceiling);
928                }
929            }
930
931            for event in ready_descriptions {
932                match event {
933                    Event::Data(cap, data) => {
934                        // Ingest new batch descriptions.
935                        for description in data {
936                            if collection_id.is_user() {
937                                trace!(
938                                    "persist_sink {collection_id}/{shard_id}: \
939                                        write_batches: \
940                                        new_description: {:?}, \
941                                        desired_frontier: {:?}, \
942                                        batch_descriptions_frontier: {:?}",
943                                    description, desired_frontier, batch_descriptions_frontier,
944                                );
945                            }
946
947                            let (lower, upper) = (&description.0, &description.1);
948                            let lower_ts = *lower
949                                .as_option()
950                                .expect("minted descriptions have a single-element lower");
951
952                            // The description that retires a commitment ends at or past the
953                            // ceiling, so rows landing between the description arriving and the
954                            // description becoming ready can still join the builder that will
955                            // finish under that description.
956                            if let Some(held) = commitment.filter(|held| held.lower == lower_ts)
957                                && let Some(upper_ts) = upper.as_option()
958                            {
959                                commitment = Some(Commitment {
960                                    ceiling: held.ceiling.max(*upper_ts),
961                                    ..held
962                                });
963                            }
964
965                            match in_flight_batches.entry(description) {
966                                std::collections::hash_map::Entry::Vacant(v) => {
967                                    // `cap` is borrowed once per description, so each
968                                    // in-flight batch holds its own copy.
969                                    v.insert(sole_capability(&cap).clone());
970                                }
971                                std::collections::hash_map::Entry::Occupied(o) => {
972                                    let (description, _) = o.remove_entry();
973                                    panic!(
974                                        "write_batches: sink {} got more than one \
975                                            batch for description {:?}, in-flight: {:?}",
976                                        collection_id, description, in_flight_batches
977                                    );
978                                }
979                            }
980                        }
981                    }
982                    Event::Progress(frontier) => {
983                        batch_descriptions_frontier = frontier;
984                    }
985                }
986            }
987
988            for event in ready_events {
989                match event {
990                    Event::Data(_cap, data) => {
991                        // Extract desired rows as positive contributions to `correction`.
992                        if collection_id.is_user() && !data.is_empty() {
993                            trace!(
994                                "persist_sink {collection_id}/{shard_id}: \
995                                    updates: {:?}, \
996                                    in-flight-batches: {:?}, \
997                                    desired_frontier: {:?}, \
998                                    batch_descriptions_frontier: {:?}",
999                                data,
1000                                in_flight_batches,
1001                                desired_frontier,
1002                                batch_descriptions_frontier,
1003                            );
1004                        }
1005
1006                        for (row, ts, diff) in data {
1007                            if write.upper().less_equal(&ts) {
1008                                // Every description this operator has emitted was covered by the
1009                                // desired frontier at the time, so no update below
1010                                // `operator_batch_lower` can still be in flight. An update that
1011                                // arrives anyway belongs to a description that is already gone: it
1012                                // matches no later description, so its batch would never be
1013                                // appended and the update would be lost unnoticed. Not a
1014                                // `debug_assert!`, which compiles out of the optimized and release
1015                                // profiles and would leave the loss silent everywhere it matters.
1016                                assert!(
1017                                    operator_batch_lower.less_equal(&ts),
1018                                    "persist_sink {collection_id}/{shard_id}: update at {ts:?} \
1019                                    arrived below the emitted batch lower {operator_batch_lower:?}",
1020                                );
1021
1022                                let inside =
1023                                    commitment.filter(|held| held.lower <= ts && ts < held.ceiling);
1024                                let builder = if let Some(held) = inside {
1025                                    // The description retiring the commitment is known to contain
1026                                    // it, so the update joins the one open builder whatever its
1027                                    // timestamp.
1028                                    open_builder.get_or_insert_with(|| {
1029                                        BatchBuilderAndMetadata::new(
1030                                            write.builder(Antichain::from_elem(held.lower)),
1031                                        )
1032                                    })
1033                                } else {
1034                                    // Nothing to group under, so the only lower this builder can
1035                                    // declare is the operator's own, the one lower at or below
1036                                    // every description that could come to cover it. That
1037                                    // declaration is what registers the batch truncated once it is
1038                                    // appended under a description's narrower bounds.
1039                                    uncovered_builders.entry(ts).or_insert_with(|| {
1040                                        BatchBuilderAndMetadata::new(
1041                                            write.builder(operator_batch_lower.clone()),
1042                                        )
1043                                    })
1044                                };
1045                                stage_update(builder, row, ts, diff).await;
1046                                source_statistics.inc_updates_staged_by(1);
1047                            }
1048                        }
1049                    }
1050                    Event::Progress(frontier) => {
1051                        desired_frontier = frontier;
1052                    }
1053                }
1054            }
1055
1056            // We may have the opportunity to commit updates, if either frontier
1057            // has moved
1058            if PartialOrder::less_equal(&processed_desired_frontier, &desired_frontier)
1059                || PartialOrder::less_equal(
1060                    &processed_descriptions_frontier,
1061                    &batch_descriptions_frontier,
1062                )
1063            {
1064                trace!(
1065                    "persist_sink {collection_id}/{shard_id}: \
1066                        CAN emit: \
1067                        processed_desired_frontier: {:?}, \
1068                        processed_descriptions_frontier: {:?}, \
1069                        desired_frontier: {:?}, \
1070                        batch_descriptions_frontier: {:?}",
1071                    processed_desired_frontier,
1072                    processed_descriptions_frontier,
1073                    desired_frontier,
1074                    batch_descriptions_frontier,
1075                );
1076
1077                trace!(
1078                    "persist_sink {collection_id}/{shard_id}: \
1079                        in-flight batches: {:?}, \
1080                        batch_descriptions_frontier: {:?}, \
1081                        desired_frontier: {:?}",
1082                    in_flight_batches, batch_descriptions_frontier, desired_frontier,
1083                );
1084
1085                // We can write updates for a given batch description when
1086                // a) the batch is not beyond `batch_descriptions_frontier`,
1087                // and b) we know that we have seen all updates that would
1088                // fall into the batch, from `desired_frontier`.
1089                let ready_batches = in_flight_batches
1090                    .keys()
1091                    .filter(|(lower, upper)| {
1092                        !PartialOrder::less_equal(&batch_descriptions_frontier, lower)
1093                            && !PartialOrder::less_than(&desired_frontier, upper)
1094                    })
1095                    .cloned()
1096                    .collect::<Vec<_>>();
1097
1098                trace!(
1099                    "persist_sink {collection_id}/{shard_id}: \
1100                        ready batches: {:?}",
1101                    ready_batches,
1102                );
1103
1104                for batch_description in ready_batches {
1105                    let cap = in_flight_batches.remove(&batch_description).unwrap();
1106
1107                    if collection_id.is_user() {
1108                        trace!(
1109                            "persist_sink {collection_id}/{shard_id}: \
1110                                emitting done batch: {:?}, cap: {:?}",
1111                            batch_description, cap
1112                        );
1113                    }
1114
1115                    let (batch_lower, batch_upper) = batch_description;
1116                    let lower = *batch_lower
1117                        .as_option()
1118                        .expect("minted descriptions have a single-element lower");
1119
1120                    let covered: Vec<_> = uncovered_builders
1121                        .keys()
1122                        .copied()
1123                        .filter(|ts| batch_lower.less_equal(ts) && !batch_upper.less_equal(ts))
1124                        .collect();
1125                    let mut batch_tokens = Vec::with_capacity(covered.len() + 1);
1126                    for ts in covered {
1127                        let builder = uncovered_builders.remove(&ts).expect("just looked up");
1128                        batch_tokens.push(
1129                            builder
1130                                .finish(batch_lower.clone(), batch_upper.clone())
1131                                .await,
1132                        );
1133                    }
1134
1135                    // If snapshotting, and the as_of = T, where T is greater than the minimum,
1136                    // the minter emits a description covering `(0, T)` and `Commitment{lower:T}`.
1137                    // The open builder must wait for the appropriate description.
1138                    if let Some(held) = commitment.filter(|held| held.lower == lower)
1139                        && let Some(builder) = open_builder.take()
1140                    {
1141                        assert!(
1142                            !batch_upper.less_than(&held.ceiling),
1143                            "persist_sink {collection_id}/{shard_id}: description upper {batch_upper:?} is below commitment ceiling {commitment:?}",
1144                        );
1145                        commitment = None;
1146                        if collection_id.is_user() {
1147                            trace!(
1148                                "persist_sink {collection_id}/{shard_id}: \
1149                                    wrote batch from worker {}: ({:?}, {:?}), containing {:?}",
1150                                worker_index, batch_lower, batch_upper, builder.metrics
1151                            );
1152                        }
1153
1154                        batch_tokens.push(
1155                            builder
1156                                .finish(batch_lower.clone(), batch_upper.clone())
1157                                .await,
1158                        );
1159                    }
1160
1161                    // The next "safe" lower for batches is the meet (max) of all the emitted
1162                    // batches. These uppers all are not beyond the `desired_frontier`, which
1163                    // means all updates received by this operator will be beyond this lower.
1164                    // Additionally, the `mint_batch_descriptions` operator ensures that
1165                    // later-received batch descriptions will start beyond these uppers as
1166                    // well.
1167                    //
1168                    // It is impossible to emit a batch description that is
1169                    // beyond a not-yet emitted description in `in_flight_batches`, as
1170                    // a that description would also have been chosen as ready above.
1171                    operator_batch_lower = operator_batch_lower.join(&batch_upper);
1172
1173                    output.give_container(&cap, &mut batch_tokens);
1174
1175                    processed_desired_frontier.clone_from(&desired_frontier);
1176                    processed_descriptions_frontier.clone_from(&batch_descriptions_frontier);
1177                }
1178            } else {
1179                trace!(
1180                    "persist_sink {collection_id}/{shard_id}: \
1181                        cannot emit: processed_desired_frontier: {:?}, \
1182                        processed_descriptions_frontier: {:?}, \
1183                        desired_frontier: {:?}",
1184                    processed_desired_frontier, processed_descriptions_frontier, desired_frontier
1185                );
1186            }
1187            drop(permit);
1188        }
1189    });
1190
1191    // Use `inspect_core` instead of `inspect`. `inspect` carries a
1192    // `where for<'a> &'a C: IntoIterator` bound, and on macOS the solver can
1193    // satisfy that bound by chasing objc2's `&Retained<T>: IntoIterator`
1194    // blanket impl into an endless `Retained<Retained<…>>` chain, overflowing
1195    // the recursion limit. `inspect_core` has no such bound, so the cascade
1196    // never starts. We iterate the container by hand to recover the per-item
1197    // callback.
1198    let output_stream = if collection_id.is_user() {
1199        Inspect::inspect_core(output_stream, |event| {
1200            if let Ok((_, data)) = event {
1201                for d in data {
1202                    trace!("batch: {:?}", d);
1203                }
1204            }
1205        })
1206    } else {
1207        output_stream
1208    };
1209
1210    (output_stream, shutdown_button.press_on_drop())
1211}
1212
1213/// Fuses written batches together and appends them to persist using one
1214/// `compare_and_append` call. Writing only happens for batch descriptions where
1215/// we know that no future batches will arrive, that is, for those batch
1216/// descriptions that are not beyond the frontier of both the
1217/// `batch_descriptions` and `batches` inputs.
1218///
1219/// This also keeps the shared frontier that is stored in `compute_state` in
1220/// sync with the upper of the persist shard, and updates various metrics
1221/// and statistics objects.
1222fn append_batches<'scope>(
1223    scope: Scope<'scope, mz_repr::Timestamp>,
1224    collection_id: GlobalId,
1225    operator_name: String,
1226    target: &CollectionMetadata,
1227    batch_descriptions: Stream<
1228        'scope,
1229        mz_repr::Timestamp,
1230        Vec<(Antichain<mz_repr::Timestamp>, Antichain<mz_repr::Timestamp>)>,
1231    >,
1232    batches: StreamVec<'scope, mz_repr::Timestamp, HollowBatchAndMetadata<mz_repr::Timestamp>>,
1233    persist_clients: Arc<PersistClientCache>,
1234    storage_state: &StorageState,
1235    metrics: SourcePersistSinkMetrics,
1236    busy_signal: Arc<Semaphore>,
1237) -> (
1238    StreamVec<'scope, mz_repr::Timestamp, ()>,
1239    StreamVec<'scope, mz_repr::Timestamp, Rc<anyhow::Error>>,
1240    PressOnDropButton,
1241) {
1242    let persist_location = target.persist_location.clone();
1243    let shard_id = target.data_shard;
1244    let target_relation_desc = target.relation_desc.clone();
1245
1246    // We can only be lenient with concurrent modifications when we know that
1247    // this source pipeline is using the feedback upsert operator, which works
1248    // correctly when multiple instances of an ingestion pipeline produce
1249    // different updates, because of concurrency/non-determinism.
1250    let use_continual_feedback_upsert = dyncfgs::STORAGE_USE_CONTINUAL_FEEDBACK_UPSERT
1251        .get(storage_state.storage_configuration.config_set());
1252    let bail_on_concurrent_modification = !use_continual_feedback_upsert;
1253
1254    let mut read_only_rx = storage_state.read_only_rx.clone();
1255
1256    let operator_name = format!("{} append_batches", operator_name);
1257    let mut append_op = AsyncOperatorBuilder::new(operator_name, scope.clone());
1258
1259    let hashed_id = collection_id.hashed();
1260    let active_worker = usize::cast_from(hashed_id) % scope.peers() == scope.index();
1261    let worker_id = scope.index();
1262
1263    // Both of these inputs are disconnected from the output capabilities of this operator, as
1264    // any output of this operator is entirely driven by the `compare_and_append`s. Currently
1265    // this operator has no outputs, but they may be added in the future, when merging with
1266    // the compute `persist_sink`.
1267    let mut descriptions_input =
1268        append_op.new_disconnected_input(batch_descriptions, Exchange::new(move |_| hashed_id));
1269    let mut batches_input =
1270        append_op.new_disconnected_input(batches, Exchange::new(move |_| hashed_id));
1271
1272    let current_upper = Rc::clone(&storage_state.source_uppers[&collection_id]);
1273    if !active_worker {
1274        // This worker is not writing, so make sure it's "taken out" of the
1275        // calculation by advancing to the empty frontier.
1276        current_upper.borrow_mut().clear();
1277    }
1278
1279    let source_statistics = storage_state
1280        .aggregated_statistics
1281        .get_source(&collection_id)
1282        .expect("statistics initialized")
1283        .clone();
1284
1285    // An output whose frontier tracks the last successful compare and append of this operator
1286    let (_upper_output, upper_stream) = append_op.new_output::<CapacityContainerBuilder<Vec<_>>>();
1287
1288    // This operator accepts the batch descriptions and tokens that represent
1289    // written batches. Written batches get appended to persist when we learn
1290    // from our input frontiers that we have seen all batches for a given batch
1291    // description.
1292
1293    let (shutdown_button, errors) = append_op.build_fallible(move |caps| Box::pin(async move {
1294        let [upper_cap_set]: &mut [_; 1] = caps.try_into().unwrap();
1295
1296        // This may SEEM unnecessary, but metrics contains extra
1297        // `DeregisterOnDrop`-wrapped fields that will NOT be moved into this
1298        // closure otherwise, dropping and destroying
1299        // those metrics. This is because rust now only moves the
1300        // explicitly-referenced fields into closures.
1301        let metrics = metrics;
1302
1303        // Contains descriptions of batches for which we know that we can
1304        // write data. We got these from the "centralized" operator that
1305        // determines batch descriptions for all writers.
1306        //
1307        // `Antichain` does not implement `Ord`, so we cannot use a `BTreeSet`. We need to search
1308        // through the set, so we cannot use the `mz_ore` wrapper either.
1309        #[allow(clippy::disallowed_types)]
1310        let mut in_flight_descriptions = std::collections::HashSet::<(
1311            Antichain<mz_repr::Timestamp>,
1312            Antichain<mz_repr::Timestamp>,
1313        )>::new();
1314
1315        // In flight batches that haven't been `compare_and_append`'d yet, plus metrics about
1316        // the batch.
1317        let mut in_flight_batches = HashMap::<
1318            (Antichain<mz_repr::Timestamp>, Antichain<mz_repr::Timestamp>),
1319            BatchSet,
1320        >::new();
1321
1322        source_statistics.initialize_rehydration_latency_ms();
1323        if !active_worker {
1324            // The non-active workers report that they are done snapshotting and hydrating.
1325            let empty_frontier = Antichain::new();
1326            source_statistics.initialize_snapshot_committed(&empty_frontier);
1327            source_statistics.update_rehydration_latency_ms(&empty_frontier);
1328            return Ok(());
1329        }
1330
1331        let persist_client = persist_clients
1332            .open(persist_location)
1333            .await?;
1334
1335        let mut write = persist_client
1336            .open_writer::<SourceData, (), mz_repr::Timestamp, StorageDiff>(
1337                shard_id,
1338                Arc::new(target_relation_desc),
1339                Arc::new(UnitSchema),
1340                Diagnostics {
1341                    shard_name:collection_id.to_string(),
1342                    handle_purpose: format!("persist_sink::append_batches {}", collection_id)
1343                },
1344            )
1345            .await?;
1346
1347        // Initialize this sink's `upper` to the `upper` of the persist shard we are writing
1348        // to. Data from the source not beyond this time will be dropped, as it has already
1349        // been persisted.
1350        // In the future, sources will avoid passing through data not beyond this upper
1351        // VERY IMPORTANT: Only the active write worker must change the
1352        // shared upper. All other workers have already cleared this
1353        // upper above.
1354        current_upper.borrow_mut().clone_from(write.upper());
1355        upper_cap_set.downgrade(current_upper.borrow().iter());
1356        source_statistics.initialize_snapshot_committed(write.upper());
1357
1358        // The current input frontiers.
1359        let mut batch_description_frontier = Antichain::from_elem(Timestamp::minimum());
1360        let mut batches_frontier = Antichain::from_elem(Timestamp::minimum());
1361
1362        loop {
1363            tokio::select! {
1364                Some(event) = descriptions_input.next() => {
1365                    match event {
1366                        Event::Data(_cap, data) => {
1367                            // Ingest new batch descriptions.
1368                            for batch_description in data {
1369                                if collection_id.is_user() {
1370                                    trace!(
1371                                        "persist_sink {collection_id}/{shard_id}: \
1372                                            append_batches: sink {}, \
1373                                            new description: {:?}, \
1374                                            batch_description_frontier: {:?}",
1375                                        collection_id,
1376                                        batch_description,
1377                                        batch_description_frontier
1378                                    );
1379                                }
1380
1381                                // This line has to be broken up, or
1382                                // rustfmt fails in the whole function :(
1383                                let is_new = in_flight_descriptions.insert(
1384                                    batch_description.clone()
1385                                );
1386
1387                                assert!(
1388                                    is_new,
1389                                    "append_batches: sink {} got more than one batch \
1390                                        for a given description in-flight: {:?}",
1391                                    collection_id, in_flight_batches
1392                                );
1393                            }
1394
1395                            continue;
1396                        }
1397                        Event::Progress(frontier) => {
1398                            batch_description_frontier = frontier;
1399                        }
1400                    }
1401                }
1402                Some(event) = batches_input.next() => {
1403                    match event {
1404                        Event::Data(_cap, data) => {
1405                            for batch in data {
1406                                let batch_description = (batch.lower.clone(), batch.upper.clone());
1407
1408                                let batches = in_flight_batches
1409                                    .entry(batch_description)
1410                                    .or_default();
1411
1412                                batches.finished.push(FinishedBatch {
1413                                    batch: write.batch_from_transmittable_batch(batch.batch),
1414                                    data_max_ts: batch.data_max_ts,
1415                                });
1416                                batches.batch_metrics += &batch.metrics;
1417                            }
1418                            continue;
1419                        }
1420                        Event::Progress(frontier) => {
1421                            batches_frontier = frontier;
1422                        }
1423                    }
1424                }
1425                else => {
1426                    // All inputs are exhausted, so we can shut down.
1427                    return Ok(());
1428                }
1429            };
1430
1431            // Peel off any batches that are not beyond the frontier
1432            // anymore.
1433            //
1434            // It is correct to consider batches that are not beyond the
1435            // `batches_frontier` because it is held back by the writer
1436            // operator as long as a) the `batch_description_frontier` did
1437            // not advance and b) as long as the `desired_frontier` has not
1438            // advanced to the `upper` of a given batch description.
1439
1440            let mut done_batches = in_flight_descriptions
1441                .iter()
1442                .filter(|(lower, _upper)| !PartialOrder::less_equal(&batches_frontier, lower))
1443                .cloned()
1444                .collect::<Vec<_>>();
1445
1446            trace!(
1447                "persist_sink {collection_id}/{shard_id}: \
1448                    append_batches: in_flight: {:?}, \
1449                    done: {:?}, \
1450                    batch_frontier: {:?}, \
1451                    batch_description_frontier: {:?}",
1452                in_flight_descriptions,
1453                done_batches,
1454                batches_frontier,
1455                batch_description_frontier
1456            );
1457
1458            // Append batches in order, to ensure that their `lower` and
1459            // `upper` line up.
1460            done_batches.sort_by(|a, b| {
1461                if PartialOrder::less_than(a, b) {
1462                    Ordering::Less
1463                } else if PartialOrder::less_than(b, a) {
1464                    Ordering::Greater
1465                } else {
1466                    Ordering::Equal
1467                }
1468            });
1469
1470            let validate_part_bounds_on_write = write.validate_part_bounds_on_write();
1471            let mut todo = VecDeque::new();
1472
1473            if validate_part_bounds_on_write {
1474                // Persist will expect each batch's bounds to match the append-time bounds; write them separately.
1475                for done_batch_metadata in done_batches.drain(..) {
1476                    in_flight_descriptions.remove(&done_batch_metadata);
1477                    let batch_set = in_flight_batches
1478                        .remove(&done_batch_metadata)
1479                        .unwrap_or_default();
1480                    todo.push_back((done_batch_metadata, batch_set));
1481                }
1482            } else {
1483                // Persist should allow batches to be written as part of a single append even when the bounds don't
1484                // match exactly; group all eligible batches together.
1485                let mut combined_batch_metadata = None;
1486                let mut combined_batch_set = BatchSet::default();
1487                for done_batch_metadata in done_batches.drain(..) {
1488                    in_flight_descriptions.remove(&done_batch_metadata);
1489                    let mut batch_set = in_flight_batches
1490                        .remove(&done_batch_metadata)
1491                        .unwrap_or_default();
1492                    match combined_batch_metadata.as_mut() {
1493                        Some((_, upper)) => *upper = done_batch_metadata.1,
1494                        None => combined_batch_metadata = Some(done_batch_metadata),
1495                    }
1496                    combined_batch_set.batch_metrics += &batch_set.batch_metrics;
1497                    combined_batch_set.finished.append(&mut batch_set.finished);
1498                }
1499                if let Some(done_batch_metadata) = combined_batch_metadata {
1500                    todo.push_back((done_batch_metadata, combined_batch_set))
1501                }
1502            };
1503
1504            while let Some((done_batch_metadata, batch_set)) = todo.pop_front() {
1505                in_flight_descriptions.remove(&done_batch_metadata);
1506
1507                let mut batches = batch_set.finished;
1508
1509                trace!(
1510                    "persist_sink {collection_id}/{shard_id}: \
1511                        done batch: {:?}, {:?}",
1512                    done_batch_metadata,
1513                    batches
1514                );
1515
1516                let (batch_lower, batch_upper) = done_batch_metadata;
1517
1518                let batch_metrics = batch_set.batch_metrics;
1519
1520                let mut to_append = batches.iter_mut().map(|b| &mut b.batch).collect::<Vec<_>>();
1521
1522                let result = {
1523                    let maybe_err = if *read_only_rx.borrow() {
1524
1525                        // We have to wait for either us coming out of read-only
1526                        // mode or someone else applying a write that covers our
1527                        // batch.
1528                        //
1529                        // If we didn't wait for the latter here, and just go
1530                        // around the loop again, we might miss a moment where
1531                        // _we_ have to write down a batch. For example when our
1532                        // input frontier advances to a state where we can
1533                        // write, and the read-write instance sees the same
1534                        // update but then crashes before it can append a batch.
1535
1536                        let maybe_err = loop {
1537                            if collection_id.is_user() {
1538                                tracing::debug!(
1539                                    %worker_id,
1540                                    %collection_id,
1541                                    %shard_id,
1542                                    ?batch_lower,
1543                                    ?batch_upper,
1544                                    ?current_upper,
1545                                    "persist_sink is in read-only mode, waiting until we come out of it or the shard upper advances"
1546                                );
1547                            }
1548
1549                            // We don't try to be smart here, and for example
1550                            // use `wait_for_upper_past()`. We'd have to use a
1551                            // select!, which would require cancel safety of
1552                            // `wait_for_upper_past()`, which it doesn't
1553                            // advertise.
1554                            let _ = tokio::time::timeout(
1555                                Duration::from_secs(1),
1556                                read_only_rx.changed(),
1557                            )
1558                            .await;
1559
1560                            if !*read_only_rx.borrow() {
1561                                if collection_id.is_user() {
1562                                    tracing::debug!(
1563                                        %worker_id,
1564                                        %collection_id,
1565                                        %shard_id,
1566                                        ?batch_lower,
1567                                        ?batch_upper,
1568                                        ?current_upper,
1569                                        "persist_sink has come out of read-only mode"
1570                                    );
1571                                }
1572
1573                                // It's okay to write now.
1574                                break Ok(());
1575                            }
1576
1577                            let current_upper = write.fetch_recent_upper().await;
1578
1579                            if PartialOrder::less_than(&batch_upper, current_upper) {
1580                                // We synthesize an `UpperMismatch` so that we can go
1581                                // through the same logic below for trimming down our
1582                                // batches.
1583                                //
1584                                // Notably, we are not trying to be smart, and teach the
1585                                // write operator about read-only mode. Writing down
1586                                // those batches does not append anything to the persist
1587                                // shard, and it would be a hassle to figure out in the
1588                                // write workers how to trim down batches in read-only
1589                                // mode, when the shard upper advances.
1590                                //
1591                                // Right here, in the logic below, we have all we need
1592                                // for figuring out how to trim our batches.
1593
1594                                if collection_id.is_user() {
1595                                    tracing::debug!(
1596                                        %worker_id,
1597                                        %collection_id,
1598                                        %shard_id,
1599                                        ?batch_lower,
1600                                        ?batch_upper,
1601                                        ?current_upper,
1602                                        "persist_sink not appending in read-only mode"
1603                                    );
1604                                }
1605
1606                                break Err(UpperMismatch {
1607                                    current: current_upper.clone(),
1608                                    expected: batch_lower.clone()}
1609                                );
1610                            }
1611                        };
1612
1613                        maybe_err
1614                    } else {
1615                        // It's okay to proceed with the write.
1616                        Ok(())
1617                    };
1618
1619                    match maybe_err {
1620                        Ok(()) => {
1621                            let _permit = busy_signal.acquire().await;
1622
1623                            write.compare_and_append_batch(
1624                                &mut to_append[..],
1625                                batch_lower.clone(),
1626                                batch_upper.clone(),
1627                                validate_part_bounds_on_write,
1628                            )
1629                            .await
1630                            .expect("Invalid usage")
1631                        },
1632                        Err(e) => {
1633                            // We forward the synthesize error message, so that
1634                            // we go though the batch cleanup logic below.
1635                            Err(e)
1636                        }
1637                    }
1638                };
1639
1640
1641                // These metrics are independent of whether it was _us_ or
1642                // _someone_ that managed to commit a batch that advanced the
1643                // upper.
1644                source_statistics.update_snapshot_committed(&batch_upper);
1645                source_statistics.update_rehydration_latency_ms(&batch_upper);
1646                metrics
1647                    .progress
1648                    .set(mz_persist_client::metrics::encode_ts_metric(&batch_upper));
1649
1650                if collection_id.is_user() {
1651                    trace!(
1652                        "persist_sink {collection_id}/{shard_id}: \
1653                            append result for batch ({:?} -> {:?}): {:?}",
1654                        batch_lower,
1655                        batch_upper,
1656                        result
1657                    );
1658                }
1659
1660                match result {
1661                    Ok(()) => {
1662                        // Only update these metrics when we know that _we_ were
1663                        // successful.
1664                        let committed =
1665                            batch_metrics.inserts + batch_metrics.retractions;
1666                        source_statistics
1667                            .inc_updates_committed_by(committed);
1668                        metrics.processed_batches.inc();
1669                        metrics.row_inserts.inc_by(batch_metrics.inserts);
1670                        metrics.row_retractions.inc_by(batch_metrics.retractions);
1671                        metrics.error_inserts.inc_by(batch_metrics.error_inserts);
1672                        metrics
1673                            .error_retractions
1674                            .inc_by(batch_metrics.error_retractions);
1675
1676                        current_upper.borrow_mut().clone_from(&batch_upper);
1677                        upper_cap_set.downgrade(current_upper.borrow().iter());
1678                    }
1679                    Err(mismatch) => {
1680                        // We tried to to a non-contiguous append, that won't work.
1681                        if PartialOrder::less_than(&mismatch.current, &batch_lower) {
1682                            // Best-effort attempt to delete unneeded batches.
1683                            future::join_all(batches.into_iter().map(|b| b.batch.delete())).await;
1684
1685                            // We always bail when this happens, regardless of
1686                            // `bail_on_concurrent_modification`.
1687                            tracing::warn!(
1688                                "persist_sink({}): invalid upper! \
1689                                    Tried to append batch ({:?} -> {:?}) but upper \
1690                                    is {:?}. This is surpising and likely indicates \
1691                                    a bug in the persist sink, but we'll restart the \
1692                                    dataflow and try again.",
1693                                collection_id, batch_lower, batch_upper, mismatch.current,
1694                            );
1695                            anyhow::bail!("collection concurrently modified. Ingestion dataflow will be restarted");
1696                        } else if PartialOrder::less_than(&mismatch.current, &batch_upper) {
1697                            // The shard's upper was ahead of our batch's lower
1698                            // but not ahead of our upper. Cut down the
1699                            // description by advancing its lower to the current
1700                            // shard upper and try again. IMPORTANT: We can only
1701                            // advance the lower, meaning we cut updates away,
1702                            // we must not "extend" the batch by changing to a
1703                            // lower that is not beyond the current lower. This
1704                            // invariant is checked by the first if branch: if
1705                            // `!(current_upper < lower)` then it holds that
1706                            // `lower <= current_upper`.
1707
1708                            // First, construct a new batch description with the
1709                            // lower advanced to the current shard upper.
1710                            let new_batch_lower = mismatch.current.clone();
1711                            let new_done_batch_metadata =
1712                                (new_batch_lower.clone(), batch_upper.clone());
1713
1714                            // Re-append every batch that still holds something we owe, under the
1715                            // narrowed description. A batch may hold data on both sides of the new
1716                            // lower: persist registers it truncated and filters the updates
1717                            // outside the registered bounds on read, so the ones the concurrent
1718                            // writer already committed do not come back. A batch entirely below
1719                            // the new lower owes nothing and is deleted instead, to keep parts
1720                            // that would be truncated away in full out of shard state.
1721                            let mut batch_delete_futures = vec![];
1722                            let mut new_batch_set = BatchSet::default();
1723                            for batch in batches {
1724                                if new_batch_lower.less_equal(&batch.data_max_ts) {
1725                                    new_batch_set.finished.push(batch);
1726                                } else {
1727                                    batch_delete_futures.push(batch.batch.delete());
1728                                }
1729                            }
1730
1731                            // Re-add the new batch to the list of batches to process.
1732                            todo.push_front((new_done_batch_metadata, new_batch_set));
1733
1734                            // Best-effort attempt to delete unneeded batches.
1735                            future::join_all(batch_delete_futures).await;
1736                        } else {
1737                            // Best-effort attempt to delete unneeded batches.
1738                            future::join_all(batches.into_iter().map(|b| b.batch.delete())).await;
1739                        }
1740
1741                        if bail_on_concurrent_modification {
1742                            tracing::warn!(
1743                                "persist_sink({}): invalid upper! \
1744                                    Tried to append batch ({:?} -> {:?}) but upper \
1745                                    is {:?}. This is not a problem, it just means \
1746                                    someone else was faster than us. We will try \
1747                                    again with a new batch description.",
1748                                collection_id, batch_lower, batch_upper, mismatch.current,
1749                            );
1750                            anyhow::bail!("collection concurrently modified. Ingestion dataflow will be restarted");
1751                        }
1752                    }
1753                }
1754            }
1755        }
1756    }));
1757
1758    (upper_stream, errors, shutdown_button.press_on_drop())
1759}
1760
1761#[cfg(test)]
1762mod tests {
1763    use std::cell::RefCell;
1764    use std::str::FromStr;
1765
1766    use mz_build_info::DUMMY_BUILD_INFO;
1767    use mz_dyncfg::{ConfigUpdates, ConfigVal};
1768    use mz_ore::metrics::MetricsRegistry;
1769    use mz_ore::now::SYSTEM_TIME;
1770    use mz_ore::url::SensitiveUrl;
1771    use mz_persist_client::PersistLocation;
1772    use mz_persist_client::cfg::PersistConfig;
1773    use mz_persist_client::rpc::PubSubClientConnection;
1774    use mz_persist_types::ShardId;
1775    use mz_repr::{Datum, RelationDesc, SqlScalarType};
1776    use mz_storage_types::sources::SourceEnvelope;
1777    use mz_storage_types::sources::envelope::{KeyEnvelope, NoneEnvelope};
1778    use timely::dataflow::operators::Input;
1779
1780    use crate::statistics::SourceStatisticsMetricDefs;
1781
1782    use super::*;
1783
1784    fn ts(t: u64) -> mz_repr::Timestamp {
1785        t.into()
1786    }
1787
1788    fn frontier(t: u64) -> Antichain<mz_repr::Timestamp> {
1789        Antichain::from_elem(ts(t))
1790    }
1791
1792    /// One step of a `write_batches` script.
1793    #[derive(Clone)]
1794    enum Step {
1795        /// Deliver a batch description, as `mint_batch_descriptions` would.
1796        Description(u64, u64),
1797        /// Deliver a commitment, as `mint_batch_descriptions` would.
1798        Commit(u64, u64),
1799        /// Deliver `count` updates at time `at`.
1800        Updates(u64, usize),
1801        /// Deliver `count` updates at time `at` with negated diffs, as the rewind of a snapshot
1802        /// does for rows the replication stream redelivers at their true offset.
1803        Retractions(u64, usize),
1804        /// Advance both input frontiers.
1805        AdvanceTo(u64),
1806        /// Advance the remap upper to `t`, as the remap operator's next binding would.
1807        Tick(u64),
1808    }
1809
1810    /// What a batch emitted by `write_batches` carries, flattened for assertions.
1811    #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
1812    struct EmittedBatch {
1813        lower: u64,
1814        upper: u64,
1815        data_max_ts: u64,
1816        inserts: u64,
1817    }
1818
1819    /// Runs a timely worker to completion on a blocking thread.
1820    ///
1821    /// The test body is a single poll of the runtime's `block_on` future, so it runs under one
1822    /// tokio cooperative budget. A worker driven inline spends that budget on the operators'
1823    /// `select!` and semaphore polls, and once it is gone every such poll returns `Pending` and
1824    /// re-wakes itself, parking the operator for good with no error. Blocking threads have no
1825    /// budget.
1826    async fn run_worker<T: Send + 'static>(
1827        worker: impl FnOnce(&mut timely::worker::Worker) -> T + Send + Sync + 'static,
1828    ) -> T {
1829        mz_ore::task::spawn_blocking(
1830            || "persist_sink_test_worker",
1831            move || timely::execute_directly(worker),
1832        )
1833        .await
1834    }
1835
1836    /// Drives `write_batches` through `script` and returns the batches it emitted, along with a
1837    /// handle to the shard so callers can append them and read the result back.
1838    async fn run_write_batches(
1839        target: CollectionMetadata,
1840        persist_clients: Arc<PersistClientCache>,
1841        script: Vec<Step>,
1842    ) -> Vec<(EmittedBatch, ProtoBatch)> {
1843        run_worker(move |worker| {
1844            // `ProtoBatch` is not `Ord`, so the captured stream is summarized on the way out
1845            // rather than going through `Capture`.
1846            let emitted = Rc::new(RefCell::new(Vec::new()));
1847
1848            let (mut descs_input, mut ceilings_input, mut data_input, button) =
1849                worker.dataflow::<mz_repr::Timestamp, _, _>(|scope| {
1850                    let (descs_input, descs) = scope.new_input();
1851                    let (ceilings_input, ceilings) = scope.new_input();
1852                    let (data_input, data) = scope.new_input();
1853
1854                    let source_id = GlobalId::User(0);
1855                    let stats_defs =
1856                        SourceStatisticsMetricDefs::register_with(&MetricsRegistry::new());
1857                    let source_statistics = SourceStatistics::new(
1858                        source_id,
1859                        0,
1860                        &stats_defs,
1861                        source_id,
1862                        &target.data_shard,
1863                        SourceEnvelope::None(NoneEnvelope {
1864                            key_envelope: KeyEnvelope::None,
1865                            key_arity: 0,
1866                        }),
1867                        Antichain::from_elem(Timestamp::minimum()),
1868                        Antichain::from_elem(Timestamp::minimum()),
1869                    );
1870
1871                    let (batches, button) = write_batches(
1872                        scope,
1873                        source_id,
1874                        "test",
1875                        &target,
1876                        descs,
1877                        ceilings,
1878                        data.as_collection(),
1879                        persist_clients,
1880                        source_statistics,
1881                        Arc::new(Semaphore::new(Semaphore::MAX_PERMITS)),
1882                    );
1883                    let sink = Rc::clone(&emitted);
1884                    Inspect::inspect_core(batches, move |event| {
1885                        if let Ok((_, data)) = event {
1886                            for b in data {
1887                                sink.borrow_mut().push((
1888                                    EmittedBatch {
1889                                        lower: b.lower.as_option().expect("single lower").into(),
1890                                        upper: b.upper.as_option().expect("single upper").into(),
1891                                        data_max_ts: b.data_max_ts.into(),
1892                                        inserts: b.metrics.inserts,
1893                                    },
1894                                    b.batch.clone(),
1895                                ));
1896                            }
1897                        }
1898                    });
1899
1900                    (descs_input, ceilings_input, data_input, button)
1901                });
1902
1903            // We want the operator to finish processing before we advance the script,
1904            // but when the operator is waiting for persist,
1905            // a plain timely `step` will find no work to do and return immediately.
1906            // So we repeatedly `step_or_park`, where `park`ing the thread forces a delay,
1907            // allowing the operator to finish waiting for persist and do its work on the next `step`.
1908            fn pump(worker: &mut timely::worker::Worker) {
1909                // no solid reason for this number, it's a selection that seems high enough to
1910                // work reliably, but not create a ton of delay (~32ms parked).
1911                for _ in 0..32 {
1912                    worker.step_or_park(Some(Duration::from_millis(1)));
1913                }
1914            }
1915
1916            // Twice, so the operator is past opening its persist handles before the script runs.
1917            pump(worker);
1918            pump(worker);
1919
1920            for step in script {
1921                match step {
1922                    Step::Description(lower, upper) => {
1923                        descs_input.send((frontier(lower), frontier(upper)));
1924                    }
1925                    Step::Commit(lower, ceiling) => ceilings_input.send(Commitment {
1926                        lower: ts(lower),
1927                        ceiling: ts(ceiling),
1928                    }),
1929                    Step::Updates(at, count) => {
1930                        for i in 0..i64::try_from(count).expect("small count") {
1931                            let row = Row::pack_slice(&[Datum::Int64(i)]);
1932                            data_input.send((Ok(row), ts(at), Diff::ONE));
1933                        }
1934                    }
1935                    Step::Retractions(at, count) => {
1936                        for i in 0..i64::try_from(count).expect("small count") {
1937                            let row = Row::pack_slice(&[Datum::Int64(i)]);
1938                            data_input.send((Ok(row), ts(at), -Diff::ONE));
1939                        }
1940                    }
1941                    Step::AdvanceTo(t) => {
1942                        descs_input.advance_to(ts(t));
1943                        ceilings_input.advance_to(ts(t));
1944                        data_input.advance_to(ts(t));
1945                    }
1946                    Step::Tick(..) => panic!("the writer takes no remap upper"),
1947                }
1948                // NOTE: `send` buffers until its container fills, so without a flush every step
1949                // before the next `advance_to` would reach the operator together, in one round.
1950                descs_input.flush();
1951                ceilings_input.flush();
1952                data_input.flush();
1953                pump(worker);
1954            }
1955
1956            descs_input.close();
1957            ceilings_input.close();
1958            data_input.close();
1959            for _ in 0..1_000 {
1960                if !worker.step_or_park(Some(Duration::from_millis(1))) {
1961                    break;
1962                }
1963            }
1964
1965            drop(button);
1966            while worker.step() {}
1967
1968            let mut emitted = emitted.borrow().clone();
1969            emitted.sort_by(|a, b| a.0.cmp(&b.0));
1970            emitted
1971        })
1972        .await
1973    }
1974
1975    /// What `mint_batch_descriptions` put on each of its two outputs.
1976    #[derive(Debug, Default, PartialEq, Eq)]
1977    struct Minted {
1978        descriptions: Vec<(u64, u64)>,
1979        commitments: Vec<(u64, u64)>,
1980    }
1981
1982    /// Drives `mint_batch_descriptions` through `script` and returns what it emitted.
1983    ///
1984    /// Closing the input mints a final description with an empty upper, which is dropped here so a
1985    /// script only sees what it drove.
1986    async fn run_mint_batch_descriptions(
1987        target: CollectionMetadata,
1988        persist_clients: Arc<PersistClientCache>,
1989        lookahead: Option<u64>,
1990        snapshot_time: Option<Antichain<mz_repr::Timestamp>>,
1991        script: Vec<Step>,
1992    ) -> Minted {
1993        run_worker(move |worker| {
1994            let emitted = Rc::new(RefCell::new(Minted::default()));
1995
1996            let (mut data_input, mut remap_input, button) = worker
1997                .dataflow::<mz_repr::Timestamp, _, _>(|scope| {
1998                    let (data_input, data) = scope.new_input();
1999                    let (remap_input, remap_upper) = scope.new_input::<Vec<()>>();
2000
2001                    let (descriptions, ceilings, _passthrough, button) = mint_batch_descriptions(
2002                        scope,
2003                        GlobalId::User(0),
2004                        "test",
2005                        &target,
2006                        data.as_collection(),
2007                        remap_upper,
2008                        persist_clients,
2009                        lookahead,
2010                        snapshot_time,
2011                    );
2012
2013                    let sink = Rc::clone(&emitted);
2014                    Inspect::inspect_core(descriptions, move |event| {
2015                        if let Ok((_, data)) = event {
2016                            for (lower, upper) in data {
2017                                let (Some(lower), Some(upper)) =
2018                                    (lower.as_option(), upper.as_option())
2019                                else {
2020                                    continue;
2021                                };
2022                                sink.borrow_mut()
2023                                    .descriptions
2024                                    .push((lower.into(), upper.into()));
2025                            }
2026                        }
2027                    });
2028
2029                    let sink = Rc::clone(&emitted);
2030                    Inspect::inspect_core(ceilings, move |event| {
2031                        if let Ok((_, data)) = event {
2032                            for Commitment { lower, ceiling } in data {
2033                                sink.borrow_mut()
2034                                    .commitments
2035                                    .push(((*lower).into(), (*ceiling).into()));
2036                            }
2037                        }
2038                    });
2039
2040                    (data_input, remap_input, button)
2041                });
2042
2043            // The operator waits on persist off the timely scheduler, so a plain `step` can find
2044            // the worker idle while the operator is still starting up.
2045            fn pump(worker: &mut timely::worker::Worker) {
2046                for _ in 0..32 {
2047                    worker.step_or_park(Some(Duration::from_millis(1)));
2048                }
2049            }
2050
2051            // Twice, so the operator has fetched the shard upper before the script runs.
2052            pump(worker);
2053            pump(worker);
2054
2055            for step in script {
2056                match step {
2057                    Step::Description(..) | Step::Commit(..) => {
2058                        panic!("the minter emits these rather than taking them")
2059                    }
2060                    Step::Updates(at, count) => {
2061                        for i in 0..i64::try_from(count).expect("small count") {
2062                            let row = Row::pack_slice(&[Datum::Int64(i)]);
2063                            data_input.send((Ok(row), ts(at), Diff::ONE));
2064                        }
2065                    }
2066                    Step::Retractions(at, count) => {
2067                        for i in 0..i64::try_from(count).expect("small count") {
2068                            let row = Row::pack_slice(&[Datum::Int64(i)]);
2069                            data_input.send((Ok(row), ts(at), -Diff::ONE));
2070                        }
2071                    }
2072                    Step::AdvanceTo(t) => data_input.advance_to(ts(t)),
2073                    Step::Tick(t) => remap_input.advance_to(ts(t)),
2074                }
2075                // Sends are buffered until the handle is flushed or its time advances, and these
2076                // scripts deliver data at times a pinned frontier never reaches.
2077                data_input.flush();
2078                pump(worker);
2079            }
2080
2081            data_input.close();
2082            remap_input.close();
2083            for _ in 0..1_000 {
2084                if !worker.step_or_park(Some(Duration::from_millis(1))) {
2085                    break;
2086                }
2087            }
2088
2089            drop(button);
2090            while worker.step() {}
2091
2092            emitted.take()
2093        })
2094        .await
2095    }
2096
2097    /// Drives the real `mint_batch_descriptions` into the real `write_batches`, so the ordering
2098    /// between a committed ceiling and the data it covers is the operators' own rather than the
2099    /// script's.
2100    async fn run_sink_pipeline(
2101        target: CollectionMetadata,
2102        persist_clients: Arc<PersistClientCache>,
2103        lookahead: Option<u64>,
2104        snapshot_time: Option<Antichain<mz_repr::Timestamp>>,
2105        script: Vec<Step>,
2106    ) -> Vec<EmittedBatch> {
2107        run_worker(move |worker| {
2108            let emitted = Rc::new(RefCell::new(Vec::new()));
2109
2110            let (mut data_input, mut remap_input, buttons) = worker
2111                .dataflow::<mz_repr::Timestamp, _, _>(|scope| {
2112                    let (data_input, data) = scope.new_input();
2113                    let (remap_input, remap_upper) = scope.new_input::<Vec<()>>();
2114
2115                    let source_id = GlobalId::User(0);
2116                    let (descriptions, ceilings, passthrough, mint_button) =
2117                        mint_batch_descriptions(
2118                            scope,
2119                            source_id,
2120                            "test",
2121                            &target,
2122                            data.as_collection(),
2123                            remap_upper,
2124                            Arc::clone(&persist_clients),
2125                            lookahead,
2126                            snapshot_time,
2127                        );
2128
2129                    let stats_defs =
2130                        SourceStatisticsMetricDefs::register_with(&MetricsRegistry::new());
2131                    let source_statistics = SourceStatistics::new(
2132                        source_id,
2133                        0,
2134                        &stats_defs,
2135                        source_id,
2136                        &target.data_shard,
2137                        SourceEnvelope::None(NoneEnvelope {
2138                            key_envelope: KeyEnvelope::None,
2139                            key_arity: 0,
2140                        }),
2141                        Antichain::from_elem(Timestamp::minimum()),
2142                        Antichain::from_elem(Timestamp::minimum()),
2143                    );
2144
2145                    let (batches, write_button) = write_batches(
2146                        scope,
2147                        source_id,
2148                        "test",
2149                        &target,
2150                        descriptions,
2151                        ceilings,
2152                        passthrough.as_collection(),
2153                        persist_clients,
2154                        source_statistics,
2155                        Arc::new(Semaphore::new(Semaphore::MAX_PERMITS)),
2156                    );
2157
2158                    let sink = Rc::clone(&emitted);
2159                    Inspect::inspect_core(batches, move |event| {
2160                        if let Ok((_, data)) = event {
2161                            for b in data {
2162                                sink.borrow_mut().push(EmittedBatch {
2163                                    lower: b.lower.as_option().expect("single lower").into(),
2164                                    upper: b.upper.as_option().expect("single upper").into(),
2165                                    data_max_ts: b.data_max_ts.into(),
2166                                    inserts: b.metrics.inserts,
2167                                });
2168                            }
2169                        }
2170                    });
2171
2172                    (data_input, remap_input, [mint_button, write_button])
2173                });
2174
2175            fn pump(worker: &mut timely::worker::Worker) {
2176                for _ in 0..32 {
2177                    worker.step_or_park(Some(Duration::from_millis(1)));
2178                }
2179            }
2180
2181            // Twice, so both operators are past opening their persist handles before the script.
2182            pump(worker);
2183            pump(worker);
2184
2185            for step in script {
2186                match step {
2187                    Step::Description(..) | Step::Commit(..) => {
2188                        panic!("the minter emits these")
2189                    }
2190                    Step::Updates(at, count) => {
2191                        for i in 0..i64::try_from(count).expect("small count") {
2192                            let row = Row::pack_slice(&[Datum::Int64(i)]);
2193                            data_input.send((Ok(row), ts(at), Diff::ONE));
2194                        }
2195                    }
2196                    Step::Retractions(at, count) => {
2197                        for i in 0..i64::try_from(count).expect("small count") {
2198                            let row = Row::pack_slice(&[Datum::Int64(i)]);
2199                            data_input.send((Ok(row), ts(at), -Diff::ONE));
2200                        }
2201                    }
2202                    Step::AdvanceTo(t) => data_input.advance_to(ts(t)),
2203                    Step::Tick(t) => remap_input.advance_to(ts(t)),
2204                }
2205                data_input.flush();
2206                pump(worker);
2207            }
2208
2209            // Closing makes every outstanding description ready. `step_or_park` reports whether
2210            // any dataflow remains, and the operators retire once their inputs are exhausted, so
2211            // stepping until it reports none drains everything they still owe.
2212            data_input.close();
2213            remap_input.close();
2214            for _ in 0..20_000 {
2215                if !worker.step_or_park(Some(Duration::from_millis(1))) {
2216                    break;
2217                }
2218            }
2219
2220            drop(buttons);
2221            while worker.step() {}
2222
2223            let mut emitted = emitted.borrow().clone();
2224            emitted.sort();
2225            emitted
2226        })
2227        .await
2228    }
2229
2230    const PIPELINE_SNAPSHOT_ROWS: u64 = 200;
2231    const PIPELINE_PER_TIME: u64 = 3;
2232    const PIPELINE_PINNED_TIMES: u64 = 40;
2233    const PIPELINE_LOOKAHEAD: u64 = 8;
2234    /// The last data time is `PIPELINE_PINNED_TIMES + 1`, and the tick that stamps it puts the
2235    /// remap upper one past that, so the last ceiling ends a lookahead past the upper. The frontier
2236    /// has to reach it for the one description to be minted, which is the tail the lookahead costs.
2237    const PIPELINE_DONE: u64 = PIPELINE_PINNED_TIMES + 2 + PIPELINE_LOOKAHEAD;
2238
2239    /// Ticks the remap upper past `t`, as the binding that stamps rows at `t` does.
2240    fn tick_past(t: u64) -> Step {
2241        Step::Tick(t + 1)
2242    }
2243
2244    /// A snapshot at the pinned time, then thin replication data at each later time, each behind
2245    /// the tick that stamps it, then the frontier advance that ends the snapshot.
2246    fn pinned_snapshot_script() -> Vec<Step> {
2247        let mut script = vec![
2248            tick_past(1),
2249            Step::AdvanceTo(1),
2250            Step::Updates(1, usize::cast_from(PIPELINE_SNAPSHOT_ROWS)),
2251        ];
2252        for t in 2..=PIPELINE_PINNED_TIMES + 1 {
2253            script.push(tick_past(t));
2254            script.push(Step::Updates(t, usize::cast_from(PIPELINE_PER_TIME)));
2255        }
2256        script.push(Step::AdvanceTo(PIPELINE_DONE));
2257        script
2258    }
2259
2260    /// The same rows as `pinned_snapshot_script`, but the snapshot's rows keep landing at the
2261    /// pinned time for as long as the snapshot runs, interleaved with the replication rows at
2262    /// later times.
2263    fn interleaved_snapshot_script() -> Vec<Step> {
2264        let per_round = usize::cast_from(PIPELINE_SNAPSHOT_ROWS / PIPELINE_PINNED_TIMES);
2265        let mut script = vec![tick_past(1), Step::AdvanceTo(1)];
2266        for t in 2..=PIPELINE_PINNED_TIMES + 1 {
2267            script.push(Step::Updates(1, per_round));
2268            script.push(tick_past(t));
2269            script.push(Step::Updates(t, usize::cast_from(PIPELINE_PER_TIME)));
2270        }
2271        script.push(Step::AdvanceTo(PIPELINE_DONE));
2272        script
2273    }
2274
2275    /// Either pinned script writes one batch, under the one description the minter emits when the
2276    /// frontier reaches its ceiling. Every row is in it, whatever its timestamp.
2277    fn one_batch() -> Vec<EmittedBatch> {
2278        vec![EmittedBatch {
2279            lower: 1,
2280            upper: PIPELINE_DONE,
2281            data_max_ts: PIPELINE_PINNED_TIMES + 1,
2282            inserts: PIPELINE_SNAPSHOT_ROWS + PIPELINE_PINNED_TIMES * PIPELINE_PER_TIME,
2283        }]
2284    }
2285
2286    /// The ceiling reaches the writer ahead of the updates it covers, so a pinned frontier groups
2287    /// every timestamp it accumulates into one batch.
2288    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2289    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2290    async fn sink_pipeline_groups_a_pinned_frontier_into_one_batch() {
2291        let emitted = run_sink_pipeline(
2292            test_target(),
2293            test_persist_clients(),
2294            Some(PIPELINE_LOOKAHEAD),
2295            Some(frontier(1)),
2296            pinned_snapshot_script(),
2297        )
2298        .await;
2299        assert_eq!(emitted, one_batch());
2300    }
2301
2302    /// The pinned time is the one time that stays incomplete for the whole snapshot, so rows at
2303    /// it arrive after rows at every later time. Grouping must not depend on data leaving a
2304    /// timestamp for good.
2305    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2306    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2307    async fn sink_pipeline_keeps_one_batch_while_snapshot_rows_interleave() {
2308        let emitted = run_sink_pipeline(
2309            test_target(),
2310            test_persist_clients(),
2311            Some(PIPELINE_LOOKAHEAD),
2312            Some(frontier(1)),
2313            interleaved_snapshot_script(),
2314        )
2315        .await;
2316        assert_eq!(emitted, one_batch());
2317    }
2318
2319    /// Without a lookahead the only description a pinned frontier gets is the one its release
2320    /// derives, which arrives behind every update, so each timestamp writes its own batch.
2321    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2322    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2323    async fn sink_pipeline_writes_one_batch_per_timestamp_without_a_lookahead() {
2324        let emitted = run_sink_pipeline(
2325            test_target(),
2326            test_persist_clients(),
2327            None,
2328            Some(frontier(1)),
2329            pinned_snapshot_script(),
2330        )
2331        .await;
2332
2333        assert_eq!(
2334            emitted.len(),
2335            usize::cast_from(PIPELINE_PINNED_TIMES + 1),
2336            "{emitted:?}"
2337        );
2338        assert!(
2339            emitted
2340                .iter()
2341                .all(|b| (b.lower, b.upper) == (1, PIPELINE_DONE)),
2342            "all finished under the one description the frontier's advance derives: {emitted:?}"
2343        );
2344        assert_eq!(
2345            emitted.iter().map(|b| b.inserts).sum::<u64>(),
2346            PIPELINE_SNAPSHOT_ROWS + PIPELINE_PINNED_TIMES * PIPELINE_PER_TIME,
2347        );
2348    }
2349
2350    fn test_target() -> CollectionMetadata {
2351        CollectionMetadata {
2352            persist_location: PersistLocation {
2353                blob_uri: SensitiveUrl::from_str("mem://").expect("invalid URL"),
2354                consensus_uri: SensitiveUrl::from_str("mem://").expect("invalid URL"),
2355            },
2356            data_shard: ShardId::new(),
2357            relation_desc: RelationDesc::builder()
2358                .with_column("a", SqlScalarType::Int64.nullable(false))
2359                .finish(),
2360            txns_shard: None,
2361        }
2362    }
2363
2364    /// Turn on part bounds validation so _append_ checks the bounds the sink writes.
2365    /// Both settings default off in code but are turned on in production.
2366    fn test_persist_clients() -> Arc<PersistClientCache> {
2367        let persist_cfg =
2368            PersistConfig::new_default_configs(&DUMMY_BUILD_INFO, SYSTEM_TIME.clone());
2369        let mut updates = ConfigUpdates::default();
2370        updates.add_dynamic(
2371            "persist_validate_part_bounds_on_write",
2372            ConfigVal::Bool(true),
2373        );
2374        updates.add_dynamic(
2375            "persist_validate_part_bounds_on_read",
2376            ConfigVal::Bool(true),
2377        );
2378        updates.apply(&persist_cfg.configs);
2379        Arc::new(PersistClientCache::new(
2380            persist_cfg,
2381            &MetricsRegistry::new(),
2382            |_, _| PubSubClientConnection::noop(),
2383        ))
2384    }
2385
2386    /// A single `compare_and_append` over `[lower, upper)` carrying every emitted batch.
2387    fn one_append(
2388        emitted: Vec<(EmittedBatch, ProtoBatch)>,
2389        lower: u64,
2390        upper: u64,
2391    ) -> Vec<(u64, u64, Vec<ProtoBatch>)> {
2392        vec![(lower, upper, emitted.into_iter().map(|(_, p)| p).collect())]
2393    }
2394
2395    /// One `compare_and_append` per description the batches were written for, ascending by lower.
2396    fn append_per_description(
2397        emitted: Vec<(EmittedBatch, ProtoBatch)>,
2398    ) -> Vec<(u64, u64, Vec<ProtoBatch>)> {
2399        let mut by_desc: BTreeMap<(u64, u64), Vec<ProtoBatch>> = BTreeMap::new();
2400        for (batch, proto) in emitted {
2401            by_desc
2402                .entry((batch.lower, batch.upper))
2403                .or_default()
2404                .push(proto);
2405        }
2406        by_desc
2407            .into_iter()
2408            .map(|((lower, upper), protos)| (lower, upper, protos))
2409            .collect()
2410    }
2411
2412    /// Applies each entry in `appends` as one `compare_and_append` over `[lower, upper)`, in order,
2413    /// then reads the shard back as of `as_of` and returns the summed diffs.
2414    ///
2415    /// Batches written for different descriptions need separate entries, because persist rejects a
2416    /// batch whose upper is below the append upper. A `lower` above a batch's own lower registers
2417    /// it truncated, which is what the sink relies on when a concurrent writer has already claimed
2418    /// part of the range.
2419    ///
2420    /// Part bounds validation is what catches a batch whose parts reach outside their registered
2421    /// bounds, so the tests append for real rather than stopping at what `write_batches` emitted.
2422    async fn append_and_read_back(
2423        target: &CollectionMetadata,
2424        persist_clients: &PersistClientCache,
2425        appends: Vec<(u64, u64, Vec<ProtoBatch>)>,
2426        as_of: u64,
2427    ) -> i64 {
2428        let persist_client = persist_clients
2429            .open(target.persist_location.clone())
2430            .await
2431            .expect("could not open persist client");
2432        let mut write = persist_client
2433            .open_writer::<SourceData, (), mz_repr::Timestamp, StorageDiff>(
2434                target.data_shard,
2435                Arc::new(target.relation_desc.clone()),
2436                Arc::new(UnitSchema),
2437                Diagnostics::for_tests(),
2438            )
2439            .await
2440            .expect("could not open persist shard");
2441
2442        assert!(
2443            write.validate_part_bounds_on_write(),
2444            "part bounds validation is off, so this append proves nothing about batch bounds"
2445        );
2446
2447        for (lower, upper, protos) in appends {
2448            let mut batches: Vec<_> = protos
2449                .into_iter()
2450                .map(|proto| write.batch_from_transmittable_batch(proto))
2451                .collect();
2452            let mut to_append: Vec<_> = batches.iter_mut().collect();
2453            write
2454                .compare_and_append_batch(
2455                    &mut to_append[..],
2456                    frontier(lower),
2457                    frontier(upper),
2458                    true,
2459                )
2460                .await
2461                .expect("invalid usage")
2462                .expect("upper mismatch");
2463
2464            assert_eq!(write.fetch_recent_upper().await, &frontier(upper));
2465        }
2466
2467        let mut read = persist_client
2468            .open_leased_reader::<SourceData, (), mz_repr::Timestamp, StorageDiff>(
2469                target.data_shard,
2470                Arc::new(target.relation_desc.clone()),
2471                Arc::new(UnitSchema),
2472                Diagnostics::for_tests(),
2473                true,
2474            )
2475            .await
2476            .expect("invalid usage");
2477        let contents = read
2478            .snapshot_and_fetch(frontier(as_of))
2479            .await
2480            .expect("since <= as_of");
2481
2482        contents.iter().map(|(_, _, d)| *d).sum()
2483    }
2484
2485    /// Several descriptions can become ready in the same pass. Each is written under its own
2486    /// bounds, so every batch holds exactly the updates its description covers however that ready
2487    /// set happens to be ordered.
2488    ///
2489    /// NOTE: `in_flight_batches` is a `HashMap`, so the ready set comes out in no particular order.
2490    /// Enough descriptions are used here that an all-ascending pass is unlikely.
2491    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2492    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2493    async fn write_batches_handles_descriptions_ready_in_one_pass() {
2494        const DESCRIPTIONS: u64 = 6;
2495        const DONE: u64 = DESCRIPTIONS * 2;
2496
2497        let persist_clients = test_persist_clients();
2498        let target = test_target();
2499
2500        // One update inside each of the tiling descriptions [0,2), [2,4), ... None of them is ready
2501        // until the frontier passes every upper, so they all come due together.
2502        let mut script = vec![];
2503        for i in 0..DESCRIPTIONS {
2504            script.push(Step::Updates(i * 2 + 1, 1));
2505        }
2506        for i in 0..DESCRIPTIONS {
2507            script.push(Step::Description(i * 2, i * 2 + 2));
2508        }
2509        script.push(Step::AdvanceTo(DONE));
2510
2511        let emitted = run_write_batches(target.clone(), Arc::clone(&persist_clients), script).await;
2512
2513        assert_eq!(
2514            emitted.len(),
2515            usize::cast_from(DESCRIPTIONS),
2516            "one batch per description, got {:?}",
2517            emitted.iter().map(|(b, _)| b).collect::<Vec<_>>()
2518        );
2519        for (batch, _) in &emitted {
2520            assert!(
2521                batch.lower <= batch.data_max_ts && batch.data_max_ts < batch.upper,
2522                "batch {batch:?} holds data outside the description it was written for"
2523            );
2524        }
2525
2526        let total = append_and_read_back(
2527            &target,
2528            &persist_clients,
2529            append_per_description(emitted),
2530            DONE - 1,
2531        )
2532        .await;
2533        assert_eq!(
2534            total,
2535            i64::try_from(DESCRIPTIONS).expect("small"),
2536            "every update should be readable exactly once"
2537        );
2538    }
2539
2540    /// Advances the shard upper to `upper` without writing data, standing in for a concurrent
2541    /// writer that reached part of the range first.
2542    async fn advance_shard_upper(
2543        target: &CollectionMetadata,
2544        persist_clients: &PersistClientCache,
2545        upper: u64,
2546    ) {
2547        let persist_client = persist_clients
2548            .open(target.persist_location.clone())
2549            .await
2550            .expect("could not open persist client");
2551        let mut write = persist_client
2552            .open_writer::<SourceData, (), mz_repr::Timestamp, StorageDiff>(
2553                target.data_shard,
2554                Arc::new(target.relation_desc.clone()),
2555                Arc::new(UnitSchema),
2556                Diagnostics::for_tests(),
2557            )
2558            .await
2559            .expect("could not open persist shard");
2560
2561        let empty: Vec<((SourceData, ()), mz_repr::Timestamp, StorageDiff)> = Vec::new();
2562        write
2563            .compare_and_append(
2564                &empty,
2565                Antichain::from_elem(Timestamp::minimum()),
2566                frontier(upper),
2567            )
2568            .await
2569            .expect("invalid usage")
2570            .expect("upper mismatch");
2571
2572        // Otherwise the handle's heartbeat task outlives the test and nextest reports a leak.
2573        write.expire().await;
2574    }
2575
2576    /// A batch spanning many timestamps stays usable when a concurrent writer has raised the shard
2577    /// upper into the middle of it. Persist registers the batch under the narrowed description and
2578    /// filters the updates outside those bounds on read, so the sink can hand a straddling batch
2579    /// over as is rather than discarding it and rebuilding from the new upper.
2580    ///
2581    /// This is the property that lets `write_batches` coalesce a pinned frontier into one batch
2582    /// without giving up the ability to recover from a concurrent append.
2583    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2584    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2585    async fn a_straddling_batch_is_usable_under_a_raised_lower() {
2586        const TIMES: u64 = 10;
2587        const DONE: u64 = TIMES + 1;
2588        // Inside the batch's data range, so the batch holds updates on both sides of it.
2589        const RAISED_LOWER: u64 = 5;
2590
2591        let persist_clients = test_persist_clients();
2592        let target = test_target();
2593
2594        // One update at each of times 1..=TIMES, all going into the builder of a description
2595        // already in hand, so they coalesce into one batch whose data spans the whole range.
2596        let emitted = run_write_batches(
2597            target.clone(),
2598            Arc::clone(&persist_clients),
2599            committed_ceiling_script(1, TIMES - 1, DONE),
2600        )
2601        .await;
2602        assert_eq!(
2603            emitted.len(),
2604            1,
2605            "expected one coalesced batch, got {:?}",
2606            emitted.iter().map(|(b, _)| b).collect::<Vec<_>>()
2607        );
2608
2609        advance_shard_upper(&target, &persist_clients, RAISED_LOWER).await;
2610
2611        // Append the straddling batch under the raised lower, as the sink does after an
2612        // `UpperMismatch` cuts the description down.
2613        let total = append_and_read_back(
2614            &target,
2615            &persist_clients,
2616            one_append(emitted, RAISED_LOWER, DONE),
2617            DONE - 1,
2618        )
2619        .await;
2620
2621        assert_eq!(
2622            total,
2623            i64::try_from(TIMES - RAISED_LOWER + 1).expect("small"),
2624            "only the updates at or above the raised lower should be readable, and all of them"
2625        );
2626    }
2627
2628    /// The rewind mechanism retracts the snapshot's copy of every row the replication stream
2629    /// redelivers at its true offset, and both land at the pinned timestamp. They share the
2630    /// covering description's builder, so the batch carries both signs and they net out on read.
2631    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2632    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2633    async fn write_batches_carries_rewind_retractions_into_the_covering_batch() {
2634        const ROWS: usize = 512;
2635        const DONE: u64 = 4;
2636
2637        let persist_clients = test_persist_clients();
2638        let target = test_target();
2639
2640        let emitted = run_write_batches(
2641            target.clone(),
2642            Arc::clone(&persist_clients),
2643            vec![
2644                Step::Commit(0, DONE),
2645                Step::Description(0, DONE),
2646                Step::AdvanceTo(1),
2647                Step::Updates(1, ROWS),
2648                Step::Retractions(1, ROWS - 1),
2649                Step::AdvanceTo(DONE),
2650            ],
2651        )
2652        .await;
2653
2654        assert_eq!(
2655            emitted.len(),
2656            1,
2657            "both signs belong to the one description, got {:?}",
2658            emitted.iter().map(|(b, _)| b).collect::<Vec<_>>()
2659        );
2660
2661        let total = append_and_read_back(
2662            &target,
2663            &persist_clients,
2664            one_append(emitted, 0, DONE),
2665            DONE - 1,
2666        )
2667        .await;
2668        assert_eq!(total, 1, "the shard should hold exactly the surviving row");
2669    }
2670
2671    /// A description that covers no updates must emit no batch, rather than open a builder that
2672    /// has no data bounds to register.
2673    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2674    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2675    async fn write_batches_emits_nothing_for_a_description_with_no_updates() {
2676        const SPLIT: u64 = 4;
2677        const DONE: u64 = 8;
2678
2679        // Two descriptions in hand, with data only in the second.
2680        let emitted = run_write_batches(
2681            test_target(),
2682            test_persist_clients(),
2683            vec![
2684                Step::Description(0, SPLIT),
2685                Step::Description(SPLIT, DONE),
2686                Step::AdvanceTo(SPLIT),
2687                Step::Updates(SPLIT, 8),
2688                Step::AdvanceTo(DONE),
2689            ],
2690        )
2691        .await;
2692
2693        assert_eq!(
2694            emitted
2695                .iter()
2696                .map(|(b, _)| (b.lower, b.upper))
2697                .collect::<Vec<_>>(),
2698            vec![(SPLIT, DONE)],
2699            "only the description holding data should produce a batch",
2700        );
2701    }
2702
2703    /// A snapshot at time 1 pinning the frontier while replication delivers one update at each of
2704    /// times 2..=`pinned_times`+1, with the description that covers the whole snapshot arriving
2705    /// only at the end. This is the sink without a budget: nothing is committed ahead of the
2706    /// frontier, so no update ever has bounds to be grouped under while it is staged.
2707    fn pinned_frontier_script(snapshot_rows: usize, pinned_times: u64, done: u64) -> Vec<Step> {
2708        let mut script = vec![Step::Updates(1, snapshot_rows)];
2709        for t in 2..=pinned_times + 1 {
2710            script.push(Step::Updates(t, 1));
2711        }
2712        // The minter holds a capability at the shard upper for the whole snapshot, so its one
2713        // description is emitted there, and the frontier then jumps past everything staged.
2714        script.push(Step::Description(0, done));
2715        script.push(Step::AdvanceTo(done));
2716        script
2717    }
2718
2719    /// The same snapshot with a ceiling committed first, which is what the minter does behind a
2720    /// frontier that is not moving. The description itself only arrives once the frontier reaches
2721    /// the ceiling, which is what ends the script.
2722    fn committed_ceiling_script(snapshot_rows: usize, pinned_times: u64, done: u64) -> Vec<Step> {
2723        let mut script = vec![
2724            Step::Commit(0, done),
2725            Step::AdvanceTo(1),
2726            Step::Updates(1, snapshot_rows),
2727        ];
2728        for t in 2..=pinned_times + 1 {
2729            script.push(Step::Updates(t, 1));
2730        }
2731        script.push(Step::Description(0, done));
2732        script.push(Step::AdvanceTo(done));
2733        script
2734    }
2735
2736    /// A snapshot pins the export's frontier at its as_of while concurrent replication keeps
2737    /// delivering updates at later times. Without a description in hand there are no bounds to
2738    /// group under, so each timestamp writes a batch of its own, all finished under the one
2739    /// description that arrives when the snapshot finishes.
2740    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2741    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2742    async fn write_batches_writes_one_batch_per_timestamp_before_a_description_covers_them() {
2743        const SNAPSHOT_ROWS: usize = 4;
2744        const PINNED_TIMES: u64 = 16;
2745        const DONE: u64 = PINNED_TIMES + 2;
2746
2747        let persist_clients = test_persist_clients();
2748        let target = test_target();
2749
2750        let emitted = run_write_batches(
2751            target.clone(),
2752            Arc::clone(&persist_clients),
2753            pinned_frontier_script(SNAPSHOT_ROWS, PINNED_TIMES, DONE),
2754        )
2755        .await;
2756
2757        // Every batch carries the description's bounds, since that is what they are finished
2758        // under, and holds a single timestamp's updates. One timestamp per batch is what makes
2759        // writing before the covering description is known safe, since a single-timestamp batch
2760        // cannot span a description boundary.
2761        let expected: Vec<_> = std::iter::once(EmittedBatch {
2762            lower: 0,
2763            upper: DONE,
2764            data_max_ts: 1,
2765            inserts: u64::cast_from(SNAPSHOT_ROWS),
2766        })
2767        .chain((2..=PINNED_TIMES + 1).map(|ts| EmittedBatch {
2768            lower: 0,
2769            upper: DONE,
2770            data_max_ts: ts,
2771            inserts: 1,
2772        }))
2773        .collect();
2774        assert_eq!(
2775            emitted.iter().map(|(b, _)| b.clone()).collect::<Vec<_>>(),
2776            expected,
2777        );
2778
2779        let total = append_and_read_back(
2780            &target,
2781            &persist_clients,
2782            one_append(emitted, 0, DONE),
2783            DONE - 1,
2784        )
2785        .await;
2786        assert_eq!(
2787            total,
2788            i64::try_from(SNAPSHOT_ROWS).expect("small")
2789                + i64::try_from(PINNED_TIMES).expect("small"),
2790            "the same updates should be readable however they were batched"
2791        );
2792    }
2793
2794    /// The writer drains descriptions ahead of data in each round, and input keeps queueing while
2795    /// it awaits persist, so rows at a timestamp can be staged after the description covering them
2796    /// was processed. They belong in the batch that timestamp already has.
2797    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2798    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2799    async fn write_batches_keeps_one_batch_per_timestamp_for_rows_after_the_description() {
2800        const ROWS: usize = 4;
2801        const DONE: u64 = 4;
2802
2803        let persist_clients = test_persist_clients();
2804        let target = test_target();
2805
2806        let emitted = run_write_batches(
2807            target.clone(),
2808            Arc::clone(&persist_clients),
2809            vec![
2810                Step::Updates(1, ROWS),
2811                Step::Description(0, DONE),
2812                Step::Updates(1, ROWS),
2813                Step::AdvanceTo(DONE),
2814            ],
2815        )
2816        .await;
2817
2818        assert_eq!(
2819            emitted.iter().map(|(b, _)| b.clone()).collect::<Vec<_>>(),
2820            vec![EmittedBatch {
2821                lower: 0,
2822                upper: DONE,
2823                data_max_ts: 1,
2824                inserts: u64::cast_from(ROWS * 2),
2825            }],
2826        );
2827
2828        let total = append_and_read_back(
2829            &target,
2830            &persist_clients,
2831            one_append(emitted, 0, DONE),
2832            DONE - 1,
2833        )
2834        .await;
2835        assert_eq!(total, i64::try_from(ROWS * 2).expect("small"));
2836    }
2837
2838    /// The remap upper is broadcast, so the minting worker paces on its own copy of it without a
2839    /// row on any worker: on the pin a lookahead past the upper, then past each tick.
2840    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2841    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2842    async fn mint_batch_descriptions_paces_on_the_remap_upper_without_data() {
2843        const PEERS: usize = 2;
2844
2845        let source_id = GlobalId::User(0);
2846        let target = test_target();
2847        let persist_clients = test_persist_clients();
2848        let emitted = Arc::new(std::sync::Mutex::new(Vec::new()));
2849        // `timely::execute` runs each worker on a thread of its own, which is outside the test's
2850        // runtime, and the operators drive persist there.
2851        let runtime = tokio::runtime::Handle::current();
2852
2853        let collected = Arc::clone(&emitted);
2854        mz_ore::task::spawn_blocking(
2855            || "persist_sink_test_workers",
2856            move || {
2857                timely::execute::execute(timely::Config::process(PEERS), move |worker| {
2858                    let _guard = runtime.enter();
2859                    let emitted = Arc::clone(&emitted);
2860                    let target = target.clone();
2861                    let persist_clients = Arc::clone(&persist_clients);
2862
2863                    let (mut data_input, mut remap_input, button) = worker
2864                        .dataflow::<mz_repr::Timestamp, _, _>(|scope| {
2865                            let (data_input, data) = scope.new_input();
2866                            let (remap_input, remap_upper) = scope.new_input::<Vec<()>>();
2867                            let (_descriptions, ceilings, _passthrough, button) =
2868                                mint_batch_descriptions(
2869                                    scope,
2870                                    source_id,
2871                                    "test",
2872                                    &target,
2873                                    data.as_collection(),
2874                                    remap_upper,
2875                                    persist_clients,
2876                                    Some(4),
2877                                    Some(frontier(1)),
2878                                );
2879
2880                            Inspect::inspect_core(ceilings, move |event| {
2881                                if let Ok((_, data)) = event {
2882                                    let mut emitted = emitted.lock().expect("not poisoned");
2883                                    for commitment in data {
2884                                        emitted.push(u64::from(commitment.ceiling));
2885                                    }
2886                                }
2887                            });
2888
2889                            (data_input, remap_input, button)
2890                        });
2891
2892                    fn pump(worker: &mut timely::worker::Worker) {
2893                        for _ in 0..64 {
2894                            worker.step_or_park(Some(Duration::from_millis(1)));
2895                        }
2896                    }
2897
2898                    pump(worker);
2899                    pump(worker);
2900
2901                    // Both frontiers are global, so every worker advances both. The data frontier
2902                    // pins at 1 for the rest of the script, and no worker ever sends a row.
2903                    remap_input.advance_to(ts(2));
2904                    remap_input.flush();
2905                    pump(worker);
2906                    data_input.advance_to(ts(1));
2907                    data_input.flush();
2908                    pump(worker);
2909
2910                    for at in [5, 9, 11] {
2911                        remap_input.advance_to(ts(at));
2912                        remap_input.flush();
2913                        pump(worker);
2914                    }
2915
2916                    data_input.close();
2917                    remap_input.close();
2918                    for _ in 0..2_000 {
2919                        if !worker.step_or_park(Some(Duration::from_millis(1))) {
2920                            break;
2921                        }
2922                    }
2923
2924                    drop(button);
2925                    while worker.step() {}
2926                })
2927                .expect("timely configuration")
2928                .join()
2929                .into_iter()
2930                .collect::<Result<Vec<_>, _>>()
2931                .expect("timely worker panicked");
2932            },
2933        )
2934        .await;
2935
2936        let collected = collected.lock().expect("not poisoned").clone();
2937        assert_eq!(collected, vec![6, 9, 13, 15]);
2938    }
2939
2940    #[mz_ore::test]
2941    fn description_lookahead_disables_at_zero_and_floors_at_the_timestamp_interval() {
2942        let secs = Duration::from_secs;
2943
2944        assert_eq!(
2945            description_lookahead(Duration::ZERO, secs(1)),
2946            None,
2947            "zero is what turns committing ahead off"
2948        );
2949        // A lookahead shorter than the step the remap upper moves in can be outrun by the next
2950        // tick's rows.
2951        assert_eq!(
2952            description_lookahead(Duration::from_millis(1), Duration::from_millis(250)),
2953            Some(250),
2954            "a fine interval floors low, so a source that ticks fast is not forced to over-commit"
2955        );
2956        assert_eq!(
2957            description_lookahead(Duration::from_millis(1), secs(10)),
2958            Some(10_000),
2959            "a coarse interval floors high, or the next tick's rows clear the ceiling"
2960        );
2961        assert_eq!(
2962            description_lookahead(secs(5), Duration::from_millis(250)),
2963            Some(5_000),
2964            "a configured lookahead above the floor stands"
2965        );
2966        assert_eq!(
2967            description_lookahead(Duration::from_micros(500), Duration::from_millis(250)),
2968            Some(250),
2969            "a sub-ms lookahead is nonzero, as the sources that enable concurrent replication read it"
2970        );
2971        assert_eq!(
2972            description_lookahead(Duration::from_micros(500), Duration::from_micros(500)),
2973            Some(1),
2974            "a sub-ms lookahead and interval still commit ahead"
2975        );
2976    }
2977
2978    /// The ceiling is committed the moment the frontier pins, a lookahead past the remap upper,
2979    /// and re-committed as the upper ticks, all under the one description's lower. The rows do not
2980    /// enter into it. Without a lookahead a pinned frontier commits nothing.
2981    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
2982    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
2983    async fn mint_batch_descriptions_commits_a_ceiling_ahead_of_the_remap_upper() {
2984        let script = vec![
2985            Step::Tick(2),
2986            // Data arrives at times the frontier never reaches, which is a pinned frontier.
2987            Step::AdvanceTo(1),
2988            Step::Tick(5),
2989            Step::Updates(4, 4),
2990            Step::Tick(9),
2991            Step::Updates(8, 4),
2992            Step::Tick(11),
2993        ];
2994
2995        assert_eq!(
2996            run_mint_batch_descriptions(
2997                test_target(),
2998                test_persist_clients(),
2999                Some(4),
3000                Some(frontier(1)),
3001                script.clone(),
3002            )
3003            .await,
3004            Minted {
3005                // Only the one the first frontier advance derives. The description spanning the
3006                // snapshot needs the frontier to reach the ceiling, which this script never does.
3007                descriptions: vec![(0, 1)],
3008                // On the pin, then a lookahead past each tick.
3009                commitments: vec![(1, 6), (1, 9), (1, 13), (1, 15)],
3010            },
3011        );
3012
3013        assert_eq!(
3014            run_mint_batch_descriptions(
3015                test_target(),
3016                test_persist_clients(),
3017                None,
3018                Some(frontier(1)),
3019                script
3020            )
3021            .await,
3022            Minted {
3023                descriptions: vec![(0, 1)],
3024                commitments: vec![],
3025            },
3026            "without a lookahead a pinned frontier commits nothing",
3027        );
3028    }
3029
3030    /// Commitments are made only while the frontier sits at the snapshot time. Once the frontier
3031    /// moves, no more commitments are made, and the outstanding commitment is honored until the
3032    /// frontier reaches it.
3033    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
3034    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
3035    async fn mint_batch_descriptions_honors_its_commitment_once_the_frontier_moves_again() {
3036        let emitted = run_mint_batch_descriptions(
3037            test_target(),
3038            test_persist_clients(),
3039            Some(4),
3040            Some(frontier(1)),
3041            vec![
3042                Step::Tick(2),
3043                Step::AdvanceTo(1),
3044                // Below the ceiling committed on the pin, so this mints nothing.
3045                Step::AdvanceTo(2),
3046                // The snapshot is over, so this commits nothing either.
3047                Step::Tick(9),
3048                Step::AdvanceTo(6),
3049            ],
3050        )
3051        .await;
3052
3053        assert_eq!(
3054            emitted,
3055            Minted {
3056                descriptions: vec![(0, 1), (1, 6)],
3057                commitments: vec![(1, 6)],
3058            },
3059            "past hydration the minter should wait for the frontier to pass its ceiling and \
3060            derive from the frontier alone after that",
3061        );
3062    }
3063
3064    /// An export that does not snapshot in this incarnation is given no snapshot time and never
3065    /// commits ahead of its frontier, so a restart of a caught up source has no window where it
3066    /// might.
3067    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
3068    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
3069    async fn mint_batch_descriptions_does_not_commit_without_a_snapshot_time() {
3070        // The shape of a restart: the frontier leaves the minimum once, and the remap upper runs
3071        // past it before it moves again.
3072        let script = vec![Step::Tick(2), Step::AdvanceTo(1), Step::Tick(9)];
3073
3074        assert_eq!(
3075            run_mint_batch_descriptions(
3076                test_target(),
3077                test_persist_clients(),
3078                Some(4),
3079                None,
3080                script.clone(),
3081            )
3082            .await,
3083            Minted {
3084                descriptions: vec![(0, 1)],
3085                commitments: vec![],
3086            },
3087        );
3088
3089        assert_eq!(
3090            run_mint_batch_descriptions(
3091                test_target(),
3092                test_persist_clients(),
3093                Some(4),
3094                Some(frontier(1)),
3095                script,
3096            )
3097            .await,
3098            Minted {
3099                descriptions: vec![(0, 1)],
3100                commitments: vec![(1, 6), (1, 13)],
3101            },
3102            "the same script commits when the export snapshots at 1",
3103        );
3104    }
3105
3106    /// A source may send rows at the snapshot time under a delayed capability before it has
3107    /// downgraded its own, so rows can reach the minter while the frontier still sits at the
3108    /// minimum, with the remap upper ticking all the while. Nothing is committed until the frontier
3109    /// reaches the snapshot time, so the commitment's lower is the snapshot time and not the shard
3110    /// upper the rows were seen under.
3111    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
3112    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
3113    async fn mint_batch_descriptions_waits_for_the_frontier_to_reach_the_snapshot_time() {
3114        let emitted = run_mint_batch_descriptions(
3115            test_target(),
3116            test_persist_clients(),
3117            Some(4),
3118            Some(frontier(1)),
3119            vec![
3120                Step::Tick(3),
3121                Step::Updates(1, 1),
3122                Step::AdvanceTo(1),
3123                Step::Tick(4),
3124                Step::AdvanceTo(8),
3125            ],
3126        )
3127        .await;
3128        assert_eq!(
3129            emitted,
3130            Minted {
3131                descriptions: vec![(0, 1), (1, 8)],
3132                // On the pin a lookahead past the upper, then re-committed past the tick.
3133                commitments: vec![(1, 7), (1, 8)],
3134            },
3135        );
3136    }
3137
3138    /// A fresh source's snapshot lands at the minimum and its frontier never leaves it during the
3139    /// snapshot, so no progress statement announces the pin. The first ceiling goes out on the
3140    /// first activation, ahead of every row and before the remap upper is known, and is raised as
3141    /// the upper ticks.
3142    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
3143    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
3144    async fn mint_batch_descriptions_commits_on_the_first_activation_at_a_minimum_snapshot_time() {
3145        // Replication rows land first, the snapshot's own rows at the pinned time arrive behind
3146        // them, and the frontier never moves.
3147        let script = vec![
3148            Step::Tick(6),
3149            Step::Updates(5, 4),
3150            Step::Updates(0, 4),
3151            Step::Tick(10),
3152            Step::Updates(9, 2),
3153        ];
3154
3155        assert_eq!(
3156            run_mint_batch_descriptions(
3157                test_target(),
3158                test_persist_clients(),
3159                Some(4),
3160                Some(frontier(0)),
3161                script.clone(),
3162            )
3163            .await,
3164            Minted {
3165                // The frontier never moves, so nothing is derived from it.
3166                descriptions: vec![],
3167                // A lookahead past the minimum before anything arrives, then past each tick. The
3168                // rows sit below the outstanding ceiling and re-commit nothing.
3169                commitments: vec![(0, 4), (0, 10), (0, 14)],
3170            },
3171        );
3172
3173        assert_eq!(
3174            run_mint_batch_descriptions(
3175                test_target(),
3176                test_persist_clients(),
3177                Some(4),
3178                Some(frontier(1)),
3179                script,
3180            )
3181            .await,
3182            Minted::default(),
3183            "a snapshot time the frontier has not reached commits nothing, however far the remap \
3184            upper runs",
3185        );
3186    }
3187
3188    #[mz_ore::test]
3189    fn next_mint_derives_from_the_frontier_or_commits_a_ceiling_ahead() {
3190        // Everything below the frontier has arrived, so that is the description, lookahead or not.
3191        assert_eq!(
3192            next_mint(&frontier(5), &frontier(7), &frontier(100), None, Some(8)),
3193            Some(Mint::Description(frontier(7)))
3194        );
3195        assert_eq!(
3196            next_mint(&frontier(5), &frontier(7), &frontier(0), None, None),
3197            Some(Mint::Description(frontier(7)))
3198        );
3199
3200        // A pinned frontier with no lookahead has nothing to commit.
3201        assert_eq!(
3202            next_mint(&frontier(5), &frontier(5), &frontier(14), None, None),
3203            None
3204        );
3205        // Pinned with nothing committed yet: the ceiling goes out a lookahead past the remap
3206        // upper.
3207        assert_eq!(
3208            next_mint(&frontier(5), &frontier(5), &frontier(6), None, Some(8)),
3209            Some(Mint::Ceiling(ts(14)))
3210        );
3211        // A ceiling already a lookahead past the upper stands, so the minter does not re-commit
3212        // every round.
3213        assert_eq!(
3214            next_mint(
3215                &frontier(5),
3216                &frontier(5),
3217                &frontier(6),
3218                Some(ts(14)),
3219                Some(8)
3220            ),
3221            None
3222        );
3223        // An upper that ticked re-commits, a lookahead past where it reached.
3224        assert_eq!(
3225            next_mint(
3226                &frontier(5),
3227                &frontier(5),
3228                &frontier(7),
3229                Some(ts(14)),
3230                Some(8)
3231            ),
3232            Some(Mint::Ceiling(ts(15)))
3233        );
3234        // An upper so far behind the shard upper that the ceiling would not clear it, as before
3235        // the first remap progress event, commits nothing.
3236        assert_eq!(
3237            next_mint(&frontier(5), &frontier(5), &frontier(0), None, Some(4)),
3238            None
3239        );
3240        // A closed remap stream has no upper to commit past.
3241        assert_eq!(
3242            next_mint(
3243                &frontier(5),
3244                &frontier(5),
3245                &Antichain::new(),
3246                Some(ts(14)),
3247                Some(8)
3248            ),
3249            None
3250        );
3251
3252        // An outstanding ceiling the frontier has not reached suppresses the description that
3253        // would otherwise be derived from it, which is what binds the commitment.
3254        assert_eq!(
3255            next_mint(&frontier(5), &frontier(9), &frontier(9), Some(ts(20)), None),
3256            None
3257        );
3258        // Once the frontier reaches it, the one description covering the whole ceiling goes out.
3259        assert_eq!(
3260            next_mint(
3261                &frontier(5),
3262                &frontier(20),
3263                &frontier(20),
3264                Some(ts(20)),
3265                None
3266            ),
3267            Some(Mint::Description(frontier(20)))
3268        );
3269
3270        // A fresh shard's minimum upper pins like any other time, and before the remap upper is
3271        // known the ceiling goes out a lookahead past the minimum.
3272        assert_eq!(
3273            next_mint(&frontier(0), &frontier(0), &frontier(0), None, Some(8)),
3274            Some(Mint::Ceiling(ts(8)))
3275        );
3276        // A collection that is done has no next description.
3277        assert_eq!(
3278            next_mint(
3279                &Antichain::new(),
3280                &frontier(7),
3281                &frontier(100),
3282                None,
3283                Some(8)
3284            ),
3285            None
3286        );
3287    }
3288
3289    /// A ceiling committed ahead of the frontier gives arriving updates a bound, so a pinned
3290    /// frontier writes one batch rather than one per timestamp, and the rows sit in a builder that
3291    /// fills to the blob target instead of many single-timestamp builders that each stay under it
3292    /// and hold their rows in memory.
3293    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
3294    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
3295    async fn write_batches_routes_updates_below_the_ceiling_into_one_builder() {
3296        const SNAPSHOT_ROWS: usize = 4;
3297        const PINNED_TIMES: u64 = 16;
3298        const DONE: u64 = PINNED_TIMES + 2;
3299
3300        let persist_clients = test_persist_clients();
3301        let target = test_target();
3302
3303        let emitted = run_write_batches(
3304            target.clone(),
3305            Arc::clone(&persist_clients),
3306            committed_ceiling_script(SNAPSHOT_ROWS, PINNED_TIMES, DONE),
3307        )
3308        .await;
3309
3310        assert_eq!(
3311            emitted.len(),
3312            1,
3313            "the whole snapshot should share the one open builder, got {:?}",
3314            emitted.iter().map(|(b, _)| b).collect::<Vec<_>>()
3315        );
3316        assert_eq!(
3317            emitted[0].0,
3318            EmittedBatch {
3319                lower: 0,
3320                upper: DONE,
3321                data_max_ts: PINNED_TIMES + 1,
3322                inserts: u64::cast_from(SNAPSHOT_ROWS) + PINNED_TIMES,
3323            }
3324        );
3325
3326        let total = append_and_read_back(
3327            &target,
3328            &persist_clients,
3329            one_append(emitted, 0, DONE),
3330            DONE - 1,
3331        )
3332        .await;
3333        assert_eq!(
3334            total,
3335            i64::try_from(SNAPSHOT_ROWS).expect("small")
3336                + i64::try_from(PINNED_TIMES).expect("small"),
3337            "grouping must not change what the shard ends up holding"
3338        );
3339    }
3340
3341    /// Rows staged before a ceiling covers their timestamp stay in that timestamp's own builder.
3342    /// Rows arriving at it once a ceiling covers it join the open builder instead, and both are
3343    /// finished under the description once it is ready.
3344    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
3345    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
3346    async fn write_batches_keeps_an_uncovered_batch_apart_from_the_open_builder() {
3347        const PINNED: u64 = 2;
3348        const BEFORE: usize = 4;
3349        const AFTER: usize = 6;
3350        const DONE: u64 = 8;
3351
3352        let persist_clients = test_persist_clients();
3353        let target = test_target();
3354
3355        let script = vec![
3356            Step::Updates(PINNED, BEFORE),
3357            Step::Commit(0, DONE),
3358            Step::Description(0, DONE),
3359            Step::Updates(PINNED, AFTER),
3360            Step::AdvanceTo(DONE),
3361        ];
3362
3363        let emitted = run_write_batches(target.clone(), Arc::clone(&persist_clients), script).await;
3364
3365        assert_eq!(
3366            emitted.iter().map(|(b, _)| b.clone()).collect::<Vec<_>>(),
3367            vec![
3368                EmittedBatch {
3369                    lower: 0,
3370                    upper: DONE,
3371                    data_max_ts: PINNED,
3372                    inserts: u64::cast_from(BEFORE),
3373                },
3374                EmittedBatch {
3375                    lower: 0,
3376                    upper: DONE,
3377                    data_max_ts: PINNED,
3378                    inserts: u64::cast_from(AFTER),
3379                },
3380            ],
3381            "rows staged before the ceiling should keep their own batch, and later rows at that \
3382            timestamp should join the open builder"
3383        );
3384
3385        let total = append_and_read_back(
3386            &target,
3387            &persist_clients,
3388            one_append(emitted, 0, DONE),
3389            DONE - 1,
3390        )
3391        .await;
3392        assert_eq!(
3393            total,
3394            i64::try_from(BEFORE + AFTER).expect("small"),
3395            "splitting a timestamp across two batches must not change what the shard ends up holding"
3396        );
3397    }
3398
3399    /// The ceiling is a bound, not a promise about the data: an update that outruns it has nothing
3400    /// to be grouped under and writes a batch of its own timestamp, which costs a batch rather
3401    /// than correctness. The description that arrives later covers both.
3402    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
3403    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait
3404    async fn write_batches_writes_a_trailblazer_past_the_ceiling_its_own_batch() {
3405        const CEILING: u64 = 9;
3406        const PINNED_TIMES: u64 = 16;
3407        const DONE: u64 = PINNED_TIMES + 2;
3408
3409        let persist_clients = test_persist_clients();
3410        let target = test_target();
3411
3412        let mut script = vec![Step::Commit(0, CEILING), Step::AdvanceTo(1)];
3413        for t in 1..=PINNED_TIMES {
3414            script.push(Step::Updates(t, 1));
3415        }
3416        script.push(Step::Description(0, DONE));
3417        script.push(Step::AdvanceTo(DONE));
3418
3419        let emitted = run_write_batches(target.clone(), Arc::clone(&persist_clients), script).await;
3420
3421        assert_eq!(
3422            emitted.len(),
3423            // One for everything below the ceiling, then one per timestamp past it.
3424            usize::cast_from(DONE - CEILING),
3425            "got {:?}",
3426            emitted.iter().map(|(b, _)| b).collect::<Vec<_>>()
3427        );
3428        assert!(
3429            emitted.iter().all(|(b, _)| (b.lower, b.upper) == (0, DONE)),
3430            "all finished under the one description, got {:?}",
3431            emitted.iter().map(|(b, _)| b).collect::<Vec<_>>()
3432        );
3433        assert_eq!(
3434            emitted[0].0,
3435            EmittedBatch {
3436                lower: 0,
3437                upper: DONE,
3438                data_max_ts: CEILING - 1,
3439                inserts: CEILING - 1,
3440            },
3441            "the grouped batch holds every timestamp below the ceiling"
3442        );
3443
3444        let total = append_and_read_back(
3445            &target,
3446            &persist_clients,
3447            one_append(emitted, 0, DONE),
3448            DONE - 1,
3449        )
3450        .await;
3451        assert_eq!(
3452            total,
3453            i64::try_from(PINNED_TIMES).expect("small"),
3454            "outrunning the ceiling must not change what the shard holds"
3455        );
3456    }
3457}