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 bounds_truncated = 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                    let mut run_meta = run_meta.clone();
695                    if bounds_truncated {
696                        run_meta.set_bounds_truncated();
697                    }
698                    run_metas.push(run_meta);
699                }
700                num_updates += batch.batch.len;
701            }
702
703            let mut flushed_inline_batch = if let Some((_, builder)) = inline_batch_builder.take() {
704                let mut finished = builder
705                    .finish(desc.upper().clone())
706                    .await
707                    .expect("invalid usage");
708                let cfg = BatchBuilderConfig::new(&self.cfg, self.shard_id());
709                finished
710                    .flush_to_blob(
711                        &cfg,
712                        &self.metrics.inline.backpressure,
713                        &self.isolated_runtime,
714                        &self.write_schemas,
715                    )
716                    .await;
717                Some(finished)
718            } else {
719                None
720            };
721
722            if let Some(batch) = &flushed_inline_batch {
723                for (run_meta, run) in batch.batch.runs() {
724                    assert!(run.len() > 0);
725                    let start_index = parts.len();
726                    if start_index != 0 {
727                        run_splits.push(start_index);
728                    }
729                    run_metas.push(run_meta.clone());
730                    parts.extend(run.iter().cloned())
731                }
732            }
733
734            let mut combined_batch =
735                HollowBatch::new(desc.clone(), parts, num_updates, run_metas, run_splits);
736
737            // The batch may have been written by a writer without a registered schema.
738            // Ensure we have a schema ID in the batch metadata before we append, to avoid type
739            // confusion later.
740            match schema_id {
741                Some(schema_id) => {
742                    ensure_batch_schema(&mut combined_batch, self.shard_id(), schema_id);
743                }
744                None => {
745                    assert!(
746                        self.fetch_recent_upper().await.is_empty(),
747                        "fetching a schema id should only fail when the shard is tombstoned"
748                    )
749                }
750            }
751
752            let res = self
753                .machine
754                .compare_and_append(&combined_batch, &self.writer_id, &self.debug_state)
755                .await;
756
757            match res {
758                CompareAndAppendRes::Success(_seqno, maintenance) => {
759                    self.upper.clone_from(desc.upper());
760                    for batch in batches.iter_mut() {
761                        batch.mark_consumed();
762                    }
763                    if let Some(batch) = &mut flushed_inline_batch {
764                        batch.mark_consumed();
765                    }
766                    break maintenance;
767                }
768                CompareAndAppendRes::InvalidUsage(invalid_usage) => {
769                    if let Some(batch) = flushed_inline_batch.take() {
770                        batch.delete().await;
771                    }
772                    return Err(invalid_usage);
773                }
774                CompareAndAppendRes::UpperMismatch(_seqno, current_upper) => {
775                    if let Some(batch) = flushed_inline_batch.take() {
776                        batch.delete().await;
777                    }
778                    // We tried to to a compare_and_append with the wrong expected upper, that
779                    // won't work. Update the cached upper to the current upper.
780                    self.upper.clone_from(&current_upper);
781                    return Ok(Err(UpperMismatch {
782                        current: current_upper,
783                        expected: expected_upper,
784                    }));
785                }
786                CompareAndAppendRes::InlineBackpressure => {
787                    // We tried to write an inline part, but there was already
788                    // too much in state. Flush it out to s3 and try again.
789                    assert_eq!(received_inline_backpressure, false);
790                    received_inline_backpressure = true;
791                    if COMBINE_INLINE_WRITES.get(&self.cfg) {
792                        inline_batch_builder = Some((
793                            self.machine.applier.schema_cache(),
794                            self.builder(desc.lower().clone()),
795                        ));
796                        continue;
797                    }
798
799                    let cfg = BatchBuilderConfig::new(&self.cfg, self.shard_id());
800                    // We could have a large number of inline parts (imagine the
801                    // sharded persist_sink), do this flushing concurrently.
802                    let flush_batches = batches
803                        .iter_mut()
804                        .map(|batch| async {
805                            batch
806                                .flush_to_blob(
807                                    &cfg,
808                                    &self.metrics.inline.backpressure,
809                                    &self.isolated_runtime,
810                                    &self.write_schemas,
811                                )
812                                .await
813                        })
814                        .collect::<FuturesUnordered<_>>();
815                    let () = flush_batches.collect::<()>().await;
816
817                    for batch in batches.iter() {
818                        assert_eq!(batch.batch.inline_bytes(), 0);
819                    }
820
821                    continue;
822                }
823            }
824        };
825
826        maintenance.start_performing(&self.machine, &self.gc, self.compact.as_ref());
827
828        Ok(Ok(()))
829    }
830
831    /// Turns the given [`ProtoBatch`] back into a [`Batch`] which can be used
832    /// to append it to this shard.
833    pub fn batch_from_transmittable_batch(&self, batch: ProtoBatch) -> Batch<K, V, T, D> {
834        let shard_id: ShardId = batch
835            .shard_id
836            .into_rust()
837            .expect("valid transmittable batch");
838        assert_eq!(shard_id, self.machine.shard_id());
839
840        let ret = Batch {
841            batch_delete_enabled: BATCH_DELETE_ENABLED.get(&self.cfg),
842            metrics: Arc::clone(&self.metrics),
843            shard_metrics: Arc::clone(&self.machine.applier.shard_metrics),
844            version: Version::parse(&batch.version).expect("valid transmittable batch"),
845            schemas: (batch.key_schema, batch.val_schema),
846            batch: batch
847                .batch
848                .into_rust_if_some("ProtoBatch::batch")
849                .expect("valid transmittable batch"),
850            blob: Arc::clone(&self.blob),
851            _phantom: std::marker::PhantomData,
852        };
853        assert_eq!(ret.shard_id(), self.machine.shard_id());
854        ret
855    }
856
857    /// Returns a [BatchBuilder] that can be used to write a batch of updates to
858    /// blob storage which can then be appended to this shard using
859    /// [Self::compare_and_append_batch] or [Self::append_batch].
860    ///
861    /// It is correct to create an empty batch, which allows for downgrading
862    /// `upper` to communicate progress. (see [Self::compare_and_append_batch]
863    /// or [Self::append_batch])
864    ///
865    /// The builder uses a bounded amount of memory, even when the number of
866    /// updates is very large. Individual records, however, should be small
867    /// enough that we can reasonably chunk them up: O(KB) is definitely fine,
868    /// O(MB) come talk to us.
869    pub fn builder(&self, lower: Antichain<T>) -> BatchBuilder<K, V, T, D> {
870        Self::builder_inner(
871            &self.cfg,
872            CompactConfig::new(&self.cfg, self.shard_id()),
873            Arc::clone(&self.metrics),
874            Arc::clone(&self.machine.applier.shard_metrics),
875            &self.metrics.user,
876            Arc::clone(&self.isolated_runtime),
877            Arc::clone(&self.blob),
878            self.shard_id(),
879            self.write_schemas.clone(),
880            lower,
881        )
882    }
883
884    /// Implementation of [Self::builder], so that we can share the
885    /// implementation in `PersistClient`.
886    pub(crate) fn builder_inner(
887        persist_cfg: &PersistConfig,
888        compact_cfg: CompactConfig,
889        metrics: Arc<Metrics>,
890        shard_metrics: Arc<ShardMetrics>,
891        user_batch_metrics: &BatchWriteMetrics,
892        isolated_runtime: Arc<IsolatedRuntime>,
893        blob: Arc<dyn Blob>,
894        shard_id: ShardId,
895        schemas: Schemas<K, V>,
896        lower: Antichain<T>,
897    ) -> BatchBuilder<K, V, T, D> {
898        let parts = if let Some(max_runs) = compact_cfg.batch.max_runs {
899            BatchParts::new_compacting::<K, V, D>(
900                compact_cfg,
901                Description::new(
902                    lower.clone(),
903                    Antichain::new(),
904                    Antichain::from_elem(T::minimum()),
905                ),
906                max_runs,
907                Arc::clone(&metrics),
908                shard_metrics,
909                shard_id,
910                Arc::clone(&blob),
911                isolated_runtime,
912                user_batch_metrics,
913                schemas.clone(),
914            )
915        } else {
916            BatchParts::new_ordered::<D>(
917                compact_cfg.batch,
918                RunOrder::Unordered,
919                Arc::clone(&metrics),
920                shard_metrics,
921                shard_id,
922                Arc::clone(&blob),
923                isolated_runtime,
924                user_batch_metrics,
925            )
926        };
927        let builder = BatchBuilderInternal::new(
928            BatchBuilderConfig::new(persist_cfg, shard_id),
929            parts,
930            metrics,
931            schemas,
932            blob,
933            shard_id,
934            persist_cfg.build_version.clone(),
935        );
936        BatchBuilder::new(
937            builder,
938            Description::new(lower, Antichain::new(), Antichain::from_elem(T::minimum())),
939        )
940    }
941
942    /// Uploads the given `updates` as one `Batch` to the blob store and returns
943    /// a handle to the batch.
944    #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
945    pub async fn batch<SB, KB, VB, TB, DB, I>(
946        &mut self,
947        updates: I,
948        lower: Antichain<T>,
949        upper: Antichain<T>,
950    ) -> Result<Batch<K, V, T, D>, InvalidUsage<T>>
951    where
952        SB: Borrow<((KB, VB), TB, DB)>,
953        KB: Borrow<K>,
954        VB: Borrow<V>,
955        TB: Borrow<T>,
956        DB: Borrow<D>,
957        I: IntoIterator<Item = SB>,
958    {
959        let iter = updates.into_iter();
960
961        let mut builder = self.builder(lower.clone());
962
963        for update in iter {
964            let ((k, v), t, d) = update.borrow();
965            let (k, v, t, d) = (k.borrow(), v.borrow(), t.borrow(), d.borrow());
966            match builder.add(k, v, t, d).await {
967                Ok(Added::Record | Added::RecordAndParts) => (),
968                Err(invalid_usage) => return Err(invalid_usage),
969            }
970        }
971
972        builder.finish(upper.clone()).await
973    }
974
975    /// Blocks until the given `frontier` is less than the upper of the shard.
976    pub async fn wait_for_upper_past(&mut self, frontier: &Antichain<T>) {
977        let mut watch = self.machine.applier.watch();
978        self.machine
979            .wait_for_upper_past(
980                frontier,
981                &mut watch,
982                None,
983                &self.metrics.retries.next_listen_batch, // TODO: new retry metrics for these?
984                next_listen_batch_retry_params(&self.cfg),
985            )
986            .await;
987        let upper = self.machine.applier.clone_upper();
988        if PartialOrder::less_than(&self.upper, &upper) {
989            self.upper.clone_from(&upper);
990        }
991        assert!(PartialOrder::less_than(frontier, &self.upper));
992    }
993
994    /// Politely expires this writer, releasing any associated state.
995    ///
996    /// There is a best-effort impl in Drop to expire a writer that wasn't
997    /// explictly expired with this method. When possible, explicit expiry is
998    /// still preferred because the Drop one is best effort and is dependant on
999    /// a tokio [Handle] being available in the TLC at the time of drop (which
1000    /// is a bit subtle). Also, explicit expiry allows for control over when it
1001    /// happens.
1002    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
1003    pub async fn expire(mut self) {
1004        let Some(expire_fn) = self.expire_fn.take() else {
1005            return;
1006        };
1007        expire_fn.0().await;
1008    }
1009
1010    fn expire_fn(
1011        machine: Machine<K, V, T, D>,
1012        gc: GarbageCollector<K, V, T, D>,
1013        writer_id: WriterId,
1014    ) -> ExpireFn {
1015        ExpireFn(Box::new(move || {
1016            Box::pin(async move {
1017                let (_, maintenance) = machine.expire_writer(&writer_id).await;
1018                maintenance.start_performing(&machine, &gc);
1019            })
1020        }))
1021    }
1022
1023    /// Test helper for an [Self::append] call that is expected to succeed.
1024    #[cfg(test)]
1025    #[track_caller]
1026    pub async fn expect_append<L, U>(&mut self, updates: &[((K, V), T, D)], lower: L, new_upper: U)
1027    where
1028        L: Into<Antichain<T>>,
1029        U: Into<Antichain<T>>,
1030        D: Send + Sync,
1031    {
1032        self.append(updates.iter(), lower.into(), new_upper.into())
1033            .await
1034            .expect("invalid usage")
1035            .expect("unexpected upper");
1036    }
1037
1038    /// Test helper for a [Self::compare_and_append] call that is expected to
1039    /// succeed.
1040    #[cfg(test)]
1041    #[track_caller]
1042    pub async fn expect_compare_and_append(
1043        &mut self,
1044        updates: &[((K, V), T, D)],
1045        expected_upper: T,
1046        new_upper: T,
1047    ) where
1048        D: Send + Sync,
1049    {
1050        self.compare_and_append(
1051            updates.iter().map(|((k, v), t, d)| ((k, v), t, d)),
1052            Antichain::from_elem(expected_upper),
1053            Antichain::from_elem(new_upper),
1054        )
1055        .await
1056        .expect("invalid usage")
1057        .expect("unexpected upper")
1058    }
1059
1060    /// Test helper for a [Self::compare_and_append_batch] call that is expected
1061    /// to succeed.
1062    #[cfg(test)]
1063    #[track_caller]
1064    pub async fn expect_compare_and_append_batch(
1065        &mut self,
1066        batches: &mut [&mut Batch<K, V, T, D>],
1067        expected_upper: T,
1068        new_upper: T,
1069    ) {
1070        self.compare_and_append_batch(
1071            batches,
1072            Antichain::from_elem(expected_upper),
1073            Antichain::from_elem(new_upper),
1074            true,
1075        )
1076        .await
1077        .expect("invalid usage")
1078        .expect("unexpected upper")
1079    }
1080
1081    /// Test helper for an [Self::append] call that is expected to succeed.
1082    #[cfg(test)]
1083    #[track_caller]
1084    pub async fn expect_batch(
1085        &mut self,
1086        updates: &[((K, V), T, D)],
1087        lower: T,
1088        upper: T,
1089    ) -> Batch<K, V, T, D> {
1090        self.batch(
1091            updates.iter(),
1092            Antichain::from_elem(lower),
1093            Antichain::from_elem(upper),
1094        )
1095        .await
1096        .expect("invalid usage")
1097    }
1098}
1099
1100impl<K: Codec, V: Codec, T, D> Drop for WriteHandle<K, V, T, D> {
1101    fn drop(&mut self) {
1102        let Some(expire_fn) = self.expire_fn.take() else {
1103            return;
1104        };
1105        let handle = match Handle::try_current() {
1106            Ok(x) => x,
1107            Err(_) => {
1108                warn!(
1109                    "WriteHandle {} dropped without being explicitly expired, falling back to lease timeout",
1110                    self.writer_id
1111                );
1112                return;
1113            }
1114        };
1115        // Spawn a best-effort task to expire this write handle. It's fine if
1116        // this doesn't run to completion, we'd just have to wait out the lease
1117        // before the shard-global since is unblocked.
1118        //
1119        // Intentionally create the span outside the task to set the parent.
1120        let expire_span = debug_span!("drop::expire");
1121        handle.spawn_named(
1122            || format!("WriteHandle::expire ({})", self.writer_id),
1123            expire_fn.0().instrument(expire_span),
1124        );
1125    }
1126}
1127
1128/// Ensure the given batch uses the given schema ID.
1129///
1130/// If the batch has no schema set, initialize it to the given one.
1131/// If the batch has a schema set, assert that it matches the given one.
1132fn ensure_batch_schema<T>(batch: &mut HollowBatch<T>, shard_id: ShardId, schema_id: SchemaId)
1133where
1134    T: Timestamp + Lattice + Codec64,
1135{
1136    let ensure = |id: &mut Option<SchemaId>| match id {
1137        Some(id) => assert_eq!(*id, schema_id, "schema ID mismatch; shard={shard_id}"),
1138        None => *id = Some(schema_id),
1139    };
1140
1141    for run_meta in &mut batch.run_meta {
1142        ensure(&mut run_meta.schema);
1143    }
1144    for part in &mut batch.parts {
1145        match part {
1146            RunPart::Single(BatchPart::Hollow(part)) => ensure(&mut part.schema_id),
1147            RunPart::Single(BatchPart::Inline { schema_id, .. }) => ensure(schema_id),
1148            RunPart::Many(_hollow_run_ref) => {
1149                // TODO: Fetch the parts in this run and rewrite them too. Alternatively, make
1150                // `run_meta` the only place we keep schema IDs, so rewriting parts isn't
1151                // necessary.
1152            }
1153        }
1154    }
1155}
1156
1157#[cfg(test)]
1158mod tests {
1159    use std::str::FromStr;
1160    use std::sync::mpsc;
1161
1162    use differential_dataflow::consolidation::consolidate_updates;
1163    use futures_util::FutureExt;
1164    use mz_dyncfg::ConfigUpdates;
1165    use mz_ore::collections::CollectionExt;
1166    use mz_ore::task;
1167    use serde_json::json;
1168
1169    use crate::cache::PersistClientCache;
1170    use crate::tests::{all_ok, new_test_client};
1171    use crate::{PersistLocation, ShardId};
1172
1173    use super::*;
1174
1175    #[mz_persist_proc::test(tokio::test)]
1176    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1177    async fn empty_batches(dyncfgs: ConfigUpdates) {
1178        let data = [
1179            (("1".to_owned(), "one".to_owned()), 1, 1),
1180            (("2".to_owned(), "two".to_owned()), 2, 1),
1181            (("3".to_owned(), "three".to_owned()), 3, 1),
1182        ];
1183
1184        let (mut write, _) = new_test_client(&dyncfgs)
1185            .await
1186            .expect_open::<String, String, u64, i64>(ShardId::new())
1187            .await;
1188        let blob = Arc::clone(&write.blob);
1189
1190        // Write an initial batch.
1191        let mut upper = 3;
1192        write.expect_append(&data[..2], vec![0], vec![upper]).await;
1193
1194        // Write a bunch of empty batches. This shouldn't write blobs, so the count should stay the same.
1195        let mut count_before = 0;
1196        blob.list_keys_and_metadata("", &mut |_| {
1197            count_before += 1;
1198        })
1199        .await
1200        .expect("list_keys failed");
1201        for _ in 0..5 {
1202            let new_upper = upper + 1;
1203            write.expect_compare_and_append(&[], upper, new_upper).await;
1204            upper = new_upper;
1205        }
1206        let mut count_after = 0;
1207        blob.list_keys_and_metadata("", &mut |_| {
1208            count_after += 1;
1209        })
1210        .await
1211        .expect("list_keys failed");
1212        assert_eq!(count_after, count_before);
1213    }
1214
1215    #[mz_persist_proc::test(tokio::test)]
1216    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1217    async fn compare_and_append_batch_multi(dyncfgs: ConfigUpdates) {
1218        let data0 = vec![
1219            (("1".to_owned(), "one".to_owned()), 1, 1),
1220            (("2".to_owned(), "two".to_owned()), 2, 1),
1221            (("4".to_owned(), "four".to_owned()), 4, 1),
1222        ];
1223        let data1 = vec![
1224            (("1".to_owned(), "one".to_owned()), 1, 1),
1225            (("2".to_owned(), "two".to_owned()), 2, 1),
1226            (("3".to_owned(), "three".to_owned()), 3, 1),
1227        ];
1228
1229        let (mut write, mut read) = new_test_client(&dyncfgs)
1230            .await
1231            .expect_open::<String, String, u64, i64>(ShardId::new())
1232            .await;
1233
1234        let mut batch0 = write.expect_batch(&data0, 0, 5).await;
1235        let mut batch1 = write.expect_batch(&data1, 0, 4).await;
1236
1237        write
1238            .expect_compare_and_append_batch(&mut [&mut batch0, &mut batch1], 0, 4)
1239            .await;
1240
1241        let batch = write
1242            .machine
1243            .unleased_snapshot(&Antichain::from_elem(3))
1244            .await
1245            .expect("just wrote this")
1246            .into_element();
1247
1248        assert!(batch.runs().count() >= 2);
1249
1250        let expected = vec![
1251            (("1".to_owned(), "one".to_owned()), 1, 2),
1252            (("2".to_owned(), "two".to_owned()), 2, 2),
1253            (("3".to_owned(), "three".to_owned()), 3, 1),
1254        ];
1255        let mut actual = read.expect_snapshot_and_fetch(3).await;
1256        consolidate_updates(&mut actual);
1257        assert_eq!(actual, all_ok(&expected, 3));
1258    }
1259
1260    #[mz_ore::test]
1261    fn writer_id_human_readable_serde() {
1262        #[derive(Debug, Serialize, Deserialize)]
1263        struct Container {
1264            writer_id: WriterId,
1265        }
1266
1267        // roundtrip through json
1268        let id = WriterId::from_str("w00000000-1234-5678-0000-000000000000").expect("valid id");
1269        assert_eq!(
1270            id,
1271            serde_json::from_value(serde_json::to_value(id.clone()).expect("serializable"))
1272                .expect("deserializable")
1273        );
1274
1275        // deserialize a serialized string directly
1276        assert_eq!(
1277            id,
1278            serde_json::from_str("\"w00000000-1234-5678-0000-000000000000\"")
1279                .expect("deserializable")
1280        );
1281
1282        // roundtrip id through a container type
1283        let json = json!({ "writer_id": id });
1284        assert_eq!(
1285            "{\"writer_id\":\"w00000000-1234-5678-0000-000000000000\"}",
1286            &json.to_string()
1287        );
1288        let container: Container = serde_json::from_value(json).expect("deserializable");
1289        assert_eq!(container.writer_id, id);
1290    }
1291
1292    #[mz_persist_proc::test(tokio::test)]
1293    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1294    async fn hollow_batch_roundtrip(dyncfgs: ConfigUpdates) {
1295        let data = vec![
1296            (("1".to_owned(), "one".to_owned()), 1, 1),
1297            (("2".to_owned(), "two".to_owned()), 2, 1),
1298            (("3".to_owned(), "three".to_owned()), 3, 1),
1299        ];
1300
1301        let (mut write, mut read) = new_test_client(&dyncfgs)
1302            .await
1303            .expect_open::<String, String, u64, i64>(ShardId::new())
1304            .await;
1305
1306        // This test is a bit more complex than it should be. It would be easier
1307        // if we could just compare the rehydrated batch to the original batch.
1308        // But a) turning a batch into a hollow batch consumes it, and b) Batch
1309        // doesn't have Eq/PartialEq.
1310        let batch = write.expect_batch(&data, 0, 4).await;
1311        let hollow_batch = batch.into_transmittable_batch();
1312        let mut rehydrated_batch = write.batch_from_transmittable_batch(hollow_batch);
1313
1314        write
1315            .expect_compare_and_append_batch(&mut [&mut rehydrated_batch], 0, 4)
1316            .await;
1317
1318        let expected = vec![
1319            (("1".to_owned(), "one".to_owned()), 1, 1),
1320            (("2".to_owned(), "two".to_owned()), 2, 1),
1321            (("3".to_owned(), "three".to_owned()), 3, 1),
1322        ];
1323        let mut actual = read.expect_snapshot_and_fetch(3).await;
1324        consolidate_updates(&mut actual);
1325        assert_eq!(actual, all_ok(&expected, 3));
1326    }
1327
1328    #[mz_persist_proc::test(tokio::test)]
1329    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1330    async fn wait_for_upper_past(dyncfgs: ConfigUpdates) {
1331        let client = new_test_client(&dyncfgs).await;
1332        let (mut write, _) = client.expect_open::<(), (), u64, i64>(ShardId::new()).await;
1333        let five = Antichain::from_elem(5);
1334
1335        // Upper is not past 5.
1336        assert_eq!(write.wait_for_upper_past(&five).now_or_never(), None);
1337
1338        // Upper is still not past 5.
1339        write
1340            .expect_compare_and_append(&[(((), ()), 1, 1)], 0, 5)
1341            .await;
1342        assert_eq!(write.wait_for_upper_past(&five).now_or_never(), None);
1343
1344        // Upper is past 5.
1345        write
1346            .expect_compare_and_append(&[(((), ()), 5, 1)], 5, 7)
1347            .await;
1348        assert_eq!(write.wait_for_upper_past(&five).now_or_never(), Some(()));
1349        assert_eq!(write.upper(), &Antichain::from_elem(7));
1350
1351        // Waiting for previous uppers does not regress the handle's cached
1352        // upper.
1353        assert_eq!(
1354            write
1355                .wait_for_upper_past(&Antichain::from_elem(2))
1356                .now_or_never(),
1357            Some(())
1358        );
1359        assert_eq!(write.upper(), &Antichain::from_elem(7));
1360    }
1361
1362    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
1363    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1364    async fn fetch_recent_upper_linearized() {
1365        type Timestamp = u64;
1366        let max_upper = 1000;
1367
1368        let shard_id = ShardId::new();
1369        let mut clients = PersistClientCache::new_no_metrics();
1370        let upper_writer_client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
1371        let (mut upper_writer, _) = upper_writer_client
1372            .expect_open::<(), (), Timestamp, i64>(shard_id)
1373            .await;
1374        // Clear the state cache between each client to maximally disconnect
1375        // them from each other.
1376        clients.clear_state_cache();
1377        let upper_reader_client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
1378        let (mut upper_reader, _) = upper_reader_client
1379            .expect_open::<(), (), Timestamp, i64>(shard_id)
1380            .await;
1381        let (tx, rx) = mpsc::channel();
1382
1383        let task = task::spawn(|| "upper-reader", async move {
1384            let mut upper = Timestamp::MIN;
1385
1386            while upper < max_upper {
1387                while let Ok(new_upper) = rx.try_recv() {
1388                    upper = new_upper;
1389                }
1390
1391                let recent_upper = upper_reader
1392                    .fetch_recent_upper()
1393                    .await
1394                    .as_option()
1395                    .cloned()
1396                    .expect("u64 is totally ordered and the shard is not finalized");
1397                assert!(
1398                    recent_upper >= upper,
1399                    "recent upper {recent_upper:?} is less than known upper {upper:?}"
1400                );
1401            }
1402        });
1403
1404        for upper in Timestamp::MIN..max_upper {
1405            let next_upper = upper + 1;
1406            upper_writer
1407                .expect_compare_and_append(&[], upper, next_upper)
1408                .await;
1409            tx.send(next_upper).expect("send failed");
1410        }
1411
1412        task.await;
1413    }
1414}