Skip to main content

mz_persist_client/
write.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//! Write capabilities and handles
11
12use std::borrow::Borrow;
13use std::fmt::Debug;
14use std::sync::Arc;
15
16use differential_dataflow::difference::Monoid;
17use differential_dataflow::lattice::Lattice;
18use differential_dataflow::trace::Description;
19use futures::StreamExt;
20use futures::stream::FuturesUnordered;
21use mz_dyncfg::{Config, ParameterScope};
22use mz_ore::task::RuntimeExt;
23use mz_ore::{instrument, soft_panic_or_log};
24use mz_persist::location::Blob;
25use mz_persist_types::schema::SchemaId;
26use mz_persist_types::{Codec, Codec64};
27use mz_proto::{IntoRustIfSome, ProtoType};
28use proptest_derive::Arbitrary;
29use semver::Version;
30use serde::{Deserialize, Serialize};
31use timely::PartialOrder;
32use timely::order::TotalOrder;
33use timely::progress::{Antichain, Timestamp};
34use tokio::runtime::Handle;
35use tracing::{Instrument, debug_span, error, info, warn};
36use uuid::Uuid;
37
38use crate::batch::{
39    Added, BATCH_DELETE_ENABLED, Batch, BatchBuilder, BatchBuilderConfig, BatchBuilderInternal,
40    BatchParts, ProtoBatch, validate_truncate_batch,
41};
42use crate::error::{InvalidUsage, UpperMismatch};
43use crate::fetch::{
44    EncodedPart, FetchBatchFilter, FetchedPart, PartDecodeFormat, VALIDATE_PART_BOUNDS_ON_READ,
45};
46use crate::internal::compact::{CompactConfig, Compactor};
47use crate::internal::encoding::{Schemas, assert_code_can_read_data};
48use crate::internal::machine::{
49    CompareAndAppendRes, ExpireFn, Machine, next_listen_batch_retry_params,
50};
51use crate::internal::metrics::{BatchWriteMetrics, Metrics, ShardMetrics};
52use crate::internal::state::{BatchPart, HandleDebugState, HollowBatch, RunOrder, RunPart};
53use crate::read::ReadHandle;
54use crate::schema::PartMigration;
55use crate::{GarbageCollector, IsolatedRuntime, PersistConfig, ShardId, parse_id};
56
57pub(crate) const COMBINE_INLINE_WRITES: Config<bool> = Config::new(
58    "persist_write_combine_inline_writes",
59    true,
60    "If set, re-encode inline writes if they don't fit into the batch metadata limits.",
61    ParameterScope::Environment,
62);
63
64pub(crate) const VALIDATE_PART_BOUNDS_ON_WRITE: Config<bool> = Config::new(
65    "persist_validate_part_bounds_on_write",
66    false,
67    "Validate the part lower <= the batch lower and the part upper <= batch upper,\
68    for the batch being appended.",
69    ParameterScope::Environment,
70);
71
72/// An opaque identifier for a writer of a persist durable TVC (aka shard).
73#[derive(
74    Arbitrary,
75    Clone,
76    PartialEq,
77    Eq,
78    PartialOrd,
79    Ord,
80    Hash,
81    Serialize,
82    Deserialize
83)]
84#[serde(try_from = "String", into = "String")]
85pub struct WriterId(pub(crate) [u8; 16]);
86
87impl std::fmt::Display for WriterId {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        write!(f, "w{}", Uuid::from_bytes(self.0))
90    }
91}
92
93impl std::fmt::Debug for WriterId {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        write!(f, "WriterId({})", Uuid::from_bytes(self.0))
96    }
97}
98
99impl std::str::FromStr for WriterId {
100    type Err = String;
101
102    fn from_str(s: &str) -> Result<Self, Self::Err> {
103        parse_id("w", "WriterId", s).map(WriterId)
104    }
105}
106
107impl From<WriterId> for String {
108    fn from(writer_id: WriterId) -> Self {
109        writer_id.to_string()
110    }
111}
112
113impl TryFrom<String> for WriterId {
114    type Error = String;
115
116    fn try_from(s: String) -> Result<Self, Self::Error> {
117        s.parse()
118    }
119}
120
121impl WriterId {
122    pub(crate) fn new() -> Self {
123        WriterId(*Uuid::new_v4().as_bytes())
124    }
125}
126
127/// A "capability" granting the ability to apply updates to some shard at times
128/// greater or equal to `self.upper()`.
129///
130/// All async methods on ReadHandle retry for as long as they are able, but the
131/// returned [std::future::Future]s implement "cancel on drop" semantics. This
132/// means that callers can add a timeout using [tokio::time::timeout] or
133/// [tokio::time::timeout_at].
134///
135/// ```rust,no_run
136/// # let mut write: mz_persist_client::write::WriteHandle<String, String, u64, i64> = unimplemented!();
137/// # let timeout: std::time::Duration = unimplemented!();
138/// # async {
139/// tokio::time::timeout(timeout, write.fetch_recent_upper()).await
140/// # };
141/// ```
142#[derive(Debug)]
143pub struct WriteHandle<K: Codec, V: Codec, T, D> {
144    pub(crate) cfg: PersistConfig,
145    pub(crate) metrics: Arc<Metrics>,
146    pub(crate) machine: Machine<K, V, T, D>,
147    pub(crate) gc: GarbageCollector<K, V, T, D>,
148    pub(crate) compact: Option<Compactor<K, V, T, D>>,
149    pub(crate) blob: Arc<dyn Blob>,
150    pub(crate) isolated_runtime: Arc<IsolatedRuntime>,
151    pub(crate) writer_id: WriterId,
152    pub(crate) debug_state: HandleDebugState,
153    pub(crate) write_schemas: Schemas<K, V>,
154
155    pub(crate) upper: Antichain<T>,
156    expire_fn: Option<ExpireFn>,
157}
158
159impl<K, V, T, D> WriteHandle<K, V, T, D>
160where
161    K: Debug + Codec,
162    V: Debug + Codec,
163    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
164    D: Monoid + Ord + Codec64 + Send + Sync,
165{
166    pub(crate) fn new(
167        cfg: PersistConfig,
168        metrics: Arc<Metrics>,
169        machine: Machine<K, V, T, D>,
170        gc: GarbageCollector<K, V, T, D>,
171        blob: Arc<dyn Blob>,
172        writer_id: WriterId,
173        purpose: &str,
174        write_schemas: Schemas<K, V>,
175    ) -> Self {
176        let isolated_runtime = Arc::clone(&machine.isolated_runtime);
177        let compact = cfg
178            .compaction_enabled
179            .then(|| Compactor::new(cfg.clone(), Arc::clone(&metrics), gc.clone()));
180        let debug_state = HandleDebugState {
181            hostname: cfg.hostname.to_owned(),
182            purpose: purpose.to_owned(),
183        };
184        let upper = machine.applier.clone_upper();
185        let expire_fn = Self::expire_fn(machine.clone(), gc.clone(), writer_id.clone());
186        WriteHandle {
187            cfg,
188            metrics,
189            machine,
190            gc,
191            compact,
192            blob,
193            isolated_runtime,
194            writer_id,
195            debug_state,
196            write_schemas,
197            upper,
198            expire_fn: Some(expire_fn),
199        }
200    }
201
202    /// Creates a [WriteHandle] for the same shard from an existing
203    /// [ReadHandle].
204    pub fn from_read(read: &ReadHandle<K, V, T, D>, purpose: &str) -> Self {
205        Self::new(
206            read.cfg.clone(),
207            Arc::clone(&read.metrics),
208            read.machine.clone(),
209            read.gc.clone(),
210            Arc::clone(&read.blob),
211            WriterId::new(),
212            purpose,
213            read.read_schemas.clone(),
214        )
215    }
216
217    /// True iff this WriteHandle supports writing without enforcing batch
218    /// bounds checks.
219    pub fn validate_part_bounds_on_write(&self) -> bool {
220        // Note that we require validation when the read checks are enabled, even if the write-time
221        // checks would otherwise be disabled, to avoid batches that would fail at read time.
222        VALIDATE_PART_BOUNDS_ON_WRITE.get(&self.cfg) || VALIDATE_PART_BOUNDS_ON_READ.get(&self.cfg)
223    }
224
225    /// This handle's shard id.
226    pub fn shard_id(&self) -> ShardId {
227        self.machine.shard_id()
228    }
229
230    /// Returns the schema of this writer.
231    pub fn schema_id(&self) -> Option<SchemaId> {
232        self.write_schemas.id
233    }
234
235    /// Registers the write schema, if it isn't already registered.
236    ///
237    /// This method expects that either the shard doesn't yet have any schema registered, or one of
238    /// the registered schemas is the same as the write schema. If all registered schemas are
239    /// different from the write schema, or the shard is a tombstone, it returns `None`.
240    pub async fn try_register_schema(&mut self) -> Option<SchemaId> {
241        let Schemas { id, key, val } = &self.write_schemas;
242
243        if let Some(id) = id {
244            return Some(*id);
245        }
246
247        let (schema_id, maintenance) = self.machine.register_schema(key, val).await;
248        maintenance.start_performing(&self.machine, &self.gc);
249
250        self.write_schemas.id = schema_id;
251        schema_id
252    }
253
254    /// A cached version of the shard-global `upper` frontier.
255    ///
256    /// This is the most recent upper discovered by this handle. It is
257    /// potentially more stale than [Self::shared_upper] but is lock-free and
258    /// allocation-free. This will always be less or equal to the shard-global
259    /// `upper`.
260    pub fn upper(&self) -> &Antichain<T> {
261        &self.upper
262    }
263
264    /// A less-stale cached version of the shard-global `upper` frontier.
265    ///
266    /// This is the most recently known upper for this shard process-wide, but
267    /// unlike [Self::upper] it requires a mutex and a clone. This will always be
268    /// less or equal to the shard-global `upper`.
269    pub fn shared_upper(&self) -> Antichain<T> {
270        self.machine.applier.clone_upper()
271    }
272
273    /// Fetches and returns a recent shard-global `upper`. Importantly, this operation is
274    /// linearized with write operations.
275    ///
276    /// This requires fetching the latest state from consensus and is therefore a potentially
277    /// expensive operation.
278    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
279    pub async fn fetch_recent_upper(&mut self) -> &Antichain<T> {
280        // TODO: Do we even need to track self.upper on WriteHandle or could
281        // WriteHandle::upper just get the one out of machine?
282        self.machine
283            .applier
284            .fetch_upper(|current_upper| self.upper.clone_from(current_upper))
285            .await;
286        &self.upper
287    }
288
289    /// Advance the shard's upper by the given frontier.
290    ///
291    /// If the provided `target` is less than or equal to the shard's upper, this is a no-op.
292    ///
293    /// In contrast to the various compare-and-append methods, this method does not require the
294    /// handle's write schema to be registered with the shard. That is, it is fine to use a dummy
295    /// schema when creating a writer just to advance a shard upper.
296    pub async fn advance_upper(&mut self, target: &Antichain<T>) {
297        // We avoid `fetch_recent_upper` here, to avoid a consensus roundtrip if the known upper is
298        // already beyond the target.
299        let mut lower = self.shared_upper().clone();
300
301        while !PartialOrder::less_equal(target, &lower) {
302            let since = Antichain::from_elem(T::minimum());
303            let desc = Description::new(lower.clone(), target.clone(), since);
304            let batch = HollowBatch::empty(desc);
305
306            let res = self
307                .machine
308                .compare_and_append(&batch, &self.writer_id, &self.debug_state)
309                .await;
310
311            use CompareAndAppendRes::*;
312            let new_upper = match res {
313                Success(_seq_no, maintenance) => {
314                    maintenance.start_performing(&self.machine, &self.gc, self.compact.as_ref());
315                    batch.desc.upper().clone()
316                }
317                UpperMismatch(_seq_no, actual_upper) => actual_upper,
318                InvalidUsage(_invalid_usage) => unreachable!("batch bounds checked above"),
319                InlineBackpressure => unreachable!("batch was empty"),
320            };
321
322            self.upper.clone_from(&new_upper);
323            lower = new_upper;
324        }
325    }
326
327    /// Applies `updates` to this shard and downgrades this handle's upper to
328    /// `upper`.
329    ///
330    /// The innermost `Result` is `Ok` if the updates were successfully written.
331    /// If not, an `Upper` err containing the current writer upper is returned.
332    /// If that happens, we also update our local `upper` to match the current
333    /// upper. This is useful in cases where a timeout happens in between a
334    /// successful write and returning that to the client.
335    ///
336    /// In contrast to [Self::compare_and_append], multiple [WriteHandle]s may
337    /// be used concurrently to write to the same shard, but in this case, the
338    /// data being written must be identical (in the sense of "definite"-ness).
339    /// It's intended for replicated use by source ingestion, sinks, etc.
340    ///
341    /// All times in `updates` must be greater or equal to `lower` and not
342    /// greater or equal to `upper`. A `upper` of the empty antichain "finishes"
343    /// this shard, promising that no more data is ever incoming.
344    ///
345    /// `updates` may be empty, which allows for downgrading `upper` to
346    /// communicate progress. It is possible to call this with `upper` equal to
347    /// `self.upper()` and an empty `updates` (making the call a no-op).
348    ///
349    /// This uses a bounded amount of memory, even when `updates` is very large.
350    /// Individual records, however, should be small enough that we can
351    /// reasonably chunk them up: O(KB) is definitely fine, O(MB) come talk to
352    /// us.
353    ///
354    /// The clunky multi-level Result is to enable more obvious error handling
355    /// in the caller. See <http://sled.rs/errors.html> for details.
356    #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
357    pub async fn append<SB, KB, VB, TB, DB, I>(
358        &mut self,
359        updates: I,
360        lower: Antichain<T>,
361        upper: Antichain<T>,
362    ) -> Result<Result<(), UpperMismatch<T>>, InvalidUsage<T>>
363    where
364        SB: Borrow<((KB, VB), TB, DB)>,
365        KB: Borrow<K>,
366        VB: Borrow<V>,
367        TB: Borrow<T>,
368        DB: Borrow<D>,
369        I: IntoIterator<Item = SB>,
370        D: Send + Sync,
371    {
372        let batch = self.batch(updates, lower.clone(), upper.clone()).await?;
373        self.append_batch(batch, lower, upper).await
374    }
375
376    /// Applies `updates` to this shard and downgrades this handle's upper to
377    /// `new_upper` iff the current global upper of this shard is
378    /// `expected_upper`.
379    ///
380    /// The innermost `Result` is `Ok` if the updates were successfully written.
381    /// If not, an `Upper` err containing the current global upper is returned.
382    ///
383    /// In contrast to [Self::append], this linearizes mutations from all
384    /// writers. It's intended for use as an atomic primitive for timestamp
385    /// bindings, SQL tables, etc.
386    ///
387    /// All times in `updates` must be greater or equal to `expected_upper` and
388    /// not greater or equal to `new_upper`. A `new_upper` of the empty
389    /// antichain "finishes" this shard, promising that no more data is ever
390    /// incoming.
391    ///
392    /// `updates` may be empty, which allows for downgrading `upper` to
393    /// communicate progress. It is possible to heartbeat a writer lease by
394    /// calling this with `new_upper` equal to `self.upper()` and an empty
395    /// `updates` (making the call a no-op).
396    ///
397    /// This uses a bounded amount of memory, even when `updates` is very large.
398    /// Individual records, however, should be small enough that we can
399    /// reasonably chunk them up: O(KB) is definitely fine, O(MB) come talk to
400    /// us.
401    ///
402    /// The clunky multi-level Result is to enable more obvious error handling
403    /// in the caller. See <http://sled.rs/errors.html> for details.
404    #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
405    pub async fn compare_and_append<SB, KB, VB, TB, DB, I>(
406        &mut self,
407        updates: I,
408        expected_upper: Antichain<T>,
409        new_upper: Antichain<T>,
410    ) -> Result<Result<(), UpperMismatch<T>>, InvalidUsage<T>>
411    where
412        SB: Borrow<((KB, VB), TB, DB)>,
413        KB: Borrow<K>,
414        VB: Borrow<V>,
415        TB: Borrow<T>,
416        DB: Borrow<D>,
417        I: IntoIterator<Item = SB>,
418        D: Send + Sync,
419    {
420        let mut batch = self
421            .batch(updates, expected_upper.clone(), new_upper.clone())
422            .await?;
423        match self
424            .compare_and_append_batch(&mut [&mut batch], expected_upper, new_upper, true)
425            .await
426        {
427            ok @ Ok(Ok(())) => ok,
428            err => {
429                // We cannot delete the batch in compare_and_append_batch()
430                // because the caller owns the batch and might want to retry
431                // with a different `expected_upper`. In this function, we
432                // control the batch, so we have to delete it.
433                batch.delete().await;
434                err
435            }
436        }
437    }
438
439    /// Appends the batch of updates to the shard and downgrades this handle's
440    /// upper to `upper`.
441    ///
442    /// The innermost `Result` is `Ok` if the updates were successfully written.
443    /// If not, an `Upper` err containing the current writer upper is returned.
444    /// If that happens, we also update our local `upper` to match the current
445    /// upper. This is useful in cases where a timeout happens in between a
446    /// successful write and returning that to the client.
447    ///
448    /// In contrast to [Self::compare_and_append_batch], multiple [WriteHandle]s
449    /// may be used concurrently to write to the same shard, but in this case,
450    /// the data being written must be identical (in the sense of
451    /// "definite"-ness). It's intended for replicated use by source ingestion,
452    /// sinks, etc.
453    ///
454    /// A `upper` of the empty antichain "finishes" this shard, promising that
455    /// no more data is ever incoming.
456    ///
457    /// The batch may be empty, which allows for downgrading `upper` to
458    /// communicate progress. It is possible to heartbeat a writer lease by
459    /// calling this with `upper` equal to `self.upper()` and an empty `updates`
460    /// (making the call a no-op).
461    ///
462    /// The clunky multi-level Result is to enable more obvious error handling
463    /// in the caller. See <http://sled.rs/errors.html> for details.
464    #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
465    pub async fn append_batch(
466        &mut self,
467        mut batch: Batch<K, V, T, D>,
468        mut lower: Antichain<T>,
469        upper: Antichain<T>,
470    ) -> Result<Result<(), UpperMismatch<T>>, InvalidUsage<T>>
471    where
472        D: Send + Sync,
473    {
474        loop {
475            let res = self
476                .compare_and_append_batch(&mut [&mut batch], lower.clone(), upper.clone(), true)
477                .await?;
478            match res {
479                Ok(()) => {
480                    self.upper = upper;
481                    return Ok(Ok(()));
482                }
483                Err(mismatch) => {
484                    // We tried to to a non-contiguous append, that won't work.
485                    if PartialOrder::less_than(&mismatch.current, &lower) {
486                        self.upper.clone_from(&mismatch.current);
487
488                        batch.delete().await;
489
490                        return Ok(Err(mismatch));
491                    } else if PartialOrder::less_than(&mismatch.current, &upper) {
492                        // Cut down the Description by advancing its lower to the current shard
493                        // upper and try again. IMPORTANT: We can only advance the lower, meaning
494                        // we cut updates away, we must not "extend" the batch by changing to a
495                        // lower that is not beyond the current lower. This invariant is checked by
496                        // the first if branch: if `!(current_upper < lower)` then it holds that
497                        // `lower <= current_upper`.
498                        lower = mismatch.current;
499                    } else {
500                        // We already have updates past this batch's upper, the append is a no-op.
501                        self.upper = mismatch.current;
502
503                        // Because we return a success result, the caller will
504                        // think that the batch was consumed or otherwise used,
505                        // so we have to delete it here.
506                        batch.delete().await;
507
508                        return Ok(Ok(()));
509                    }
510                }
511            }
512        }
513    }
514
515    /// Appends the batch of updates to the shard and downgrades this handle's
516    /// upper to `new_upper` iff the current global upper of this shard is
517    /// `expected_upper`.
518    ///
519    /// The innermost `Result` is `Ok` if the batch was successfully written. If
520    /// not, an `Upper` err containing the current global upper is returned.
521    ///
522    /// In contrast to [Self::append_batch], this linearizes mutations from all
523    /// writers. It's intended for use as an atomic primitive for timestamp
524    /// bindings, SQL tables, etc.
525    ///
526    /// A `new_upper` of the empty antichain "finishes" this shard, promising
527    /// that no more data is ever incoming.
528    ///
529    /// The batch may be empty, which allows for downgrading `upper` to
530    /// communicate progress. It is possible to heartbeat a writer lease by
531    /// calling this with `new_upper` equal to `self.upper()` and an empty
532    /// `updates` (making the call a no-op).
533    ///
534    /// IMPORTANT: In case of an erroneous result the caller is responsible for
535    /// the lifecycle of the `batch`. It can be deleted or it can be used to
536    /// retry with adjusted frontiers.
537    ///
538    /// The clunky multi-level Result is to enable more obvious error handling
539    /// in the caller. See <http://sled.rs/errors.html> for details.
540    ///
541    /// If the `enforce_matching_batch_boundaries` flag is set to `false`:
542    /// We no longer validate that every batch covers the entire range between
543    /// the expected and new uppers, as we wish to allow combining batches that
544    /// cover different subsets of that range, including subsets of that range
545    /// that include no data at all. The caller is responsible for guaranteeing
546    /// that the set of batches provided collectively include all updates for
547    /// the entire range between the expected and new upper.
548    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
549    pub async fn compare_and_append_batch(
550        &mut self,
551        batches: &mut [&mut Batch<K, V, T, D>],
552        expected_upper: Antichain<T>,
553        new_upper: Antichain<T>,
554        validate_part_bounds_on_write: bool,
555    ) -> Result<Result<(), UpperMismatch<T>>, InvalidUsage<T>>
556    where
557        D: Send + Sync,
558    {
559        // Before we append any data, we require a registered write schema.
560        // We expect the caller to ensure our schema is already present... unless this shard is a
561        // tombstone, in which case this write is either a noop or will fail gracefully.
562        let schema_id = self.try_register_schema().await;
563
564        for batch in batches.iter() {
565            if self.machine.shard_id() != batch.shard_id() {
566                return Err(InvalidUsage::BatchNotFromThisShard {
567                    batch_shard: batch.shard_id(),
568                    handle_shard: self.machine.shard_id(),
569                });
570            }
571            assert_code_can_read_data(&self.cfg.build_version, &batch.version);
572            if self.cfg.build_version > batch.version {
573                info!(
574                    shard_id =? self.machine.shard_id(),
575                    batch_version =? batch.version,
576                    writer_version =? self.cfg.build_version,
577                    "Appending batch from the past. This is fine but should be rare. \
578                    TODO: Error on very old versions once the leaked blob detector exists."
579                )
580            }
581            fn assert_schema<A: Codec>(writer_schema: &A::Schema, batch_schema: &bytes::Bytes) {
582                if batch_schema.is_empty() {
583                    // Schema is either trivial or missing!
584                    return;
585                }
586                let batch_schema: A::Schema = A::decode_schema(batch_schema);
587                if *writer_schema != batch_schema {
588                    error!(
589                        ?writer_schema,
590                        ?batch_schema,
591                        "writer and batch schemas should be identical"
592                    );
593                    soft_panic_or_log!("writer and batch schemas should be identical");
594                }
595            }
596            assert_schema::<K>(&*self.write_schemas.key, &batch.schemas.0);
597            assert_schema::<V>(&*self.write_schemas.val, &batch.schemas.1);
598        }
599
600        let lower = expected_upper.clone();
601        let upper = new_upper;
602        let since = Antichain::from_elem(T::minimum());
603        let desc = Description::new(lower, upper, since);
604
605        let mut received_inline_backpressure = false;
606        // Every hollow part must belong to some batch, so we can clean it up when the batch is dropped...
607        // but if we need to merge all our inline parts to a single run in S3, it's not correct to
608        // associate that with any of our individual input batches.
609        // At first, we'll try and put all the inline parts we receive into state... but if we
610        // get backpressured, we retry with this builder set to `Some`, put all our inline data into
611        // it, and ensure it's flushed out to S3 before including it in the batch.
612        let mut inline_batch_builder: Option<(_, BatchBuilder<K, V, T, D>)> = None;
613        let maintenance = loop {
614            let any_batch_rewrite = batches
615                .iter()
616                .any(|x| x.batch.parts.iter().any(|x| x.ts_rewrite().is_some()));
617            let (mut parts, mut num_updates, mut run_splits, mut run_metas) =
618                (vec![], 0, vec![], vec![]);
619            let mut key_storage = None;
620            let mut val_storage = None;
621            for batch in batches.iter() {
622                let () = validate_truncate_batch(
623                    &batch.batch,
624                    &desc,
625                    any_batch_rewrite,
626                    validate_part_bounds_on_write,
627                )?;
628                for (run_meta, run) in batch.batch.runs() {
629                    let start_index = parts.len();
630                    for part in run {
631                        if let (
632                            RunPart::Single(
633                                batch_part @ BatchPart::Inline {
634                                    updates,
635                                    ts_rewrite,
636                                    schema_id: _,
637                                    deprecated_schema_id: _,
638                                },
639                            ),
640                            Some((schema_cache, builder)),
641                        ) = (part, &mut inline_batch_builder)
642                        {
643                            let schema_migration = PartMigration::new(
644                                batch_part,
645                                self.write_schemas.clone(),
646                                schema_cache,
647                            )
648                            .await
649                            .expect("schemas for inline user part");
650
651                            let encoded_part = EncodedPart::from_inline(
652                                &crate::fetch::FetchConfig::from_persist_config(&self.cfg),
653                                &*self.metrics,
654                                self.metrics.read.compaction.clone(),
655                                desc.clone(),
656                                updates,
657                                ts_rewrite.as_ref(),
658                            );
659                            let mut fetched_part = FetchedPart::new(
660                                Arc::clone(&self.metrics),
661                                encoded_part,
662                                schema_migration,
663                                FetchBatchFilter::Compaction {
664                                    since: desc.since().clone(),
665                                },
666                                false,
667                                PartDecodeFormat::Arrow,
668                                None,
669                            );
670
671                            while let Some(((k, v), t, d)) =
672                                fetched_part.next_with_storage(&mut key_storage, &mut val_storage)
673                            {
674                                builder
675                                    .add(&k, &v, &t, &d)
676                                    .await
677                                    .expect("re-encoding just-decoded data");
678                            }
679                        } else {
680                            parts.push(part.clone())
681                        }
682                    }
683
684                    let end_index = parts.len();
685
686                    if start_index == end_index {
687                        continue;
688                    }
689
690                    // Mark the boundary if this is not the first run in the batch.
691                    if start_index != 0 {
692                        run_splits.push(start_index);
693                    }
694                    run_metas.push(run_meta.clone());
695                }
696                num_updates += batch.batch.len;
697            }
698
699            let mut flushed_inline_batch = if let Some((_, builder)) = inline_batch_builder.take() {
700                let mut finished = builder
701                    .finish(desc.upper().clone())
702                    .await
703                    .expect("invalid usage");
704                let cfg = BatchBuilderConfig::new(&self.cfg, self.shard_id());
705                finished
706                    .flush_to_blob(
707                        &cfg,
708                        &self.metrics.inline.backpressure,
709                        &self.isolated_runtime,
710                        &self.write_schemas,
711                    )
712                    .await;
713                Some(finished)
714            } else {
715                None
716            };
717
718            if let Some(batch) = &flushed_inline_batch {
719                for (run_meta, run) in batch.batch.runs() {
720                    assert!(run.len() > 0);
721                    let start_index = parts.len();
722                    if start_index != 0 {
723                        run_splits.push(start_index);
724                    }
725                    run_metas.push(run_meta.clone());
726                    parts.extend(run.iter().cloned())
727                }
728            }
729
730            let mut combined_batch =
731                HollowBatch::new(desc.clone(), parts, num_updates, run_metas, run_splits);
732
733            // The batch may have been written by a writer without a registered schema.
734            // Ensure we have a schema ID in the batch metadata before we append, to avoid type
735            // confusion later.
736            match schema_id {
737                Some(schema_id) => {
738                    ensure_batch_schema(&mut combined_batch, self.shard_id(), schema_id);
739                }
740                None => {
741                    assert!(
742                        self.fetch_recent_upper().await.is_empty(),
743                        "fetching a schema id should only fail when the shard is tombstoned"
744                    )
745                }
746            }
747
748            let res = self
749                .machine
750                .compare_and_append(&combined_batch, &self.writer_id, &self.debug_state)
751                .await;
752
753            match res {
754                CompareAndAppendRes::Success(_seqno, maintenance) => {
755                    self.upper.clone_from(desc.upper());
756                    for batch in batches.iter_mut() {
757                        batch.mark_consumed();
758                    }
759                    if let Some(batch) = &mut flushed_inline_batch {
760                        batch.mark_consumed();
761                    }
762                    break maintenance;
763                }
764                CompareAndAppendRes::InvalidUsage(invalid_usage) => {
765                    if let Some(batch) = flushed_inline_batch.take() {
766                        batch.delete().await;
767                    }
768                    return Err(invalid_usage);
769                }
770                CompareAndAppendRes::UpperMismatch(_seqno, current_upper) => {
771                    if let Some(batch) = flushed_inline_batch.take() {
772                        batch.delete().await;
773                    }
774                    // We tried to to a compare_and_append with the wrong expected upper, that
775                    // won't work. Update the cached upper to the current upper.
776                    self.upper.clone_from(&current_upper);
777                    return Ok(Err(UpperMismatch {
778                        current: current_upper,
779                        expected: expected_upper,
780                    }));
781                }
782                CompareAndAppendRes::InlineBackpressure => {
783                    // We tried to write an inline part, but there was already
784                    // too much in state. Flush it out to s3 and try again.
785                    assert_eq!(received_inline_backpressure, false);
786                    received_inline_backpressure = true;
787                    if COMBINE_INLINE_WRITES.get(&self.cfg) {
788                        inline_batch_builder = Some((
789                            self.machine.applier.schema_cache(),
790                            self.builder(desc.lower().clone()),
791                        ));
792                        continue;
793                    }
794
795                    let cfg = BatchBuilderConfig::new(&self.cfg, self.shard_id());
796                    // We could have a large number of inline parts (imagine the
797                    // sharded persist_sink), do this flushing concurrently.
798                    let flush_batches = batches
799                        .iter_mut()
800                        .map(|batch| async {
801                            batch
802                                .flush_to_blob(
803                                    &cfg,
804                                    &self.metrics.inline.backpressure,
805                                    &self.isolated_runtime,
806                                    &self.write_schemas,
807                                )
808                                .await
809                        })
810                        .collect::<FuturesUnordered<_>>();
811                    let () = flush_batches.collect::<()>().await;
812
813                    for batch in batches.iter() {
814                        assert_eq!(batch.batch.inline_bytes(), 0);
815                    }
816
817                    continue;
818                }
819            }
820        };
821
822        maintenance.start_performing(&self.machine, &self.gc, self.compact.as_ref());
823
824        Ok(Ok(()))
825    }
826
827    /// Turns the given [`ProtoBatch`] back into a [`Batch`] which can be used
828    /// to append it to this shard.
829    pub fn batch_from_transmittable_batch(&self, batch: ProtoBatch) -> Batch<K, V, T, D> {
830        let shard_id: ShardId = batch
831            .shard_id
832            .into_rust()
833            .expect("valid transmittable batch");
834        assert_eq!(shard_id, self.machine.shard_id());
835
836        let ret = Batch {
837            batch_delete_enabled: BATCH_DELETE_ENABLED.get(&self.cfg),
838            metrics: Arc::clone(&self.metrics),
839            shard_metrics: Arc::clone(&self.machine.applier.shard_metrics),
840            version: Version::parse(&batch.version).expect("valid transmittable batch"),
841            schemas: (batch.key_schema, batch.val_schema),
842            batch: batch
843                .batch
844                .into_rust_if_some("ProtoBatch::batch")
845                .expect("valid transmittable batch"),
846            blob: Arc::clone(&self.blob),
847            _phantom: std::marker::PhantomData,
848        };
849        assert_eq!(ret.shard_id(), self.machine.shard_id());
850        ret
851    }
852
853    /// Returns a [BatchBuilder] that can be used to write a batch of updates to
854    /// blob storage which can then be appended to this shard using
855    /// [Self::compare_and_append_batch] or [Self::append_batch].
856    ///
857    /// It is correct to create an empty batch, which allows for downgrading
858    /// `upper` to communicate progress. (see [Self::compare_and_append_batch]
859    /// or [Self::append_batch])
860    ///
861    /// The builder uses a bounded amount of memory, even when the number of
862    /// updates is very large. Individual records, however, should be small
863    /// enough that we can reasonably chunk them up: O(KB) is definitely fine,
864    /// O(MB) come talk to us.
865    pub fn builder(&self, lower: Antichain<T>) -> BatchBuilder<K, V, T, D> {
866        Self::builder_inner(
867            &self.cfg,
868            CompactConfig::new(&self.cfg, self.shard_id()),
869            Arc::clone(&self.metrics),
870            Arc::clone(&self.machine.applier.shard_metrics),
871            &self.metrics.user,
872            Arc::clone(&self.isolated_runtime),
873            Arc::clone(&self.blob),
874            self.shard_id(),
875            self.write_schemas.clone(),
876            lower,
877        )
878    }
879
880    /// Implementation of [Self::builder], so that we can share the
881    /// implementation in `PersistClient`.
882    pub(crate) fn builder_inner(
883        persist_cfg: &PersistConfig,
884        compact_cfg: CompactConfig,
885        metrics: Arc<Metrics>,
886        shard_metrics: Arc<ShardMetrics>,
887        user_batch_metrics: &BatchWriteMetrics,
888        isolated_runtime: Arc<IsolatedRuntime>,
889        blob: Arc<dyn Blob>,
890        shard_id: ShardId,
891        schemas: Schemas<K, V>,
892        lower: Antichain<T>,
893    ) -> BatchBuilder<K, V, T, D> {
894        let parts = if let Some(max_runs) = compact_cfg.batch.max_runs {
895            BatchParts::new_compacting::<K, V, D>(
896                compact_cfg,
897                Description::new(
898                    lower.clone(),
899                    Antichain::new(),
900                    Antichain::from_elem(T::minimum()),
901                ),
902                max_runs,
903                Arc::clone(&metrics),
904                shard_metrics,
905                shard_id,
906                Arc::clone(&blob),
907                isolated_runtime,
908                user_batch_metrics,
909                schemas.clone(),
910            )
911        } else {
912            BatchParts::new_ordered::<D>(
913                compact_cfg.batch,
914                RunOrder::Unordered,
915                Arc::clone(&metrics),
916                shard_metrics,
917                shard_id,
918                Arc::clone(&blob),
919                isolated_runtime,
920                user_batch_metrics,
921            )
922        };
923        let builder = BatchBuilderInternal::new(
924            BatchBuilderConfig::new(persist_cfg, shard_id),
925            parts,
926            metrics,
927            schemas,
928            blob,
929            shard_id,
930            persist_cfg.build_version.clone(),
931        );
932        BatchBuilder::new(
933            builder,
934            Description::new(lower, Antichain::new(), Antichain::from_elem(T::minimum())),
935        )
936    }
937
938    /// Uploads the given `updates` as one `Batch` to the blob store and returns
939    /// a handle to the batch.
940    #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
941    pub async fn batch<SB, KB, VB, TB, DB, I>(
942        &mut self,
943        updates: I,
944        lower: Antichain<T>,
945        upper: Antichain<T>,
946    ) -> Result<Batch<K, V, T, D>, InvalidUsage<T>>
947    where
948        SB: Borrow<((KB, VB), TB, DB)>,
949        KB: Borrow<K>,
950        VB: Borrow<V>,
951        TB: Borrow<T>,
952        DB: Borrow<D>,
953        I: IntoIterator<Item = SB>,
954    {
955        let iter = updates.into_iter();
956
957        let mut builder = self.builder(lower.clone());
958
959        for update in iter {
960            let ((k, v), t, d) = update.borrow();
961            let (k, v, t, d) = (k.borrow(), v.borrow(), t.borrow(), d.borrow());
962            match builder.add(k, v, t, d).await {
963                Ok(Added::Record | Added::RecordAndParts) => (),
964                Err(invalid_usage) => return Err(invalid_usage),
965            }
966        }
967
968        builder.finish(upper.clone()).await
969    }
970
971    /// Blocks until the given `frontier` is less than the upper of the shard.
972    pub async fn wait_for_upper_past(&mut self, frontier: &Antichain<T>) {
973        let mut watch = self.machine.applier.watch();
974        self.machine
975            .wait_for_upper_past(
976                frontier,
977                &mut watch,
978                None,
979                &self.metrics.retries.next_listen_batch, // TODO: new retry metrics for these?
980                next_listen_batch_retry_params(&self.cfg),
981            )
982            .await;
983        let upper = self.machine.applier.clone_upper();
984        if PartialOrder::less_than(&self.upper, &upper) {
985            self.upper.clone_from(&upper);
986        }
987        assert!(PartialOrder::less_than(frontier, &self.upper));
988    }
989
990    /// Politely expires this writer, releasing any associated state.
991    ///
992    /// There is a best-effort impl in Drop to expire a writer that wasn't
993    /// explictly expired with this method. When possible, explicit expiry is
994    /// still preferred because the Drop one is best effort and is dependant on
995    /// a tokio [Handle] being available in the TLC at the time of drop (which
996    /// is a bit subtle). Also, explicit expiry allows for control over when it
997    /// happens.
998    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
999    pub async fn expire(mut self) {
1000        let Some(expire_fn) = self.expire_fn.take() else {
1001            return;
1002        };
1003        expire_fn.0().await;
1004    }
1005
1006    fn expire_fn(
1007        machine: Machine<K, V, T, D>,
1008        gc: GarbageCollector<K, V, T, D>,
1009        writer_id: WriterId,
1010    ) -> ExpireFn {
1011        ExpireFn(Box::new(move || {
1012            Box::pin(async move {
1013                let (_, maintenance) = machine.expire_writer(&writer_id).await;
1014                maintenance.start_performing(&machine, &gc);
1015            })
1016        }))
1017    }
1018
1019    /// Test helper for an [Self::append] call that is expected to succeed.
1020    #[cfg(test)]
1021    #[track_caller]
1022    pub async fn expect_append<L, U>(&mut self, updates: &[((K, V), T, D)], lower: L, new_upper: U)
1023    where
1024        L: Into<Antichain<T>>,
1025        U: Into<Antichain<T>>,
1026        D: Send + Sync,
1027    {
1028        self.append(updates.iter(), lower.into(), new_upper.into())
1029            .await
1030            .expect("invalid usage")
1031            .expect("unexpected upper");
1032    }
1033
1034    /// Test helper for a [Self::compare_and_append] call that is expected to
1035    /// succeed.
1036    #[cfg(test)]
1037    #[track_caller]
1038    pub async fn expect_compare_and_append(
1039        &mut self,
1040        updates: &[((K, V), T, D)],
1041        expected_upper: T,
1042        new_upper: T,
1043    ) where
1044        D: Send + Sync,
1045    {
1046        self.compare_and_append(
1047            updates.iter().map(|((k, v), t, d)| ((k, v), t, d)),
1048            Antichain::from_elem(expected_upper),
1049            Antichain::from_elem(new_upper),
1050        )
1051        .await
1052        .expect("invalid usage")
1053        .expect("unexpected upper")
1054    }
1055
1056    /// Test helper for a [Self::compare_and_append_batch] call that is expected
1057    /// to succeed.
1058    #[cfg(test)]
1059    #[track_caller]
1060    pub async fn expect_compare_and_append_batch(
1061        &mut self,
1062        batches: &mut [&mut Batch<K, V, T, D>],
1063        expected_upper: T,
1064        new_upper: T,
1065    ) {
1066        self.compare_and_append_batch(
1067            batches,
1068            Antichain::from_elem(expected_upper),
1069            Antichain::from_elem(new_upper),
1070            true,
1071        )
1072        .await
1073        .expect("invalid usage")
1074        .expect("unexpected upper")
1075    }
1076
1077    /// Test helper for an [Self::append] call that is expected to succeed.
1078    #[cfg(test)]
1079    #[track_caller]
1080    pub async fn expect_batch(
1081        &mut self,
1082        updates: &[((K, V), T, D)],
1083        lower: T,
1084        upper: T,
1085    ) -> Batch<K, V, T, D> {
1086        self.batch(
1087            updates.iter(),
1088            Antichain::from_elem(lower),
1089            Antichain::from_elem(upper),
1090        )
1091        .await
1092        .expect("invalid usage")
1093    }
1094}
1095
1096impl<K: Codec, V: Codec, T, D> Drop for WriteHandle<K, V, T, D> {
1097    fn drop(&mut self) {
1098        let Some(expire_fn) = self.expire_fn.take() else {
1099            return;
1100        };
1101        let handle = match Handle::try_current() {
1102            Ok(x) => x,
1103            Err(_) => {
1104                warn!(
1105                    "WriteHandle {} dropped without being explicitly expired, falling back to lease timeout",
1106                    self.writer_id
1107                );
1108                return;
1109            }
1110        };
1111        // Spawn a best-effort task to expire this write handle. It's fine if
1112        // this doesn't run to completion, we'd just have to wait out the lease
1113        // before the shard-global since is unblocked.
1114        //
1115        // Intentionally create the span outside the task to set the parent.
1116        let expire_span = debug_span!("drop::expire");
1117        handle.spawn_named(
1118            || format!("WriteHandle::expire ({})", self.writer_id),
1119            expire_fn.0().instrument(expire_span),
1120        );
1121    }
1122}
1123
1124/// Ensure the given batch uses the given schema ID.
1125///
1126/// If the batch has no schema set, initialize it to the given one.
1127/// If the batch has a schema set, assert that it matches the given one.
1128fn ensure_batch_schema<T>(batch: &mut HollowBatch<T>, shard_id: ShardId, schema_id: SchemaId)
1129where
1130    T: Timestamp + Lattice + Codec64,
1131{
1132    let ensure = |id: &mut Option<SchemaId>| match id {
1133        Some(id) => assert_eq!(*id, schema_id, "schema ID mismatch; shard={shard_id}"),
1134        None => *id = Some(schema_id),
1135    };
1136
1137    for run_meta in &mut batch.run_meta {
1138        ensure(&mut run_meta.schema);
1139    }
1140    for part in &mut batch.parts {
1141        match part {
1142            RunPart::Single(BatchPart::Hollow(part)) => ensure(&mut part.schema_id),
1143            RunPart::Single(BatchPart::Inline { schema_id, .. }) => ensure(schema_id),
1144            RunPart::Many(_hollow_run_ref) => {
1145                // TODO: Fetch the parts in this run and rewrite them too. Alternatively, make
1146                // `run_meta` the only place we keep schema IDs, so rewriting parts isn't
1147                // necessary.
1148            }
1149        }
1150    }
1151}
1152
1153#[cfg(test)]
1154mod tests {
1155    use std::str::FromStr;
1156    use std::sync::mpsc;
1157
1158    use differential_dataflow::consolidation::consolidate_updates;
1159    use futures_util::FutureExt;
1160    use mz_dyncfg::ConfigUpdates;
1161    use mz_ore::collections::CollectionExt;
1162    use mz_ore::task;
1163    use serde_json::json;
1164
1165    use crate::cache::PersistClientCache;
1166    use crate::tests::{all_ok, new_test_client};
1167    use crate::{PersistLocation, ShardId};
1168
1169    use super::*;
1170
1171    #[mz_persist_proc::test(tokio::test)]
1172    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1173    async fn empty_batches(dyncfgs: ConfigUpdates) {
1174        let data = [
1175            (("1".to_owned(), "one".to_owned()), 1, 1),
1176            (("2".to_owned(), "two".to_owned()), 2, 1),
1177            (("3".to_owned(), "three".to_owned()), 3, 1),
1178        ];
1179
1180        let (mut write, _) = new_test_client(&dyncfgs)
1181            .await
1182            .expect_open::<String, String, u64, i64>(ShardId::new())
1183            .await;
1184        let blob = Arc::clone(&write.blob);
1185
1186        // Write an initial batch.
1187        let mut upper = 3;
1188        write.expect_append(&data[..2], vec![0], vec![upper]).await;
1189
1190        // Write a bunch of empty batches. This shouldn't write blobs, so the count should stay the same.
1191        let mut count_before = 0;
1192        blob.list_keys_and_metadata("", &mut |_| {
1193            count_before += 1;
1194        })
1195        .await
1196        .expect("list_keys failed");
1197        for _ in 0..5 {
1198            let new_upper = upper + 1;
1199            write.expect_compare_and_append(&[], upper, new_upper).await;
1200            upper = new_upper;
1201        }
1202        let mut count_after = 0;
1203        blob.list_keys_and_metadata("", &mut |_| {
1204            count_after += 1;
1205        })
1206        .await
1207        .expect("list_keys failed");
1208        assert_eq!(count_after, count_before);
1209    }
1210
1211    #[mz_persist_proc::test(tokio::test)]
1212    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1213    async fn compare_and_append_batch_multi(dyncfgs: ConfigUpdates) {
1214        let data0 = vec![
1215            (("1".to_owned(), "one".to_owned()), 1, 1),
1216            (("2".to_owned(), "two".to_owned()), 2, 1),
1217            (("4".to_owned(), "four".to_owned()), 4, 1),
1218        ];
1219        let data1 = vec![
1220            (("1".to_owned(), "one".to_owned()), 1, 1),
1221            (("2".to_owned(), "two".to_owned()), 2, 1),
1222            (("3".to_owned(), "three".to_owned()), 3, 1),
1223        ];
1224
1225        let (mut write, mut read) = new_test_client(&dyncfgs)
1226            .await
1227            .expect_open::<String, String, u64, i64>(ShardId::new())
1228            .await;
1229
1230        let mut batch0 = write.expect_batch(&data0, 0, 5).await;
1231        let mut batch1 = write.expect_batch(&data1, 0, 4).await;
1232
1233        write
1234            .expect_compare_and_append_batch(&mut [&mut batch0, &mut batch1], 0, 4)
1235            .await;
1236
1237        let batch = write
1238            .machine
1239            .unleased_snapshot(&Antichain::from_elem(3))
1240            .await
1241            .expect("just wrote this")
1242            .into_element();
1243
1244        assert!(batch.runs().count() >= 2);
1245
1246        let expected = vec![
1247            (("1".to_owned(), "one".to_owned()), 1, 2),
1248            (("2".to_owned(), "two".to_owned()), 2, 2),
1249            (("3".to_owned(), "three".to_owned()), 3, 1),
1250        ];
1251        let mut actual = read.expect_snapshot_and_fetch(3).await;
1252        consolidate_updates(&mut actual);
1253        assert_eq!(actual, all_ok(&expected, 3));
1254    }
1255
1256    #[mz_ore::test]
1257    fn writer_id_human_readable_serde() {
1258        #[derive(Debug, Serialize, Deserialize)]
1259        struct Container {
1260            writer_id: WriterId,
1261        }
1262
1263        // roundtrip through json
1264        let id = WriterId::from_str("w00000000-1234-5678-0000-000000000000").expect("valid id");
1265        assert_eq!(
1266            id,
1267            serde_json::from_value(serde_json::to_value(id.clone()).expect("serializable"))
1268                .expect("deserializable")
1269        );
1270
1271        // deserialize a serialized string directly
1272        assert_eq!(
1273            id,
1274            serde_json::from_str("\"w00000000-1234-5678-0000-000000000000\"")
1275                .expect("deserializable")
1276        );
1277
1278        // roundtrip id through a container type
1279        let json = json!({ "writer_id": id });
1280        assert_eq!(
1281            "{\"writer_id\":\"w00000000-1234-5678-0000-000000000000\"}",
1282            &json.to_string()
1283        );
1284        let container: Container = serde_json::from_value(json).expect("deserializable");
1285        assert_eq!(container.writer_id, id);
1286    }
1287
1288    #[mz_persist_proc::test(tokio::test)]
1289    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1290    async fn hollow_batch_roundtrip(dyncfgs: ConfigUpdates) {
1291        let data = vec![
1292            (("1".to_owned(), "one".to_owned()), 1, 1),
1293            (("2".to_owned(), "two".to_owned()), 2, 1),
1294            (("3".to_owned(), "three".to_owned()), 3, 1),
1295        ];
1296
1297        let (mut write, mut read) = new_test_client(&dyncfgs)
1298            .await
1299            .expect_open::<String, String, u64, i64>(ShardId::new())
1300            .await;
1301
1302        // This test is a bit more complex than it should be. It would be easier
1303        // if we could just compare the rehydrated batch to the original batch.
1304        // But a) turning a batch into a hollow batch consumes it, and b) Batch
1305        // doesn't have Eq/PartialEq.
1306        let batch = write.expect_batch(&data, 0, 4).await;
1307        let hollow_batch = batch.into_transmittable_batch();
1308        let mut rehydrated_batch = write.batch_from_transmittable_batch(hollow_batch);
1309
1310        write
1311            .expect_compare_and_append_batch(&mut [&mut rehydrated_batch], 0, 4)
1312            .await;
1313
1314        let expected = vec![
1315            (("1".to_owned(), "one".to_owned()), 1, 1),
1316            (("2".to_owned(), "two".to_owned()), 2, 1),
1317            (("3".to_owned(), "three".to_owned()), 3, 1),
1318        ];
1319        let mut actual = read.expect_snapshot_and_fetch(3).await;
1320        consolidate_updates(&mut actual);
1321        assert_eq!(actual, all_ok(&expected, 3));
1322    }
1323
1324    #[mz_persist_proc::test(tokio::test)]
1325    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1326    async fn wait_for_upper_past(dyncfgs: ConfigUpdates) {
1327        let client = new_test_client(&dyncfgs).await;
1328        let (mut write, _) = client.expect_open::<(), (), u64, i64>(ShardId::new()).await;
1329        let five = Antichain::from_elem(5);
1330
1331        // Upper is not past 5.
1332        assert_eq!(write.wait_for_upper_past(&five).now_or_never(), None);
1333
1334        // Upper is still not past 5.
1335        write
1336            .expect_compare_and_append(&[(((), ()), 1, 1)], 0, 5)
1337            .await;
1338        assert_eq!(write.wait_for_upper_past(&five).now_or_never(), None);
1339
1340        // Upper is past 5.
1341        write
1342            .expect_compare_and_append(&[(((), ()), 5, 1)], 5, 7)
1343            .await;
1344        assert_eq!(write.wait_for_upper_past(&five).now_or_never(), Some(()));
1345        assert_eq!(write.upper(), &Antichain::from_elem(7));
1346
1347        // Waiting for previous uppers does not regress the handle's cached
1348        // upper.
1349        assert_eq!(
1350            write
1351                .wait_for_upper_past(&Antichain::from_elem(2))
1352                .now_or_never(),
1353            Some(())
1354        );
1355        assert_eq!(write.upper(), &Antichain::from_elem(7));
1356    }
1357
1358    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
1359    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1360    async fn fetch_recent_upper_linearized() {
1361        type Timestamp = u64;
1362        let max_upper = 1000;
1363
1364        let shard_id = ShardId::new();
1365        let mut clients = PersistClientCache::new_no_metrics();
1366        let upper_writer_client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
1367        let (mut upper_writer, _) = upper_writer_client
1368            .expect_open::<(), (), Timestamp, i64>(shard_id)
1369            .await;
1370        // Clear the state cache between each client to maximally disconnect
1371        // them from each other.
1372        clients.clear_state_cache();
1373        let upper_reader_client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
1374        let (mut upper_reader, _) = upper_reader_client
1375            .expect_open::<(), (), Timestamp, i64>(shard_id)
1376            .await;
1377        let (tx, rx) = mpsc::channel();
1378
1379        let task = task::spawn(|| "upper-reader", async move {
1380            let mut upper = Timestamp::MIN;
1381
1382            while upper < max_upper {
1383                while let Ok(new_upper) = rx.try_recv() {
1384                    upper = new_upper;
1385                }
1386
1387                let recent_upper = upper_reader
1388                    .fetch_recent_upper()
1389                    .await
1390                    .as_option()
1391                    .cloned()
1392                    .expect("u64 is totally ordered and the shard is not finalized");
1393                assert!(
1394                    recent_upper >= upper,
1395                    "recent upper {recent_upper:?} is less than known upper {upper:?}"
1396                );
1397            }
1398        });
1399
1400        for upper in Timestamp::MIN..max_upper {
1401            let next_upper = upper + 1;
1402            upper_writer
1403                .expect_compare_and_append(&[], upper, next_upper)
1404                .await;
1405            tx.send(next_upper).expect("send failed");
1406        }
1407
1408        task.await;
1409    }
1410}