Skip to main content

mz_persist_client/
lib.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//! An abstraction presenting as a durable time-varying collection (aka shard)
11
12// The `fuzzing` feature re-exports internal types (see `fuzz_exports`) that are
13// intentionally undocumented. Don't require docs/Debug for them in that
14// test-only build. The normal public API is still linted.
15#![cfg_attr(
16    not(feature = "fuzzing"),
17    warn(missing_docs, missing_debug_implementations)
18)]
19// #[track_caller] is currently a no-op on async functions, but that hopefully won't be the case
20// forever. So we already annotate those functions now and ignore the compiler warning until
21// https://github.com/rust-lang/rust/issues/87417 pans out.
22#![allow(ungated_async_fn_track_caller)]
23
24use std::fmt::Debug;
25use std::marker::PhantomData;
26use std::sync::Arc;
27
28use differential_dataflow::difference::Monoid;
29use differential_dataflow::lattice::Lattice;
30use itertools::Itertools;
31use mz_build_info::{BuildInfo, build_info};
32use mz_dyncfg::ConfigSet;
33use mz_ore::instrument;
34use mz_persist::location::{Blob, Consensus, ExternalError};
35use mz_persist_types::schema::SchemaId;
36use mz_persist_types::{Codec, Codec64};
37use mz_proto::{IntoRustIfSome, ProtoType};
38use semver::Version;
39use timely::order::TotalOrder;
40use timely::progress::{Antichain, Timestamp};
41
42use crate::async_runtime::IsolatedRuntime;
43use crate::batch::{BATCH_DELETE_ENABLED, Batch, BatchBuilder, ProtoBatch};
44use crate::cache::{PersistClientCache, StateCache};
45use crate::cfg::PersistConfig;
46use crate::critical::{CriticalReaderId, Opaque, SinceHandle};
47use crate::error::InvalidUsage;
48use crate::fetch::{BatchFetcher, BatchFetcherConfig};
49use crate::internal::compact::{CompactConfig, Compactor};
50use crate::internal::encoding::parse_id;
51use crate::internal::gc::GarbageCollector;
52use crate::internal::machine::{Machine, retry_external};
53use crate::internal::state_versions::StateVersions;
54use crate::metrics::Metrics;
55use crate::read::{
56    Cursor, LazyPartStats, LeasedReaderId, READER_LEASE_DURATION, ReadHandle, Since,
57};
58use crate::rpc::PubSubSender;
59use crate::schema::CaESchema;
60use crate::write::{WriteHandle, WriterId};
61
62pub mod async_runtime;
63pub mod batch;
64pub mod cache;
65pub mod cfg;
66pub mod cli {
67    //! Persist command-line utilities
68    pub mod admin;
69    pub mod args;
70    pub mod bench;
71    pub mod inspect;
72}
73pub mod critical;
74pub mod error;
75pub mod fetch;
76pub mod internals_bench;
77pub mod iter;
78pub mod metrics {
79    //! Utilities related to metrics.
80    pub use crate::internal::metrics::{
81        BackpressureMetrics, Metrics, SinkMetrics, SinkWorkerMetrics, UpdateDelta, encode_ts_metric,
82    };
83}
84pub mod operators {
85    //! [timely] operators for reading and writing persist Shards.
86
87    use mz_dyncfg::{Config, ParameterScope};
88
89    pub mod shard_source;
90
91    // TODO(cfg): Move this next to the use.
92    pub(crate) const STORAGE_SOURCE_DECODE_FUEL: Config<usize> = Config::new(
93        "storage_source_decode_fuel",
94        100_000,
95        "\
96        The maximum amount of work to do in the persist_source mfp_and_decode \
97        operator before yielding.",
98        ParameterScope::Environment,
99    );
100}
101pub mod read;
102pub mod rpc;
103pub mod schema;
104pub mod stats;
105pub mod usage;
106pub mod write;
107
108/// Internal durable-state types re-exported under `cfg(feature = "fuzzing")` so
109/// the fuzz crate can drive their proto round-trips (`ProtoRollup`/`ProtoStateDiff`
110/// are decoded from blob/consensus on every state load). Not part of the public
111/// API.
112#[cfg(feature = "fuzzing")]
113pub mod fuzz_exports {
114    pub use crate::internal::encoding::Rollup;
115    pub use crate::internal::state::{ProtoRollup, ProtoStateDiff};
116    pub use crate::internal::state_diff::StateDiff;
117}
118
119/// An implementation of the public crate interface.
120mod internal {
121    pub mod apply;
122    pub mod cache;
123    pub mod compact;
124    pub mod encoding;
125    pub mod gc;
126    pub mod machine;
127    pub mod maintenance;
128    pub mod merge;
129    pub mod metrics;
130    pub mod paths;
131    pub mod restore;
132    pub mod service;
133    pub mod state;
134    pub mod state_diff;
135    pub mod state_versions;
136    pub mod trace;
137    pub mod watch;
138
139    #[cfg(test)]
140    pub mod datadriven;
141}
142
143/// Persist build information.
144pub const BUILD_INFO: BuildInfo = build_info!();
145
146// Re-export for convenience.
147pub use mz_persist_types::{PersistLocation, ShardId};
148
149pub use crate::internal::encoding::Schemas;
150
151/// Additional diagnostic information used within Persist
152/// e.g. for logging, metric labels, etc.
153#[derive(Clone, Debug)]
154pub struct Diagnostics {
155    /// A user-friendly name for the shard.
156    pub shard_name: String,
157    /// A purpose for the handle.
158    pub handle_purpose: String,
159}
160
161impl Diagnostics {
162    /// Create a new `Diagnostics` from `handle_purpose`.
163    pub fn from_purpose(handle_purpose: &str) -> Self {
164        Self {
165            shard_name: "unknown".to_string(),
166            handle_purpose: handle_purpose.to_string(),
167        }
168    }
169
170    /// Create a new `Diagnostics` for testing.
171    pub fn for_tests() -> Self {
172        Self {
173            shard_name: "test-shard-name".to_string(),
174            handle_purpose: "test-purpose".to_string(),
175        }
176    }
177}
178
179/// A handle for interacting with the set of persist shard made durable at a
180/// single [PersistLocation].
181///
182/// All async methods on PersistClient retry for as long as they are able, but
183/// the returned [std::future::Future]s implement "cancel on drop" semantics.
184/// This means that callers can add a timeout using [tokio::time::timeout] or
185/// [tokio::time::timeout_at].
186///
187/// ```rust,no_run
188/// # use std::sync::Arc;
189/// # use mz_persist_types::codec_impls::StringSchema;
190/// # let client: mz_persist_client::PersistClient = unimplemented!();
191/// # let timeout: std::time::Duration = unimplemented!();
192/// # let id = mz_persist_client::ShardId::new();
193/// # let diagnostics = mz_persist_client::Diagnostics { shard_name: "".into(), handle_purpose: "".into() };
194/// # async {
195/// tokio::time::timeout(timeout, client.open::<String, String, u64, i64>(id,
196///     Arc::new(StringSchema),Arc::new(StringSchema),diagnostics, true)).await
197/// # };
198/// ```
199#[derive(Debug, Clone)]
200pub struct PersistClient {
201    cfg: PersistConfig,
202    blob: Arc<dyn Blob>,
203    consensus: Arc<dyn Consensus>,
204    metrics: Arc<Metrics>,
205    isolated_runtime: Arc<IsolatedRuntime>,
206    shared_states: Arc<StateCache>,
207    pubsub_sender: Arc<dyn PubSubSender>,
208}
209
210impl PersistClient {
211    /// Returns a new client for interfacing with persist shards made durable to
212    /// the given [Blob] and [Consensus].
213    ///
214    /// This is exposed mostly for testing. Persist users likely want
215    /// [crate::cache::PersistClientCache::open].
216    pub fn new(
217        cfg: PersistConfig,
218        blob: Arc<dyn Blob>,
219        consensus: Arc<dyn Consensus>,
220        metrics: Arc<Metrics>,
221        isolated_runtime: Arc<IsolatedRuntime>,
222        shared_states: Arc<StateCache>,
223        pubsub_sender: Arc<dyn PubSubSender>,
224    ) -> Result<Self, ExternalError> {
225        // TODO: Verify somehow that blob matches consensus to prevent
226        // accidental misuse.
227        Ok(PersistClient {
228            cfg,
229            blob,
230            consensus,
231            metrics,
232            isolated_runtime,
233            shared_states,
234            pubsub_sender,
235        })
236    }
237
238    /// Returns a new in-mem [PersistClient] for tests and examples.
239    pub async fn new_for_tests() -> Self {
240        let cache = PersistClientCache::new_no_metrics();
241        cache
242            .open(PersistLocation::new_in_mem())
243            .await
244            .expect("in-mem location is valid")
245    }
246
247    /// Returns persist's [ConfigSet].
248    pub fn dyncfgs(&self) -> &ConfigSet {
249        &self.cfg.configs
250    }
251
252    async fn make_machine<K, V, T, D>(
253        &self,
254        shard_id: ShardId,
255        diagnostics: Diagnostics,
256    ) -> Result<Machine<K, V, T, D>, InvalidUsage<T>>
257    where
258        K: Debug + Codec,
259        V: Debug + Codec,
260        T: Timestamp + Lattice + Codec64 + Sync,
261        D: Monoid + Codec64 + Send + Sync,
262    {
263        let state_versions = StateVersions::new(
264            self.cfg.clone(),
265            Arc::clone(&self.consensus),
266            Arc::clone(&self.blob),
267            Arc::clone(&self.metrics),
268        );
269        let machine = Machine::<K, V, T, D>::new(
270            self.cfg.clone(),
271            shard_id,
272            Arc::clone(&self.metrics),
273            Arc::new(state_versions),
274            Arc::clone(&self.shared_states),
275            Arc::clone(&self.pubsub_sender),
276            Arc::clone(&self.isolated_runtime),
277            diagnostics.clone(),
278        )
279        .await?;
280        Ok(machine)
281    }
282
283    /// Provides capabilities for the durable TVC identified by `shard_id` at
284    /// its current since and upper frontiers.
285    ///
286    /// This method is a best-effort attempt to regain control of the frontiers
287    /// of a shard. Its most common uses are to recover capabilities that have
288    /// expired (leases) or to attempt to read a TVC that one did not create (or
289    /// otherwise receive capabilities for). If the frontiers have been fully
290    /// released by all other parties, this call may result in capabilities with
291    /// empty frontiers (which are useless).
292    ///
293    /// If `shard_id` has never been used before, initializes a new shard and
294    /// returns handles with `since` and `upper` frontiers set to initial values
295    /// of `Antichain::from_elem(T::minimum())`.
296    ///
297    /// The `schema` parameter is currently unused, but should be an object
298    /// that represents the schema of the data in the shard. This will be required
299    /// in the future.
300    #[instrument(level = "debug", fields(shard = %shard_id))]
301    pub async fn open<K, V, T, D>(
302        &self,
303        shard_id: ShardId,
304        key_schema: Arc<K::Schema>,
305        val_schema: Arc<V::Schema>,
306        diagnostics: Diagnostics,
307        use_critical_since: bool,
308    ) -> Result<(WriteHandle<K, V, T, D>, ReadHandle<K, V, T, D>), InvalidUsage<T>>
309    where
310        K: Debug + Codec,
311        V: Debug + Codec,
312        T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
313        D: Monoid + Ord + Codec64 + Send + Sync,
314    {
315        Ok((
316            self.open_writer(
317                shard_id,
318                Arc::clone(&key_schema),
319                Arc::clone(&val_schema),
320                diagnostics.clone(),
321            )
322            .await?,
323            self.open_leased_reader(
324                shard_id,
325                key_schema,
326                val_schema,
327                diagnostics,
328                use_critical_since,
329            )
330            .await?,
331        ))
332    }
333
334    /// [Self::open], but returning only a [ReadHandle].
335    ///
336    /// Use this to save latency and a bit of persist traffic if you're just
337    /// going to immediately drop or expire the [WriteHandle].
338    ///
339    /// The `_schema` parameter is currently unused, but should be an object
340    /// that represents the schema of the data in the shard. This will be required
341    /// in the future.
342    #[instrument(level = "debug", fields(shard = %shard_id))]
343    pub async fn open_leased_reader<K, V, T, D>(
344        &self,
345        shard_id: ShardId,
346        key_schema: Arc<K::Schema>,
347        val_schema: Arc<V::Schema>,
348        diagnostics: Diagnostics,
349        use_critical_since: bool,
350    ) -> Result<ReadHandle<K, V, T, D>, InvalidUsage<T>>
351    where
352        K: Debug + Codec,
353        V: Debug + Codec,
354        T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
355        D: Monoid + Codec64 + Send + Sync,
356    {
357        let machine = self.make_machine(shard_id, diagnostics.clone()).await?;
358        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
359
360        let reader_id = LeasedReaderId::new();
361        let (reader_state, maintenance) = machine
362            .register_leased_reader(
363                &reader_id,
364                &diagnostics.handle_purpose,
365                READER_LEASE_DURATION.get(&self.cfg),
366                use_critical_since,
367            )
368            .await;
369        maintenance.start_performing(&machine, &gc);
370        let schemas = Schemas {
371            id: None,
372            key: key_schema,
373            val: val_schema,
374        };
375        let reader = ReadHandle::new(
376            self.cfg.clone(),
377            Arc::clone(&self.metrics),
378            machine,
379            gc,
380            Arc::clone(&self.blob),
381            reader_id,
382            schemas,
383            reader_state,
384        )
385        .await;
386
387        Ok(reader)
388    }
389
390    /// Creates and returns a [BatchFetcher] for the given shard id.
391    #[instrument(level = "debug", fields(shard = %shard_id))]
392    pub async fn create_batch_fetcher<K, V, T, D>(
393        &self,
394        shard_id: ShardId,
395        key_schema: Arc<K::Schema>,
396        val_schema: Arc<V::Schema>,
397        is_transient: bool,
398        diagnostics: Diagnostics,
399    ) -> Result<BatchFetcher<K, V, T, D>, InvalidUsage<T>>
400    where
401        K: Debug + Codec,
402        V: Debug + Codec,
403        T: Timestamp + Lattice + Codec64 + Sync,
404        D: Monoid + Codec64 + Send + Sync,
405    {
406        let machine = self.make_machine(shard_id, diagnostics.clone()).await?;
407        let read_schemas = Schemas {
408            id: None,
409            key: key_schema,
410            val: val_schema,
411        };
412        let schema_cache = machine.applier.schema_cache();
413        let fetcher = BatchFetcher {
414            cfg: BatchFetcherConfig::new(&self.cfg),
415            blob: Arc::clone(&self.blob),
416            metrics: Arc::clone(&self.metrics),
417            shard_metrics: Arc::clone(&machine.applier.shard_metrics),
418            shard_id,
419            read_schemas,
420            schema_cache,
421            is_transient,
422            _phantom: PhantomData,
423        };
424
425        Ok(fetcher)
426    }
427
428    /// A convenience [CriticalReaderId] for Materialize controllers.
429    ///
430    /// For most (soon to be all?) shards in Materialize, a centralized
431    /// "controller" is the authority for when a user no longer needs to read at
432    /// a given frontier. (Other uses are temporary holds where correctness of
433    /// the overall system can be maintained through a lease timeout.) To make
434    /// [SinceHandle] easier to work with, we offer this convenience id for
435    /// Materialize controllers, so they don't have to durably record it.
436    ///
437    /// TODO: We're still shaking out whether the controller should be the only
438    /// critical since hold or if there are other places we want them. If the
439    /// former, we should remove [CriticalReaderId] and bake in the singular
440    /// nature of the controller critical handle.
441    ///
442    /// ```rust
443    /// // This prints as something that is not 0 but is visually recognizable.
444    /// assert_eq!(
445    ///     mz_persist_client::PersistClient::CONTROLLER_CRITICAL_SINCE.to_string(),
446    ///     "c00000000-1111-2222-3333-444444444444",
447    /// )
448    /// ```
449    pub const CONTROLLER_CRITICAL_SINCE: CriticalReaderId =
450        CriticalReaderId([0, 0, 0, 0, 17, 17, 34, 34, 51, 51, 68, 68, 68, 68, 68, 68]);
451
452    /// Provides a capability for the durable TVC identified by `shard_id` at
453    /// its current since frontier.
454    ///
455    /// In contrast to the time-leased [ReadHandle] returned by [Self::open] and
456    /// [Self::open_leased_reader], this handle and its associated capability
457    /// are not leased. A [SinceHandle] does not release its since capability;
458    /// downgrade to the empty antichain to hold back the since.
459    /// Also unlike `ReadHandle`, the handle is not expired on drop.
460    /// This is less ergonomic, but useful for "critical" since
461    /// holds which must survive even lease timeouts.
462    ///
463    /// **IMPORTANT**: The above means that if a SinceHandle is registered and
464    /// then lost, the shard's since will be permanently "stuck", forever
465    /// preventing logical compaction. Users are advised to durably record
466    /// (preferably in code) the intended [CriticalReaderId] _before_ registering
467    /// a SinceHandle (in case the process crashes at the wrong time).
468    ///
469    /// If `shard_id` has never been used before, initializes a new shard and
470    /// return a handle with its `since` frontier set to the initial value of
471    /// `Antichain::from_elem(T::minimum())`.
472    #[instrument(level = "debug", fields(shard = %shard_id))]
473    pub async fn open_critical_since<K, V, T, D>(
474        &self,
475        shard_id: ShardId,
476        reader_id: CriticalReaderId,
477        default_opaque: Opaque,
478        diagnostics: Diagnostics,
479    ) -> Result<SinceHandle<K, V, T, D>, InvalidUsage<T>>
480    where
481        K: Debug + Codec,
482        V: Debug + Codec,
483        T: Timestamp + Lattice + Codec64 + Sync,
484        D: Monoid + Codec64 + Send + Sync,
485    {
486        let machine = self.make_machine(shard_id, diagnostics.clone()).await?;
487        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
488
489        let (state, maintenance) = machine
490            .register_critical_reader(&reader_id, default_opaque, &diagnostics.handle_purpose)
491            .await;
492        maintenance.start_performing(&machine, &gc);
493        let handle = SinceHandle::new(machine, gc, reader_id, state.since, state.opaque);
494
495        Ok(handle)
496    }
497
498    /// [Self::open], but returning only a [WriteHandle].
499    ///
500    /// Use this to save latency and a bit of persist traffic if you're just
501    /// going to immediately drop or expire the [ReadHandle].
502    #[instrument(level = "debug", fields(shard = %shard_id))]
503    pub async fn open_writer<K, V, T, D>(
504        &self,
505        shard_id: ShardId,
506        key_schema: Arc<K::Schema>,
507        val_schema: Arc<V::Schema>,
508        diagnostics: Diagnostics,
509    ) -> Result<WriteHandle<K, V, T, D>, InvalidUsage<T>>
510    where
511        K: Debug + Codec,
512        V: Debug + Codec,
513        T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
514        D: Monoid + Ord + Codec64 + Send + Sync,
515    {
516        let machine = self.make_machine(shard_id, diagnostics.clone()).await?;
517        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
518
519        // We defer registering the schema until write time, to allow opening
520        // write handles in a "read-only" mode where they don't implicitly
521        // modify persist state. But it might already be registered, in which
522        // case we can fetch its ID.
523        let schema_id = machine.find_schema(&*key_schema, &*val_schema);
524
525        let writer_id = WriterId::new();
526        let schemas = Schemas {
527            id: schema_id,
528            key: key_schema,
529            val: val_schema,
530        };
531        let writer = WriteHandle::new(
532            self.cfg.clone(),
533            Arc::clone(&self.metrics),
534            machine,
535            gc,
536            Arc::clone(&self.blob),
537            writer_id,
538            &diagnostics.handle_purpose,
539            schemas,
540        );
541        Ok(writer)
542    }
543
544    /// Returns a [BatchBuilder] that can be used to write a batch of updates to
545    /// blob storage which can then be appended to the given shard using
546    /// [WriteHandle::compare_and_append_batch] or [WriteHandle::append_batch],
547    /// or which can be read using [PersistClient::read_batches_consolidated].
548    ///
549    /// The builder uses a bounded amount of memory, even when the number of
550    /// updates is very large. Individual records, however, should be small
551    /// enough that we can reasonably chunk them up: O(KB) is definitely fine,
552    /// O(MB) come talk to us.
553    ///
554    /// `max_runs` overrides `persist_batch_max_runs` when `Some`. With `None`, the builder honors
555    /// the configured value, as [WriteHandle::builder] does.
556    #[instrument(level = "debug", fields(shard = %shard_id))]
557    pub fn batch_builder<K, V, T, D>(
558        &self,
559        shard_id: ShardId,
560        shard_name: &str,
561        write_schemas: Schemas<K, V>,
562        lower: Antichain<T>,
563        max_runs: Option<usize>,
564    ) -> BatchBuilder<K, V, T, D>
565    where
566        K: Debug + Codec,
567        V: Debug + Codec,
568        T: Timestamp + Lattice + Codec64 + TotalOrder + Sync,
569        D: Monoid + Ord + Codec64 + Send + Sync,
570    {
571        let mut compact_cfg = CompactConfig::new(&self.cfg, shard_id);
572        if let Some(max_runs) = max_runs {
573            compact_cfg.batch.max_runs = Some(max_runs);
574        }
575        WriteHandle::builder_inner(
576            &self.cfg,
577            compact_cfg,
578            Arc::clone(&self.metrics),
579            self.metrics.shards.shard(&shard_id, shard_name),
580            &self.metrics.user,
581            Arc::clone(&self.isolated_runtime),
582            Arc::clone(&self.blob),
583            shard_id,
584            write_schemas,
585            lower,
586        )
587    }
588
589    /// Turns the given [`ProtoBatch`] back into a [`Batch`] which can be used
590    /// to append it to the given shard or to read it via
591    /// [PersistClient::read_batches_consolidated]
592    ///
593    /// CAUTION: This API allows turning a [ProtoBatch] into a [Batch] multiple
594    /// times, but if a batch is deleted the backing data goes away, so at that
595    /// point all in-memory copies of a batch become invalid and cannot be read
596    /// anymore.
597    pub fn batch_from_transmittable_batch<K, V, T, D>(
598        &self,
599        shard_id: &ShardId,
600        batch: ProtoBatch,
601    ) -> Batch<K, V, T, D>
602    where
603        K: Debug + Codec,
604        V: Debug + Codec,
605        T: Timestamp + Lattice + Codec64 + Sync,
606        D: Monoid + Ord + Codec64 + Send + Sync,
607    {
608        let batch_shard_id: ShardId = batch
609            .shard_id
610            .into_rust()
611            .expect("valid transmittable batch");
612        assert_eq!(&batch_shard_id, shard_id);
613
614        let shard_metrics = self.metrics.shards.shard(shard_id, "peek_stash");
615
616        let ret = Batch {
617            batch_delete_enabled: BATCH_DELETE_ENABLED.get(&self.cfg),
618            metrics: Arc::clone(&self.metrics),
619            shard_metrics,
620            version: Version::parse(&batch.version).expect("valid transmittable batch"),
621            schemas: (batch.key_schema, batch.val_schema),
622            batch: batch
623                .batch
624                .into_rust_if_some("ProtoBatch::batch")
625                .expect("valid transmittable batch"),
626            blob: Arc::clone(&self.blob),
627            _phantom: std::marker::PhantomData,
628        };
629
630        assert_eq!(&ret.shard_id(), shard_id);
631        ret
632    }
633
634    /// Returns a [Cursor] for reading the given batches. Yielded updates are
635    /// consolidated if the given batches contain sorted runs, which is true
636    /// when they have been written using a [BatchBuilder].
637    ///
638    /// To keep memory usage down when reading a snapshot that consolidates
639    /// well, this consolidates as it goes. However, note that only the
640    /// serialized data is consolidated: the deserialized data will only be
641    /// consolidated if your K/V codecs are one-to-one.
642    ///
643    /// CAUTION: The caller needs to make sure that the given batches are
644    /// readable and they have to remain readable for the lifetime of the
645    /// returned [Cursor]. The caller is also responsible for the lifecycle of
646    /// the batches: once the cursor and the batches are no longer needed you
647    /// must call [Cursor::into_lease] to get back the batches and delete them.
648    #[allow(clippy::unused_async)]
649    pub async fn read_batches_consolidated<K, V, T, D>(
650        &mut self,
651        shard_id: ShardId,
652        as_of: Antichain<T>,
653        read_schemas: Schemas<K, V>,
654        batches: Vec<Batch<K, V, T, D>>,
655        should_fetch_part: impl for<'a> Fn(Option<&'a LazyPartStats>) -> bool,
656        memory_budget_bytes: usize,
657    ) -> Result<Cursor<K, V, T, D, Vec<Batch<K, V, T, D>>>, Since<T>>
658    where
659        K: Debug + Codec + Ord,
660        V: Debug + Codec + Ord,
661        T: Timestamp + Lattice + Codec64 + TotalOrder + Sync,
662        D: Monoid + Ord + Codec64 + Send + Sync,
663    {
664        let shard_metrics = self.metrics.shards.shard(&shard_id, "peek_stash");
665
666        let hollow_batches = batches.iter().map(|b| b.batch.clone()).collect_vec();
667
668        ReadHandle::read_batches_consolidated(
669            &self.cfg,
670            Arc::clone(&self.metrics),
671            shard_metrics,
672            self.metrics.read.snapshot.clone(),
673            Arc::clone(&self.blob),
674            shard_id,
675            as_of,
676            read_schemas,
677            &hollow_batches,
678            batches,
679            should_fetch_part,
680            memory_budget_bytes,
681        )
682    }
683
684    /// Returns the requested schema, if known at the current state.
685    pub async fn get_schema<K, V, T, D>(
686        &self,
687        shard_id: ShardId,
688        schema_id: SchemaId,
689        diagnostics: Diagnostics,
690    ) -> Result<Option<(K::Schema, V::Schema)>, InvalidUsage<T>>
691    where
692        K: Debug + Codec,
693        V: Debug + Codec,
694        T: Timestamp + Lattice + Codec64 + Sync,
695        D: Monoid + Codec64 + Send + Sync,
696    {
697        let machine = self
698            .make_machine::<K, V, T, D>(shard_id, diagnostics)
699            .await?;
700        Ok(machine.get_schema(schema_id))
701    }
702
703    /// Returns the latest schema registered at the current state.
704    pub async fn latest_schema<K, V, T, D>(
705        &self,
706        shard_id: ShardId,
707        diagnostics: Diagnostics,
708    ) -> Result<Option<(SchemaId, K::Schema, V::Schema)>, InvalidUsage<T>>
709    where
710        K: Debug + Codec,
711        V: Debug + Codec,
712        T: Timestamp + Lattice + Codec64 + Sync,
713        D: Monoid + Codec64 + Send + Sync,
714    {
715        let machine = self
716            .make_machine::<K, V, T, D>(shard_id, diagnostics)
717            .await?;
718        Ok(machine.latest_schema())
719    }
720
721    /// Fetches and returns a recent shard-global `upper`, without requiring a
722    /// [`WriteHandle`].
723    ///
724    /// Importantly, this operation is linearized with write operations, giving
725    /// the same guarantee as [`WriteHandle::fetch_recent_upper`]. It requires
726    /// fetching the latest state from consensus and is therefore a potentially
727    /// expensive operation.
728    ///
729    /// If `shard_id` has never been used before, initializes the shard and
730    /// returns an upper of `Antichain::from_elem(T::minimum())`.
731    pub async fn recent_upper<K, V, T, D>(
732        &self,
733        shard_id: ShardId,
734        diagnostics: Diagnostics,
735    ) -> Result<Antichain<T>, InvalidUsage<T>>
736    where
737        K: Debug + Codec,
738        V: Debug + Codec,
739        T: Timestamp + Lattice + Codec64 + Sync,
740        D: Monoid + Codec64 + Send + Sync,
741    {
742        let machine = self
743            .make_machine::<K, V, T, D>(shard_id, diagnostics)
744            .await?;
745        Ok(machine.applier.fetch_upper(|upper| upper.clone()).await)
746    }
747
748    /// Registers a schema for the given shard.
749    ///
750    /// Returns the new schema ID if the registration succeeds, and `None`
751    /// otherwise. Schema registration succeeds in two cases:
752    ///  a) No schema was currently registered for the shard.
753    ///  b) The given schema is already registered for the shard.
754    ///
755    /// To evolve an existing schema instead, use
756    /// [PersistClient::compare_and_evolve_schema].
757    //
758    // TODO: unify with `compare_and_evolve_schema`
759    pub async fn register_schema<K, V, T, D>(
760        &self,
761        shard_id: ShardId,
762        key_schema: &K::Schema,
763        val_schema: &V::Schema,
764        diagnostics: Diagnostics,
765    ) -> Result<Option<SchemaId>, InvalidUsage<T>>
766    where
767        K: Debug + Codec,
768        V: Debug + Codec,
769        T: Timestamp + Lattice + Codec64 + Sync,
770        D: Monoid + Codec64 + Send + Sync,
771    {
772        let machine = self
773            .make_machine::<K, V, T, D>(shard_id, diagnostics)
774            .await?;
775        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
776
777        let (schema_id, maintenance) = machine.register_schema(key_schema, val_schema).await;
778        maintenance.start_performing(&machine, &gc);
779
780        Ok(schema_id)
781    }
782
783    /// Registers a new latest schema for the given shard.
784    ///
785    /// This new schema must be [backward_compatible] with all previous schemas
786    /// for this shard. If it's not, [CaESchema::Incompatible] is returned.
787    ///
788    /// [backward_compatible]: mz_persist_types::schema::backward_compatible
789    ///
790    /// To prevent races, the caller must declare what it believes to be the
791    /// latest schema id. A request that already matches the current schema
792    /// returns [CaESchema::Ok] even if `expected` is stale, so retried or
793    /// duplicated evolutions are idempotent. [CaESchema::ExpectedMismatch]
794    /// means a genuine divergence: the current schema matches neither the
795    /// request nor `expected`.
796    pub async fn compare_and_evolve_schema<K, V, T, D>(
797        &self,
798        shard_id: ShardId,
799        expected: SchemaId,
800        key_schema: &K::Schema,
801        val_schema: &V::Schema,
802        diagnostics: Diagnostics,
803    ) -> Result<CaESchema<K, V>, InvalidUsage<T>>
804    where
805        K: Debug + Codec,
806        V: Debug + Codec,
807        T: Timestamp + Lattice + Codec64 + Sync,
808        D: Monoid + Codec64 + Send + Sync,
809    {
810        let machine = self
811            .make_machine::<K, V, T, D>(shard_id, diagnostics)
812            .await?;
813        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
814        let (res, maintenance) = machine
815            .compare_and_evolve_schema(expected, key_schema, val_schema)
816            .await;
817        maintenance.start_performing(&machine, &gc);
818        Ok(res)
819    }
820
821    /// Check if the given shard is in a finalized state; ie. it can no longer be
822    /// read, any data that was written to it is no longer accessible, and we've
823    /// discarded references to that data from state.
824    pub async fn is_finalized<K, V, T, D>(
825        &self,
826        shard_id: ShardId,
827        diagnostics: Diagnostics,
828    ) -> Result<bool, InvalidUsage<T>>
829    where
830        K: Debug + Codec,
831        V: Debug + Codec,
832        T: Timestamp + Lattice + Codec64 + Sync,
833        D: Monoid + Codec64 + Send + Sync,
834    {
835        let machine = self
836            .make_machine::<K, V, T, D>(shard_id, diagnostics)
837            .await?;
838        Ok(machine.is_finalized())
839    }
840
841    /// If a shard is guaranteed to never be used again, finalize it to delete
842    /// the associated data and release any associated resources. (Except for a
843    /// little state in consensus we use to represent the tombstone.)
844    ///
845    /// The caller should ensure that both the `since` and `upper` of the shard
846    /// have been advanced to `[]`: ie. the shard is no longer writable or readable.
847    /// Otherwise an error is returned.
848    ///
849    /// Once `finalize_shard` has been called, the result of future operations on
850    /// the shard are not defined. They may return errors or succeed as a noop.
851    #[instrument(level = "debug", fields(shard = %shard_id))]
852    pub async fn finalize_shard<K, V, T, D>(
853        &self,
854        shard_id: ShardId,
855        diagnostics: Diagnostics,
856    ) -> Result<(), InvalidUsage<T>>
857    where
858        K: Debug + Codec,
859        V: Debug + Codec,
860        T: Timestamp + Lattice + Codec64 + Sync,
861        D: Monoid + Codec64 + Send + Sync,
862    {
863        let machine = self
864            .make_machine::<K, V, T, D>(shard_id, diagnostics)
865            .await?;
866
867        let maintenance = machine.become_tombstone().await?;
868        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
869
870        let () = maintenance.perform(&machine, &gc).await;
871
872        Ok(())
873    }
874
875    /// Upgrade the state to the latest version. This should only be called once we will no longer
876    /// need to interoperate with older versions, like after a successful upgrade.
877    pub async fn upgrade_version<K, V, T, D>(
878        &self,
879        shard_id: ShardId,
880        diagnostics: Diagnostics,
881    ) -> Result<(), InvalidUsage<T>>
882    where
883        K: Debug + Codec,
884        V: Debug + Codec,
885        T: Timestamp + Lattice + Codec64 + Sync,
886        D: Monoid + Codec64 + Send + Sync,
887    {
888        let machine = self
889            .make_machine::<K, V, T, D>(shard_id, diagnostics)
890            .await?;
891
892        match machine.upgrade_version().await {
893            Ok(maintenance) => {
894                let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
895                let () = maintenance.perform(&machine, &gc).await;
896                Ok(())
897            }
898            Err(version) => Err(InvalidUsage::IncompatibleVersion { version }),
899        }
900    }
901
902    /// Returns the internal state of the shard for debugging and QA.
903    ///
904    /// We'll be thoughtful about making unnecessary changes, but the **output
905    /// of this method needs to be gated from users**, so that it's not subject
906    /// to our backward compatibility guarantees.
907    pub async fn inspect_shard<T: Timestamp + Lattice + Codec64>(
908        &self,
909        shard_id: &ShardId,
910    ) -> Result<impl serde::Serialize, anyhow::Error> {
911        let state_versions = StateVersions::new(
912            self.cfg.clone(),
913            Arc::clone(&self.consensus),
914            Arc::clone(&self.blob),
915            Arc::clone(&self.metrics),
916        );
917        // TODO: Don't fetch all live diffs. Feels like we should pull out a new
918        // method in StateVersions for fetching the latest version of State of a
919        // shard that might or might not exist.
920        let versions = state_versions.fetch_all_live_diffs(shard_id).await;
921        if versions.is_empty() {
922            return Err(anyhow::anyhow!("{} does not exist", shard_id));
923        }
924        let state = state_versions
925            .fetch_current_state::<T>(shard_id, versions)
926            .await;
927        let state = state.check_ts_codec(shard_id)?;
928        Ok(state)
929    }
930
931    /// Test helper for a [Self::open] call that is expected to succeed.
932    #[cfg(test)]
933    #[track_caller]
934    pub async fn expect_open<K, V, T, D>(
935        &self,
936        shard_id: ShardId,
937    ) -> (WriteHandle<K, V, T, D>, ReadHandle<K, V, T, D>)
938    where
939        K: Debug + Codec,
940        V: Debug + Codec,
941        T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
942        D: Monoid + Ord + Codec64 + Send + Sync,
943        K::Schema: Default,
944        V::Schema: Default,
945    {
946        self.open(
947            shard_id,
948            Arc::new(K::Schema::default()),
949            Arc::new(V::Schema::default()),
950            Diagnostics::for_tests(),
951            true,
952        )
953        .await
954        .expect("codec mismatch")
955    }
956
957    /// Return the metrics being used by this client.
958    ///
959    /// Only exposed for tests, persistcli, and benchmarks.
960    pub fn metrics(&self) -> &Arc<Metrics> {
961        &self.metrics
962    }
963}
964
965#[cfg(test)]
966mod tests {
967    use std::future::Future;
968    use std::pin::Pin;
969    use std::task::Context;
970    use std::time::Duration;
971
972    use differential_dataflow::consolidation::consolidate_updates;
973    use differential_dataflow::lattice::Lattice;
974    use futures_task::noop_waker;
975    use mz_dyncfg::ConfigUpdates;
976    use mz_ore::assert_ok;
977    use mz_persist::indexed::encoding::BlobTraceBatchPart;
978    use mz_persist::workload::DataGenerator;
979    use mz_persist_types::codec_impls::{StringSchema, VecU8Schema};
980    use mz_proto::protobuf_roundtrip;
981    use proptest::prelude::*;
982    use timely::order::PartialOrder;
983    use timely::progress::Antichain;
984
985    use crate::batch::BLOB_TARGET_SIZE;
986    use crate::cache::PersistClientCache;
987    use crate::cfg::BATCH_BUILDER_MAX_OUTSTANDING_PARTS;
988    use crate::critical::Opaque;
989    use crate::error::{CodecConcreteType, CodecMismatch, UpperMismatch};
990    use crate::internal::paths::BlobKey;
991    use crate::read::ListenEvent;
992
993    use super::*;
994
995    pub fn new_test_client_cache(dyncfgs: &ConfigUpdates) -> PersistClientCache {
996        // Configure an aggressively small blob_target_size so we get some
997        // amount of coverage of that in tests. Similarly, for max_outstanding.
998        let mut cache = PersistClientCache::new_no_metrics();
999        cache.cfg.set_config(&BLOB_TARGET_SIZE, 10);
1000        cache
1001            .cfg
1002            .set_config(&BATCH_BUILDER_MAX_OUTSTANDING_PARTS, 1);
1003        dyncfgs.apply(cache.cfg());
1004
1005        // Enable compaction in tests to ensure we get coverage.
1006        cache.cfg.compaction_enabled = true;
1007        cache
1008    }
1009
1010    pub async fn new_test_client(dyncfgs: &ConfigUpdates) -> PersistClient {
1011        let cache = new_test_client_cache(dyncfgs);
1012        cache
1013            .open(PersistLocation::new_in_mem())
1014            .await
1015            .expect("client construction failed")
1016    }
1017
1018    pub fn all_ok<'a, K, V, T, D, I>(iter: I, as_of: T) -> Vec<((K, V), T, D)>
1019    where
1020        K: Ord + Clone + 'a,
1021        V: Ord + Clone + 'a,
1022        T: Timestamp + Lattice + Clone + 'a,
1023        D: Monoid + Clone + 'a,
1024        I: IntoIterator<Item = &'a ((K, V), T, D)>,
1025    {
1026        let as_of = Antichain::from_elem(as_of);
1027        let mut ret = iter
1028            .into_iter()
1029            .map(|((k, v), t, d)| {
1030                let mut t = t.clone();
1031                t.advance_by(as_of.borrow());
1032                ((k.clone(), v.clone()), t, d.clone())
1033            })
1034            .collect();
1035        consolidate_updates(&mut ret);
1036        ret
1037    }
1038
1039    pub async fn expect_fetch_part<K, V, T, D>(
1040        blob: &dyn Blob,
1041        key: &BlobKey,
1042        metrics: &Metrics,
1043        read_schemas: &Schemas<K, V>,
1044    ) -> (BlobTraceBatchPart<T>, Vec<((K, V), T, D)>)
1045    where
1046        K: Codec + Clone,
1047        V: Codec + Clone,
1048        T: Timestamp + Codec64,
1049        D: Codec64,
1050    {
1051        let value = blob
1052            .get(key)
1053            .await
1054            .expect("failed to fetch part")
1055            .expect("missing part");
1056        let mut part =
1057            BlobTraceBatchPart::decode(&value, &metrics.columnar).expect("failed to decode part");
1058        let structured = part
1059            .updates
1060            .into_part::<K, V>(&*read_schemas.key, &*read_schemas.val);
1061        let updates = structured
1062            .decode_iter::<K, V, T, D>(&*read_schemas.key, &*read_schemas.val)
1063            .expect("structured data")
1064            .collect();
1065        (part, updates)
1066    }
1067
1068    #[mz_persist_proc::test(tokio::test)]
1069    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1070    async fn sanity_check(dyncfgs: ConfigUpdates) {
1071        let data = [
1072            (("1".to_owned(), "one".to_owned()), 1, 1),
1073            (("2".to_owned(), "two".to_owned()), 2, 1),
1074            (("3".to_owned(), "three".to_owned()), 3, 1),
1075        ];
1076
1077        let (mut write, mut read) = new_test_client(&dyncfgs)
1078            .await
1079            .expect_open::<String, String, u64, i64>(ShardId::new())
1080            .await;
1081        assert_eq!(write.upper(), &Antichain::from_elem(u64::minimum()));
1082        assert_eq!(read.since(), &Antichain::from_elem(u64::minimum()));
1083
1084        // Write a [0,3) batch.
1085        write
1086            .expect_append(&data[..2], write.upper().clone(), vec![3])
1087            .await;
1088        assert_eq!(write.upper(), &Antichain::from_elem(3));
1089
1090        // Grab a snapshot and listener as_of 1. Snapshot should only have part of what we wrote.
1091        assert_eq!(
1092            read.expect_snapshot_and_fetch(1).await,
1093            all_ok(&data[..1], 1)
1094        );
1095
1096        let mut listen = read.clone("").await.expect_listen(1).await;
1097
1098        // Write a [3,4) batch.
1099        write
1100            .expect_append(&data[2..], write.upper().clone(), vec![4])
1101            .await;
1102        assert_eq!(write.upper(), &Antichain::from_elem(4));
1103
1104        // Listen should have part of the initial write plus the new one.
1105        assert_eq!(
1106            listen.read_until(&4).await,
1107            (all_ok(&data[1..], 1), Antichain::from_elem(4))
1108        );
1109
1110        // Downgrading the since is tracked locally (but otherwise is a no-op).
1111        read.downgrade_since(&Antichain::from_elem(2)).await;
1112        assert_eq!(read.since(), &Antichain::from_elem(2));
1113    }
1114
1115    // Sanity check that the open_reader and open_writer calls work.
1116    #[mz_persist_proc::test(tokio::test)]
1117    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1118    async fn open_reader_writer(dyncfgs: ConfigUpdates) {
1119        let data = vec![
1120            (("1".to_owned(), "one".to_owned()), 1, 1),
1121            (("2".to_owned(), "two".to_owned()), 2, 1),
1122            (("3".to_owned(), "three".to_owned()), 3, 1),
1123        ];
1124
1125        let shard_id = ShardId::new();
1126        let client = new_test_client(&dyncfgs).await;
1127        let mut write1 = client
1128            .open_writer::<String, String, u64, i64>(
1129                shard_id,
1130                Arc::new(StringSchema),
1131                Arc::new(StringSchema),
1132                Diagnostics::for_tests(),
1133            )
1134            .await
1135            .expect("codec mismatch");
1136        let mut read1 = client
1137            .open_leased_reader::<String, String, u64, i64>(
1138                shard_id,
1139                Arc::new(StringSchema),
1140                Arc::new(StringSchema),
1141                Diagnostics::for_tests(),
1142                true,
1143            )
1144            .await
1145            .expect("codec mismatch");
1146        let mut read2 = client
1147            .open_leased_reader::<String, String, u64, i64>(
1148                shard_id,
1149                Arc::new(StringSchema),
1150                Arc::new(StringSchema),
1151                Diagnostics::for_tests(),
1152                true,
1153            )
1154            .await
1155            .expect("codec mismatch");
1156        let mut write2 = client
1157            .open_writer::<String, String, u64, i64>(
1158                shard_id,
1159                Arc::new(StringSchema),
1160                Arc::new(StringSchema),
1161                Diagnostics::for_tests(),
1162            )
1163            .await
1164            .expect("codec mismatch");
1165
1166        write2.expect_compare_and_append(&data[..1], 0, 2).await;
1167        assert_eq!(
1168            read2.expect_snapshot_and_fetch(1).await,
1169            all_ok(&data[..1], 1)
1170        );
1171        write1.expect_compare_and_append(&data[1..], 2, 4).await;
1172        assert_eq!(read1.expect_snapshot_and_fetch(3).await, all_ok(&data, 3));
1173    }
1174
1175    #[mz_persist_proc::test(tokio::test)]
1176    #[cfg_attr(miri, ignore)] // too slow
1177    async fn invalid_usage(dyncfgs: ConfigUpdates) {
1178        let data = vec![
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 shard_id0 = "s00000000-0000-0000-0000-000000000000"
1185            .parse::<ShardId>()
1186            .expect("invalid shard id");
1187        let mut client = new_test_client(&dyncfgs).await;
1188
1189        let (mut write0, mut read0) = client
1190            .expect_open::<String, String, u64, i64>(shard_id0)
1191            .await;
1192
1193        write0.expect_compare_and_append(&data, 0, 4).await;
1194
1195        // InvalidUsage from PersistClient methods.
1196        {
1197            fn codecs(
1198                k: &str,
1199                v: &str,
1200                t: &str,
1201                d: &str,
1202            ) -> (String, String, String, String, Option<CodecConcreteType>) {
1203                (k.to_owned(), v.to_owned(), t.to_owned(), d.to_owned(), None)
1204            }
1205
1206            client.shared_states = Arc::new(StateCache::new_no_metrics());
1207            assert_eq!(
1208                client
1209                    .open::<Vec<u8>, String, u64, i64>(
1210                        shard_id0,
1211                        Arc::new(VecU8Schema),
1212                        Arc::new(StringSchema),
1213                        Diagnostics::for_tests(),
1214                        true,
1215                    )
1216                    .await
1217                    .unwrap_err(),
1218                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1219                    requested: codecs("Vec<u8>", "String", "u64", "i64"),
1220                    actual: codecs("String", "String", "u64", "i64"),
1221                }))
1222            );
1223            assert_eq!(
1224                client
1225                    .open::<String, Vec<u8>, u64, i64>(
1226                        shard_id0,
1227                        Arc::new(StringSchema),
1228                        Arc::new(VecU8Schema),
1229                        Diagnostics::for_tests(),
1230                        true,
1231                    )
1232                    .await
1233                    .unwrap_err(),
1234                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1235                    requested: codecs("String", "Vec<u8>", "u64", "i64"),
1236                    actual: codecs("String", "String", "u64", "i64"),
1237                }))
1238            );
1239            assert_eq!(
1240                client
1241                    .open::<String, String, i64, i64>(
1242                        shard_id0,
1243                        Arc::new(StringSchema),
1244                        Arc::new(StringSchema),
1245                        Diagnostics::for_tests(),
1246                        true,
1247                    )
1248                    .await
1249                    .unwrap_err(),
1250                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1251                    requested: codecs("String", "String", "i64", "i64"),
1252                    actual: codecs("String", "String", "u64", "i64"),
1253                }))
1254            );
1255            assert_eq!(
1256                client
1257                    .open::<String, String, u64, u64>(
1258                        shard_id0,
1259                        Arc::new(StringSchema),
1260                        Arc::new(StringSchema),
1261                        Diagnostics::for_tests(),
1262                        true,
1263                    )
1264                    .await
1265                    .unwrap_err(),
1266                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1267                    requested: codecs("String", "String", "u64", "u64"),
1268                    actual: codecs("String", "String", "u64", "i64"),
1269                }))
1270            );
1271
1272            // open_reader and open_writer end up using the same checks, so just
1273            // verify one type each to verify the plumbing instead of the full
1274            // set.
1275            assert_eq!(
1276                client
1277                    .open_leased_reader::<Vec<u8>, String, u64, i64>(
1278                        shard_id0,
1279                        Arc::new(VecU8Schema),
1280                        Arc::new(StringSchema),
1281                        Diagnostics::for_tests(),
1282                        true,
1283                    )
1284                    .await
1285                    .unwrap_err(),
1286                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1287                    requested: codecs("Vec<u8>", "String", "u64", "i64"),
1288                    actual: codecs("String", "String", "u64", "i64"),
1289                }))
1290            );
1291            assert_eq!(
1292                client
1293                    .open_writer::<Vec<u8>, String, u64, i64>(
1294                        shard_id0,
1295                        Arc::new(VecU8Schema),
1296                        Arc::new(StringSchema),
1297                        Diagnostics::for_tests(),
1298                    )
1299                    .await
1300                    .unwrap_err(),
1301                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1302                    requested: codecs("Vec<u8>", "String", "u64", "i64"),
1303                    actual: codecs("String", "String", "u64", "i64"),
1304                }))
1305            );
1306        }
1307
1308        // InvalidUsage from ReadHandle methods.
1309        {
1310            let snap = read0
1311                .snapshot(Antichain::from_elem(3))
1312                .await
1313                .expect("cannot serve requested as_of");
1314
1315            let shard_id1 = "s11111111-1111-1111-1111-111111111111"
1316                .parse::<ShardId>()
1317                .expect("invalid shard id");
1318            let mut fetcher1 = client
1319                .create_batch_fetcher::<String, String, u64, i64>(
1320                    shard_id1,
1321                    Default::default(),
1322                    Default::default(),
1323                    false,
1324                    Diagnostics::for_tests(),
1325                )
1326                .await
1327                .unwrap();
1328            for part in snap {
1329                let (part, _lease) = part.into_exchangeable_part();
1330                let res = fetcher1.fetch_leased_part(part).await;
1331                assert_eq!(
1332                    res.unwrap_err(),
1333                    InvalidUsage::BatchNotFromThisShard {
1334                        batch_shard: shard_id0,
1335                        handle_shard: shard_id1,
1336                    }
1337                );
1338            }
1339        }
1340
1341        // InvalidUsage from WriteHandle methods.
1342        {
1343            let ts3 = &data[2];
1344            assert_eq!(ts3.1, 3);
1345            let ts3 = vec![ts3.clone()];
1346
1347            // WriteHandle::append also covers append_batch,
1348            // compare_and_append_batch, compare_and_append.
1349            assert_eq!(
1350                write0
1351                    .append(&ts3, Antichain::from_elem(4), Antichain::from_elem(5))
1352                    .await
1353                    .unwrap_err(),
1354                InvalidUsage::UpdateNotBeyondLower {
1355                    ts: 3,
1356                    lower: Antichain::from_elem(4),
1357                },
1358            );
1359            assert_eq!(
1360                write0
1361                    .append(&ts3, Antichain::from_elem(2), Antichain::from_elem(3))
1362                    .await
1363                    .unwrap_err(),
1364                InvalidUsage::UpdateBeyondUpper {
1365                    ts: 3,
1366                    expected_upper: Antichain::from_elem(3),
1367                },
1368            );
1369            // NB unlike the previous tests, this one has empty updates.
1370            assert_eq!(
1371                write0
1372                    .append(&data[..0], Antichain::from_elem(3), Antichain::from_elem(2))
1373                    .await
1374                    .unwrap_err(),
1375                InvalidUsage::InvalidBounds {
1376                    lower: Antichain::from_elem(3),
1377                    upper: Antichain::from_elem(2),
1378                },
1379            );
1380
1381            // Tests for the BatchBuilder.
1382            assert_eq!(
1383                write0
1384                    .builder(Antichain::from_elem(3))
1385                    .finish(Antichain::from_elem(2))
1386                    .await
1387                    .unwrap_err(),
1388                InvalidUsage::InvalidBounds {
1389                    lower: Antichain::from_elem(3),
1390                    upper: Antichain::from_elem(2)
1391                },
1392            );
1393            let batch = write0
1394                .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1395                .await
1396                .expect("invalid usage");
1397            assert_eq!(
1398                write0
1399                    .append_batch(batch, Antichain::from_elem(4), Antichain::from_elem(5))
1400                    .await
1401                    .unwrap_err(),
1402                InvalidUsage::InvalidBatchBounds {
1403                    batch_lower: Antichain::from_elem(3),
1404                    batch_upper: Antichain::from_elem(4),
1405                    append_lower: Antichain::from_elem(4),
1406                    append_upper: Antichain::from_elem(5),
1407                },
1408            );
1409            let batch = write0
1410                .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1411                .await
1412                .expect("invalid usage");
1413            assert_eq!(
1414                write0
1415                    .append_batch(batch, Antichain::from_elem(2), Antichain::from_elem(3))
1416                    .await
1417                    .unwrap_err(),
1418                InvalidUsage::InvalidBatchBounds {
1419                    batch_lower: Antichain::from_elem(3),
1420                    batch_upper: Antichain::from_elem(4),
1421                    append_lower: Antichain::from_elem(2),
1422                    append_upper: Antichain::from_elem(3),
1423                },
1424            );
1425            let batch = write0
1426                .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1427                .await
1428                .expect("invalid usage");
1429            // NB unlike the others, this one uses matches! because it's
1430            // non-deterministic (the key)
1431            assert!(matches!(
1432                write0
1433                    .append_batch(batch, Antichain::from_elem(3), Antichain::from_elem(3))
1434                    .await
1435                    .unwrap_err(),
1436                InvalidUsage::InvalidEmptyTimeInterval { .. }
1437            ));
1438        }
1439    }
1440
1441    #[mz_persist_proc::test(tokio::test)]
1442    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1443    async fn multiple_shards(dyncfgs: ConfigUpdates) {
1444        let data1 = [
1445            (("1".to_owned(), "one".to_owned()), 1, 1),
1446            (("2".to_owned(), "two".to_owned()), 2, 1),
1447        ];
1448
1449        let data2 = [(("1".to_owned(), ()), 1, 1), (("2".to_owned(), ()), 2, 1)];
1450
1451        let client = new_test_client(&dyncfgs).await;
1452
1453        let (mut write1, mut read1) = client
1454            .expect_open::<String, String, u64, i64>(ShardId::new())
1455            .await;
1456
1457        // Different types, so that checks would fail in case we were not separating these
1458        // collections internally.
1459        let (mut write2, mut read2) = client
1460            .expect_open::<String, (), u64, i64>(ShardId::new())
1461            .await;
1462
1463        write1
1464            .expect_compare_and_append(&data1[..], u64::minimum(), 3)
1465            .await;
1466
1467        write2
1468            .expect_compare_and_append(&data2[..], u64::minimum(), 3)
1469            .await;
1470
1471        assert_eq!(
1472            read1.expect_snapshot_and_fetch(2).await,
1473            all_ok(&data1[..], 2)
1474        );
1475
1476        assert_eq!(
1477            read2.expect_snapshot_and_fetch(2).await,
1478            all_ok(&data2[..], 2)
1479        );
1480    }
1481
1482    #[mz_persist_proc::test(tokio::test)]
1483    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1484    async fn fetch_upper(dyncfgs: ConfigUpdates) {
1485        let data = [
1486            (("1".to_owned(), "one".to_owned()), 1, 1),
1487            (("2".to_owned(), "two".to_owned()), 2, 1),
1488        ];
1489
1490        let client = new_test_client(&dyncfgs).await;
1491
1492        let shard_id = ShardId::new();
1493
1494        let (mut write1, _read1) = client
1495            .expect_open::<String, String, u64, i64>(shard_id)
1496            .await;
1497
1498        let (mut write2, _read2) = client
1499            .expect_open::<String, String, u64, i64>(shard_id)
1500            .await;
1501
1502        write1
1503            .expect_append(&data[..], write1.upper().clone(), vec![3])
1504            .await;
1505
1506        // The shard-global upper does advance, even if this writer didn't advance its local upper.
1507        assert_eq!(write2.fetch_recent_upper().await, &Antichain::from_elem(3));
1508
1509        // The writer-local upper should advance, even if it was another writer
1510        // that advanced the frontier.
1511        assert_eq!(write2.upper(), &Antichain::from_elem(3));
1512    }
1513
1514    #[mz_persist_proc::test(tokio::test)]
1515    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1516    async fn append_with_invalid_upper(dyncfgs: ConfigUpdates) {
1517        let data = [
1518            (("1".to_owned(), "one".to_owned()), 1, 1),
1519            (("2".to_owned(), "two".to_owned()), 2, 1),
1520        ];
1521
1522        let client = new_test_client(&dyncfgs).await;
1523
1524        let shard_id = ShardId::new();
1525
1526        let (mut write, _read) = client
1527            .expect_open::<String, String, u64, i64>(shard_id)
1528            .await;
1529
1530        write
1531            .expect_append(&data[..], write.upper().clone(), vec![3])
1532            .await;
1533
1534        let data = [
1535            (("5".to_owned(), "fünf".to_owned()), 5, 1),
1536            (("6".to_owned(), "sechs".to_owned()), 6, 1),
1537        ];
1538        let res = write
1539            .append(
1540                data.iter(),
1541                Antichain::from_elem(5),
1542                Antichain::from_elem(7),
1543            )
1544            .await;
1545        assert_eq!(
1546            res,
1547            Ok(Err(UpperMismatch {
1548                expected: Antichain::from_elem(5),
1549                current: Antichain::from_elem(3)
1550            }))
1551        );
1552
1553        // Writing with an outdated upper updates the write handle's upper to the correct upper.
1554        assert_eq!(write.upper(), &Antichain::from_elem(3));
1555    }
1556
1557    // Make sure that the API structs are Sync + Send, so that they can be used in async tasks.
1558    // NOTE: This is a compile-time only test. If it compiles, we're good.
1559    #[allow(unused)]
1560    async fn sync_send(dyncfgs: ConfigUpdates) {
1561        mz_ore::test::init_logging();
1562
1563        fn is_send_sync<T: Send + Sync>(_x: T) -> bool {
1564            true
1565        }
1566
1567        let client = new_test_client(&dyncfgs).await;
1568
1569        let (write, read) = client
1570            .expect_open::<String, String, u64, i64>(ShardId::new())
1571            .await;
1572
1573        assert!(is_send_sync(client));
1574        assert!(is_send_sync(write));
1575        assert!(is_send_sync(read));
1576    }
1577
1578    #[mz_persist_proc::test(tokio::test)]
1579    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1580    async fn compare_and_append(dyncfgs: ConfigUpdates) {
1581        let data = vec![
1582            (("1".to_owned(), "one".to_owned()), 1, 1),
1583            (("2".to_owned(), "two".to_owned()), 2, 1),
1584            (("3".to_owned(), "three".to_owned()), 3, 1),
1585        ];
1586
1587        let id = ShardId::new();
1588        let client = new_test_client(&dyncfgs).await;
1589        let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1590
1591        let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1592
1593        assert_eq!(write1.upper(), &Antichain::from_elem(u64::minimum()));
1594        assert_eq!(write2.upper(), &Antichain::from_elem(u64::minimum()));
1595        assert_eq!(read.since(), &Antichain::from_elem(u64::minimum()));
1596
1597        // Write a [0,3) batch.
1598        write1
1599            .expect_compare_and_append(&data[..2], u64::minimum(), 3)
1600            .await;
1601        assert_eq!(write1.upper(), &Antichain::from_elem(3));
1602
1603        assert_eq!(
1604            read.expect_snapshot_and_fetch(2).await,
1605            all_ok(&data[..2], 2)
1606        );
1607
1608        // Try and write with a wrong expected upper.
1609        let res = write2
1610            .compare_and_append(
1611                &data[..2],
1612                Antichain::from_elem(u64::minimum()),
1613                Antichain::from_elem(3),
1614            )
1615            .await;
1616        assert_eq!(
1617            res,
1618            Ok(Err(UpperMismatch {
1619                expected: Antichain::from_elem(u64::minimum()),
1620                current: Antichain::from_elem(3)
1621            }))
1622        );
1623
1624        // A failed write updates our local cache of the shard upper.
1625        assert_eq!(write2.upper(), &Antichain::from_elem(3));
1626
1627        // Try again with a good expected upper.
1628        write2.expect_compare_and_append(&data[2..], 3, 4).await;
1629
1630        assert_eq!(write2.upper(), &Antichain::from_elem(4));
1631
1632        assert_eq!(read.expect_snapshot_and_fetch(3).await, all_ok(&data, 3));
1633    }
1634
1635    #[mz_persist_proc::test(tokio::test)]
1636    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1637    async fn overlapping_append(dyncfgs: ConfigUpdates) {
1638        mz_ore::test::init_logging_default("info");
1639
1640        let data = vec![
1641            (("1".to_owned(), "one".to_owned()), 1, 1),
1642            (("2".to_owned(), "two".to_owned()), 2, 1),
1643            (("3".to_owned(), "three".to_owned()), 3, 1),
1644            (("4".to_owned(), "vier".to_owned()), 4, 1),
1645            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1646        ];
1647
1648        let id = ShardId::new();
1649        let client = new_test_client(&dyncfgs).await;
1650
1651        let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1652
1653        let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1654
1655        // Grab a listener before we do any writing
1656        let mut listen = read.clone("").await.expect_listen(0).await;
1657
1658        // Write a [0,3) batch.
1659        write1
1660            .expect_append(&data[..2], write1.upper().clone(), vec![3])
1661            .await;
1662        assert_eq!(write1.upper(), &Antichain::from_elem(3));
1663
1664        // Write a [0,5) batch with the second writer.
1665        write2
1666            .expect_append(&data[..4], write2.upper().clone(), vec![5])
1667            .await;
1668        assert_eq!(write2.upper(), &Antichain::from_elem(5));
1669
1670        // Write a [3,6) batch with the first writer.
1671        write1
1672            .expect_append(&data[2..5], write1.upper().clone(), vec![6])
1673            .await;
1674        assert_eq!(write1.upper(), &Antichain::from_elem(6));
1675
1676        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1677
1678        assert_eq!(
1679            listen.read_until(&6).await,
1680            (all_ok(&data[..], 1), Antichain::from_elem(6))
1681        );
1682    }
1683
1684    // Appends need to be contiguous for a shard, meaning the lower of an appended batch must not
1685    // be in advance of the current shard upper.
1686    #[mz_persist_proc::test(tokio::test)]
1687    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1688    async fn contiguous_append(dyncfgs: ConfigUpdates) {
1689        let data = vec![
1690            (("1".to_owned(), "one".to_owned()), 1, 1),
1691            (("2".to_owned(), "two".to_owned()), 2, 1),
1692            (("3".to_owned(), "three".to_owned()), 3, 1),
1693            (("4".to_owned(), "vier".to_owned()), 4, 1),
1694            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1695        ];
1696
1697        let id = ShardId::new();
1698        let client = new_test_client(&dyncfgs).await;
1699
1700        let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1701
1702        // Write a [0,3) batch.
1703        write
1704            .expect_append(&data[..2], write.upper().clone(), vec![3])
1705            .await;
1706        assert_eq!(write.upper(), &Antichain::from_elem(3));
1707
1708        // Appending a non-contiguous batch should fail.
1709        // Write a [5,6) batch with the second writer.
1710        let result = write
1711            .append(
1712                &data[4..5],
1713                Antichain::from_elem(5),
1714                Antichain::from_elem(6),
1715            )
1716            .await;
1717        assert_eq!(
1718            result,
1719            Ok(Err(UpperMismatch {
1720                expected: Antichain::from_elem(5),
1721                current: Antichain::from_elem(3)
1722            }))
1723        );
1724
1725        // Fixing the lower to make the write contiguous should make the append succeed.
1726        write.expect_append(&data[2..5], vec![3], vec![6]).await;
1727        assert_eq!(write.upper(), &Antichain::from_elem(6));
1728
1729        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1730    }
1731
1732    // Per-writer appends can be non-contiguous, as long as appends to the shard from all writers
1733    // combined are contiguous.
1734    #[mz_persist_proc::test(tokio::test)]
1735    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1736    async fn noncontiguous_append_per_writer(dyncfgs: ConfigUpdates) {
1737        let data = vec![
1738            (("1".to_owned(), "one".to_owned()), 1, 1),
1739            (("2".to_owned(), "two".to_owned()), 2, 1),
1740            (("3".to_owned(), "three".to_owned()), 3, 1),
1741            (("4".to_owned(), "vier".to_owned()), 4, 1),
1742            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1743        ];
1744
1745        let id = ShardId::new();
1746        let client = new_test_client(&dyncfgs).await;
1747
1748        let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1749
1750        let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1751
1752        // Write a [0,3) batch with writer 1.
1753        write1
1754            .expect_append(&data[..2], write1.upper().clone(), vec![3])
1755            .await;
1756        assert_eq!(write1.upper(), &Antichain::from_elem(3));
1757
1758        // Write a [3,5) batch with writer 2.
1759        write2.upper = Antichain::from_elem(3);
1760        write2
1761            .expect_append(&data[2..4], write2.upper().clone(), vec![5])
1762            .await;
1763        assert_eq!(write2.upper(), &Antichain::from_elem(5));
1764
1765        // Write a [5,6) batch with writer 1.
1766        write1.upper = Antichain::from_elem(5);
1767        write1
1768            .expect_append(&data[4..5], write1.upper().clone(), vec![6])
1769            .await;
1770        assert_eq!(write1.upper(), &Antichain::from_elem(6));
1771
1772        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1773    }
1774
1775    // Compare_and_appends need to be contiguous for a shard, meaning the lower of an appended
1776    // batch needs to match the current shard upper.
1777    #[mz_persist_proc::test(tokio::test)]
1778    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1779    async fn contiguous_compare_and_append(dyncfgs: ConfigUpdates) {
1780        let data = vec![
1781            (("1".to_owned(), "one".to_owned()), 1, 1),
1782            (("2".to_owned(), "two".to_owned()), 2, 1),
1783            (("3".to_owned(), "three".to_owned()), 3, 1),
1784            (("4".to_owned(), "vier".to_owned()), 4, 1),
1785            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1786        ];
1787
1788        let id = ShardId::new();
1789        let client = new_test_client(&dyncfgs).await;
1790
1791        let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1792
1793        // Write a [0,3) batch.
1794        write.expect_compare_and_append(&data[..2], 0, 3).await;
1795        assert_eq!(write.upper(), &Antichain::from_elem(3));
1796
1797        // Appending a non-contiguous batch should fail.
1798        // Write a [5,6) batch with the second writer.
1799        let result = write
1800            .compare_and_append(
1801                &data[4..5],
1802                Antichain::from_elem(5),
1803                Antichain::from_elem(6),
1804            )
1805            .await;
1806        assert_eq!(
1807            result,
1808            Ok(Err(UpperMismatch {
1809                expected: Antichain::from_elem(5),
1810                current: Antichain::from_elem(3)
1811            }))
1812        );
1813
1814        // Writing with the correct expected upper to make the write contiguous should make the
1815        // append succeed.
1816        write.expect_compare_and_append(&data[2..5], 3, 6).await;
1817        assert_eq!(write.upper(), &Antichain::from_elem(6));
1818
1819        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1820    }
1821
1822    // Per-writer compare_and_appends can be non-contiguous, as long as appends to the shard from
1823    // all writers combined are contiguous.
1824    #[mz_persist_proc::test(tokio::test)]
1825    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1826    async fn noncontiguous_compare_and_append_per_writer(dyncfgs: ConfigUpdates) {
1827        let data = vec![
1828            (("1".to_owned(), "one".to_owned()), 1, 1),
1829            (("2".to_owned(), "two".to_owned()), 2, 1),
1830            (("3".to_owned(), "three".to_owned()), 3, 1),
1831            (("4".to_owned(), "vier".to_owned()), 4, 1),
1832            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1833        ];
1834
1835        let id = ShardId::new();
1836        let client = new_test_client(&dyncfgs).await;
1837
1838        let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1839
1840        let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1841
1842        // Write a [0,3) batch with writer 1.
1843        write1.expect_compare_and_append(&data[..2], 0, 3).await;
1844        assert_eq!(write1.upper(), &Antichain::from_elem(3));
1845
1846        // Write a [3,5) batch with writer 2.
1847        write2.expect_compare_and_append(&data[2..4], 3, 5).await;
1848        assert_eq!(write2.upper(), &Antichain::from_elem(5));
1849
1850        // Write a [5,6) batch with writer 1.
1851        write1.expect_compare_and_append(&data[4..5], 5, 6).await;
1852        assert_eq!(write1.upper(), &Antichain::from_elem(6));
1853
1854        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1855    }
1856
1857    #[mz_ore::test]
1858    fn fmt_ids() {
1859        assert_eq!(
1860            format!("{}", LeasedReaderId([0u8; 16])),
1861            "r00000000-0000-0000-0000-000000000000"
1862        );
1863        assert_eq!(
1864            format!("{:?}", LeasedReaderId([0u8; 16])),
1865            "LeasedReaderId(00000000-0000-0000-0000-000000000000)"
1866        );
1867    }
1868
1869    #[mz_persist_proc::test(tokio::test(flavor = "multi_thread"))]
1870    #[cfg_attr(miri, ignore)] // error: unsupported operation: integer-to-pointer casts and `ptr::from_exposed_addr` are not supported with `-Zmiri-strict-provenance`
1871    async fn concurrency(dyncfgs: ConfigUpdates) {
1872        let data = DataGenerator::small();
1873
1874        const NUM_WRITERS: usize = 2;
1875        let id = ShardId::new();
1876        let client = new_test_client(&dyncfgs).await;
1877        let mut handles = Vec::<mz_ore::task::JoinHandle<()>>::new();
1878        for idx in 0..NUM_WRITERS {
1879            let (data, client) = (data.clone(), client.clone());
1880
1881            let (batch_tx, mut batch_rx) = tokio::sync::mpsc::channel(1);
1882
1883            let client1 = client.clone();
1884            let handle = mz_ore::task::spawn(|| format!("writer-{}", idx), async move {
1885                let (write, _) = client1.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1886                let mut current_upper = 0;
1887                for batch in data.batches() {
1888                    let new_upper = match batch.get(batch.len() - 1) {
1889                        Some((_, max_ts, _)) => u64::decode(max_ts) + 1,
1890                        None => continue,
1891                    };
1892                    // Because we (intentionally) call open inside the task,
1893                    // some other writer may have raced ahead and already
1894                    // appended some data before this one was registered. As a
1895                    // result, this writer may not be starting with an upper of
1896                    // the initial empty antichain. This is nice because it
1897                    // mimics how a real HA source would work, but it means we
1898                    // have to skip any batches that have already been committed
1899                    // (otherwise our new_upper would be before our upper).
1900                    //
1901                    // Note however, that unlike a real source, our
1902                    // DataGenerator-derived batches are guaranteed to be
1903                    // chunked along the same boundaries. This means we don't
1904                    // have to consider partial batches when generating the
1905                    // updates below.
1906                    if PartialOrder::less_equal(&Antichain::from_elem(new_upper), write.upper()) {
1907                        continue;
1908                    }
1909
1910                    let current_upper_chain = Antichain::from_elem(current_upper);
1911                    current_upper = new_upper;
1912                    let new_upper_chain = Antichain::from_elem(new_upper);
1913                    let mut builder = write.builder(current_upper_chain);
1914
1915                    for ((k, v), t, d) in batch.iter() {
1916                        builder
1917                            .add(&k.to_vec(), &v.to_vec(), &u64::decode(t), &i64::decode(d))
1918                            .await
1919                            .expect("invalid usage");
1920                    }
1921
1922                    let batch = builder
1923                        .finish(new_upper_chain)
1924                        .await
1925                        .expect("invalid usage");
1926
1927                    match batch_tx.send(batch).await {
1928                        Ok(_) => (),
1929                        Err(e) => panic!("send error: {}", e),
1930                    }
1931                }
1932            });
1933            handles.push(handle);
1934
1935            let handle = mz_ore::task::spawn(|| format!("appender-{}", idx), async move {
1936                let (mut write, _) = client.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1937
1938                while let Some(batch) = batch_rx.recv().await {
1939                    let lower = batch.lower().clone();
1940                    let upper = batch.upper().clone();
1941                    write
1942                        .append_batch(batch, lower, upper)
1943                        .await
1944                        .expect("invalid usage")
1945                        .expect("unexpected upper");
1946                }
1947            });
1948            handles.push(handle);
1949        }
1950
1951        for handle in handles {
1952            let () = handle.await;
1953        }
1954
1955        let expected = data.records().collect::<Vec<_>>();
1956        let max_ts = expected.last().map(|(_, t, _)| *t).unwrap_or_default();
1957        let (_, mut read) = client.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1958        assert_eq!(
1959            read.expect_snapshot_and_fetch(max_ts).await,
1960            all_ok(expected.iter(), max_ts)
1961        );
1962    }
1963
1964    // Regression test for database-issues#3523. Snapshot with as_of >= upper would
1965    // immediately return the data currently available instead of waiting for
1966    // upper to advance past as_of.
1967    #[mz_persist_proc::test(tokio::test)]
1968    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1969    async fn regression_blocking_reads(dyncfgs: ConfigUpdates) {
1970        let waker = noop_waker();
1971        let mut cx = Context::from_waker(&waker);
1972
1973        let data = [
1974            (("1".to_owned(), "one".to_owned()), 1, 1),
1975            (("2".to_owned(), "two".to_owned()), 2, 1),
1976            (("3".to_owned(), "three".to_owned()), 3, 1),
1977        ];
1978
1979        let id = ShardId::new();
1980        let client = new_test_client(&dyncfgs).await;
1981        let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1982
1983        // Grab a listener as_of (aka gt) 1, which is not yet closed out.
1984        let mut listen = read.clone("").await.expect_listen(1).await;
1985        let mut listen_next = Box::pin(listen.fetch_next());
1986        // Intentionally don't await the listen_next, but instead manually poke
1987        // it for a while and assert that it doesn't resolve yet. See below for
1988        // discussion of some alternative ways of writing this unit test.
1989        for _ in 0..100 {
1990            assert!(
1991                Pin::new(&mut listen_next).poll(&mut cx).is_pending(),
1992                "listen::next unexpectedly ready"
1993            );
1994        }
1995
1996        // Write a [0,3) batch.
1997        write
1998            .expect_compare_and_append(&data[..2], u64::minimum(), 3)
1999            .await;
2000
2001        // The initial listen_next call should now be able to return data at 2.
2002        // It doesn't get 1 because the as_of was 1 and listen is strictly gt.
2003        assert_eq!(
2004            listen_next.await,
2005            vec![
2006                ListenEvent::Updates(vec![(("2".to_owned(), "two".to_owned()), 2, 1)]),
2007                ListenEvent::Progress(Antichain::from_elem(3)),
2008            ]
2009        );
2010
2011        // Grab a snapshot as_of 3, which is not yet closed out. Intentionally
2012        // don't await the snap, but instead manually poke it for a while and
2013        // assert that it doesn't resolve yet.
2014        //
2015        // An alternative to this would be to run it in a task and poll the task
2016        // with some timeout, but this would introduce a fixed test execution
2017        // latency of the timeout in the happy case. Plus, it would be
2018        // non-deterministic.
2019        //
2020        // Another alternative (that's potentially quite interesting!) would be
2021        // to separate creating a snapshot immediately (which would fail if
2022        // as_of was >= upper) from a bit of logic that retries until that case
2023        // is ready.
2024        let mut snap = Box::pin(read.expect_snapshot_and_fetch(3));
2025        for _ in 0..100 {
2026            assert!(
2027                Pin::new(&mut snap).poll(&mut cx).is_pending(),
2028                "snapshot unexpectedly ready"
2029            );
2030        }
2031
2032        // Now add the data at 3 and also unblock the snapshot.
2033        write.expect_compare_and_append(&data[2..], 3, 4).await;
2034
2035        // Read the snapshot and check that it got all the appropriate data.
2036        assert_eq!(snap.await, all_ok(&data[..], 3));
2037    }
2038
2039    #[mz_persist_proc::test(tokio::test)]
2040    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
2041    async fn heartbeat_task_shutdown(dyncfgs: ConfigUpdates) {
2042        // Verify that the ReadHandle and WriteHandle background heartbeat tasks
2043        // shut down cleanly after the handle is expired.
2044        let mut cache = new_test_client_cache(&dyncfgs);
2045        cache
2046            .cfg
2047            .set_config(&READER_LEASE_DURATION, Duration::from_millis(1));
2048        cache.cfg.writer_lease_duration = Duration::from_millis(1);
2049        let (_write, mut read) = cache
2050            .open(PersistLocation::new_in_mem())
2051            .await
2052            .expect("client construction failed")
2053            .expect_open::<(), (), u64, i64>(ShardId::new())
2054            .await;
2055        let read_unexpired_state = read
2056            .unexpired_state
2057            .take()
2058            .expect("handle should have unexpired state");
2059        read.expire().await;
2060        read_unexpired_state.heartbeat_task.await
2061    }
2062
2063    /// Verify that shard finalization works with empty shards, shards that have
2064    /// an empty write up to the empty upper Antichain.
2065    #[mz_persist_proc::test(tokio::test)]
2066    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
2067    async fn finalize_empty_shard(dyncfgs: ConfigUpdates) {
2068        let persist_client = new_test_client(&dyncfgs).await;
2069
2070        let shard_id = ShardId::new();
2071        pub const CRITICAL_SINCE: CriticalReaderId =
2072            CriticalReaderId([0, 0, 0, 0, 17, 17, 34, 34, 51, 51, 68, 68, 68, 68, 68, 68]);
2073
2074        let (mut write, mut read) = persist_client
2075            .expect_open::<(), (), u64, i64>(shard_id)
2076            .await;
2077
2078        // Advance since and upper to empty, which is a pre-requisite for
2079        // finalization/tombstoning.
2080        let () = read.downgrade_since(&Antichain::new()).await;
2081        let () = write.advance_upper(&Antichain::new()).await;
2082
2083        let mut since_handle: SinceHandle<(), (), u64, i64> = persist_client
2084            .open_critical_since(
2085                shard_id,
2086                CRITICAL_SINCE,
2087                Opaque::encode(&0u64),
2088                Diagnostics::for_tests(),
2089            )
2090            .await
2091            .expect("invalid persist usage");
2092
2093        let epoch = since_handle.opaque().clone();
2094        let new_since = Antichain::new();
2095        let downgrade = since_handle
2096            .compare_and_downgrade_since(&epoch, (&epoch, &new_since))
2097            .await;
2098
2099        assert!(
2100            downgrade.is_ok(),
2101            "downgrade of critical handle must succeed"
2102        );
2103
2104        let finalize = persist_client
2105            .finalize_shard::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2106            .await;
2107
2108        assert_ok!(finalize, "finalization must succeed");
2109
2110        let is_finalized = persist_client
2111            .is_finalized::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2112            .await
2113            .expect("invalid persist usage");
2114        assert!(is_finalized, "shard must still be finalized");
2115    }
2116
2117    /// Verify that shard finalization works with shards that had some data
2118    /// written to them, plus then an empty batch to bring their upper to the
2119    /// empty Antichain.
2120    #[mz_persist_proc::test(tokio::test)]
2121    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
2122    async fn finalize_shard(dyncfgs: ConfigUpdates) {
2123        const DATA: &[(((), ()), u64, i64)] = &[(((), ()), 0, 1)];
2124        let persist_client = new_test_client(&dyncfgs).await;
2125
2126        let shard_id = ShardId::new();
2127        pub const CRITICAL_SINCE: CriticalReaderId =
2128            CriticalReaderId([0, 0, 0, 0, 17, 17, 34, 34, 51, 51, 68, 68, 68, 68, 68, 68]);
2129
2130        let (mut write, mut read) = persist_client
2131            .expect_open::<(), (), u64, i64>(shard_id)
2132            .await;
2133
2134        // Write some data.
2135        let () = write
2136            .compare_and_append(DATA, Antichain::from_elem(0), Antichain::from_elem(1))
2137            .await
2138            .expect("usage should be valid")
2139            .expect("upper should match");
2140
2141        // Advance since and upper to empty, which is a pre-requisite for
2142        // finalization/tombstoning.
2143        let () = read.downgrade_since(&Antichain::new()).await;
2144        let () = write.advance_upper(&Antichain::new()).await;
2145
2146        let mut since_handle: SinceHandle<(), (), u64, i64> = persist_client
2147            .open_critical_since(
2148                shard_id,
2149                CRITICAL_SINCE,
2150                Opaque::encode(&0u64),
2151                Diagnostics::for_tests(),
2152            )
2153            .await
2154            .expect("invalid persist usage");
2155
2156        let epoch = since_handle.opaque().clone();
2157        let new_since = Antichain::new();
2158        let downgrade = since_handle
2159            .compare_and_downgrade_since(&epoch, (&epoch, &new_since))
2160            .await;
2161
2162        assert!(
2163            downgrade.is_ok(),
2164            "downgrade of critical handle must succeed"
2165        );
2166
2167        let finalize = persist_client
2168            .finalize_shard::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2169            .await;
2170
2171        assert_ok!(finalize, "finalization must succeed");
2172
2173        let is_finalized = persist_client
2174            .is_finalized::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2175            .await
2176            .expect("invalid persist usage");
2177        assert!(is_finalized, "shard must still be finalized");
2178    }
2179
2180    proptest! {
2181        #![proptest_config(ProptestConfig::with_cases(4096))]
2182
2183        #[mz_ore::test]
2184        #[cfg_attr(miri, ignore)] // too slow
2185        fn shard_id_protobuf_roundtrip(expect in any::<ShardId>() ) {
2186            let actual = protobuf_roundtrip::<_, String>(&expect);
2187            assert_ok!(actual);
2188            assert_eq!(actual.unwrap(), expect);
2189        }
2190    }
2191}