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