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