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        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    #[instrument(level = "debug", fields(shard = %shard_id))]
554    pub async fn batch_builder<K, V, T, D>(
555        &self,
556        shard_id: ShardId,
557        write_schemas: Schemas<K, V>,
558        lower: Antichain<T>,
559        max_runs: Option<usize>,
560    ) -> BatchBuilder<K, V, T, D>
561    where
562        K: Debug + Codec,
563        V: Debug + Codec,
564        T: Timestamp + Lattice + Codec64 + TotalOrder + Sync,
565        D: Monoid + Ord + Codec64 + Send + Sync,
566    {
567        let mut compact_cfg = CompactConfig::new(&self.cfg, shard_id);
568        compact_cfg.batch.max_runs = max_runs;
569        WriteHandle::builder_inner(
570            &self.cfg,
571            compact_cfg,
572            Arc::clone(&self.metrics),
573            self.metrics.shards.shard(&shard_id, "peek_stash"),
574            &self.metrics.user,
575            Arc::clone(&self.isolated_runtime),
576            Arc::clone(&self.blob),
577            shard_id,
578            write_schemas,
579            lower,
580        )
581    }
582
583    /// Turns the given [`ProtoBatch`] back into a [`Batch`] which can be used
584    /// to append it to the given shard or to read it via
585    /// [PersistClient::read_batches_consolidated]
586    ///
587    /// CAUTION: This API allows turning a [ProtoBatch] into a [Batch] multiple
588    /// times, but if a batch is deleted the backing data goes away, so at that
589    /// point all in-memory copies of a batch become invalid and cannot be read
590    /// anymore.
591    pub fn batch_from_transmittable_batch<K, V, T, D>(
592        &self,
593        shard_id: &ShardId,
594        batch: ProtoBatch,
595    ) -> Batch<K, V, T, D>
596    where
597        K: Debug + Codec,
598        V: Debug + Codec,
599        T: Timestamp + Lattice + Codec64 + Sync,
600        D: Monoid + Ord + Codec64 + Send + Sync,
601    {
602        let batch_shard_id: ShardId = batch
603            .shard_id
604            .into_rust()
605            .expect("valid transmittable batch");
606        assert_eq!(&batch_shard_id, shard_id);
607
608        let shard_metrics = self.metrics.shards.shard(shard_id, "peek_stash");
609
610        let ret = Batch {
611            batch_delete_enabled: BATCH_DELETE_ENABLED.get(&self.cfg),
612            metrics: Arc::clone(&self.metrics),
613            shard_metrics,
614            version: Version::parse(&batch.version).expect("valid transmittable batch"),
615            schemas: (batch.key_schema, batch.val_schema),
616            batch: batch
617                .batch
618                .into_rust_if_some("ProtoBatch::batch")
619                .expect("valid transmittable batch"),
620            blob: Arc::clone(&self.blob),
621            _phantom: std::marker::PhantomData,
622        };
623
624        assert_eq!(&ret.shard_id(), shard_id);
625        ret
626    }
627
628    /// Returns a [Cursor] for reading the given batches. Yielded updates are
629    /// consolidated if the given batches contain sorted runs, which is true
630    /// when they have been written using a [BatchBuilder].
631    ///
632    /// To keep memory usage down when reading a snapshot that consolidates
633    /// well, this consolidates as it goes. However, note that only the
634    /// serialized data is consolidated: the deserialized data will only be
635    /// consolidated if your K/V codecs are one-to-one.
636    ///
637    /// CAUTION: The caller needs to make sure that the given batches are
638    /// readable and they have to remain readable for the lifetime of the
639    /// returned [Cursor]. The caller is also responsible for the lifecycle of
640    /// the batches: once the cursor and the batches are no longer needed you
641    /// must call [Cursor::into_lease] to get back the batches and delete them.
642    #[allow(clippy::unused_async)]
643    pub async fn read_batches_consolidated<K, V, T, D>(
644        &mut self,
645        shard_id: ShardId,
646        as_of: Antichain<T>,
647        read_schemas: Schemas<K, V>,
648        batches: Vec<Batch<K, V, T, D>>,
649        should_fetch_part: impl for<'a> Fn(Option<&'a LazyPartStats>) -> bool,
650        memory_budget_bytes: usize,
651    ) -> Result<Cursor<K, V, T, D, Vec<Batch<K, V, T, D>>>, Since<T>>
652    where
653        K: Debug + Codec + Ord,
654        V: Debug + Codec + Ord,
655        T: Timestamp + Lattice + Codec64 + TotalOrder + Sync,
656        D: Monoid + Ord + Codec64 + Send + Sync,
657    {
658        let shard_metrics = self.metrics.shards.shard(&shard_id, "peek_stash");
659
660        let hollow_batches = batches.iter().map(|b| b.batch.clone()).collect_vec();
661
662        ReadHandle::read_batches_consolidated(
663            &self.cfg,
664            Arc::clone(&self.metrics),
665            shard_metrics,
666            self.metrics.read.snapshot.clone(),
667            Arc::clone(&self.blob),
668            shard_id,
669            as_of,
670            read_schemas,
671            &hollow_batches,
672            batches,
673            should_fetch_part,
674            memory_budget_bytes,
675        )
676    }
677
678    /// Returns the requested schema, if known at the current state.
679    pub async fn get_schema<K, V, T, D>(
680        &self,
681        shard_id: ShardId,
682        schema_id: SchemaId,
683        diagnostics: Diagnostics,
684    ) -> Result<Option<(K::Schema, V::Schema)>, InvalidUsage<T>>
685    where
686        K: Debug + Codec,
687        V: Debug + Codec,
688        T: Timestamp + Lattice + Codec64 + Sync,
689        D: Monoid + Codec64 + Send + Sync,
690    {
691        let machine = self
692            .make_machine::<K, V, T, D>(shard_id, diagnostics)
693            .await?;
694        Ok(machine.get_schema(schema_id))
695    }
696
697    /// Returns the latest schema registered at the current state.
698    pub async fn latest_schema<K, V, T, D>(
699        &self,
700        shard_id: ShardId,
701        diagnostics: Diagnostics,
702    ) -> Result<Option<(SchemaId, K::Schema, V::Schema)>, InvalidUsage<T>>
703    where
704        K: Debug + Codec,
705        V: Debug + Codec,
706        T: Timestamp + Lattice + Codec64 + Sync,
707        D: Monoid + Codec64 + Send + Sync,
708    {
709        let machine = self
710            .make_machine::<K, V, T, D>(shard_id, diagnostics)
711            .await?;
712        Ok(machine.latest_schema())
713    }
714
715    /// Fetches and returns a recent shard-global `upper`, without requiring a
716    /// [`WriteHandle`].
717    ///
718    /// Importantly, this operation is linearized with write operations, giving
719    /// the same guarantee as [`WriteHandle::fetch_recent_upper`]. It requires
720    /// fetching the latest state from consensus and is therefore a potentially
721    /// expensive operation.
722    ///
723    /// If `shard_id` has never been used before, initializes the shard and
724    /// returns an upper of `Antichain::from_elem(T::minimum())`.
725    pub async fn recent_upper<K, V, T, D>(
726        &self,
727        shard_id: ShardId,
728        diagnostics: Diagnostics,
729    ) -> Result<Antichain<T>, InvalidUsage<T>>
730    where
731        K: Debug + Codec,
732        V: Debug + Codec,
733        T: Timestamp + Lattice + Codec64 + Sync,
734        D: Monoid + Codec64 + Send + Sync,
735    {
736        let machine = self
737            .make_machine::<K, V, T, D>(shard_id, diagnostics)
738            .await?;
739        Ok(machine.applier.fetch_upper(|upper| upper.clone()).await)
740    }
741
742    /// Registers a schema for the given shard.
743    ///
744    /// Returns the new schema ID if the registration succeeds, and `None`
745    /// otherwise. Schema registration succeeds in two cases:
746    ///  a) No schema was currently registered for the shard.
747    ///  b) The given schema is already registered for the shard.
748    ///
749    /// To evolve an existing schema instead, use
750    /// [PersistClient::compare_and_evolve_schema].
751    //
752    // TODO: unify with `compare_and_evolve_schema`
753    pub async fn register_schema<K, V, T, D>(
754        &self,
755        shard_id: ShardId,
756        key_schema: &K::Schema,
757        val_schema: &V::Schema,
758        diagnostics: Diagnostics,
759    ) -> Result<Option<SchemaId>, InvalidUsage<T>>
760    where
761        K: Debug + Codec,
762        V: Debug + Codec,
763        T: Timestamp + Lattice + Codec64 + Sync,
764        D: Monoid + Codec64 + Send + Sync,
765    {
766        let machine = self
767            .make_machine::<K, V, T, D>(shard_id, diagnostics)
768            .await?;
769        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
770
771        let (schema_id, maintenance) = machine.register_schema(key_schema, val_schema).await;
772        maintenance.start_performing(&machine, &gc);
773
774        Ok(schema_id)
775    }
776
777    /// Registers a new latest schema for the given shard.
778    ///
779    /// This new schema must be [backward_compatible] with all previous schemas
780    /// for this shard. If it's not, [CaESchema::Incompatible] is returned.
781    ///
782    /// [backward_compatible]: mz_persist_types::schema::backward_compatible
783    ///
784    /// To prevent races, the caller must declare what it believes to be the
785    /// latest schema id. If this doesn't match reality,
786    /// [CaESchema::ExpectedMismatch] is returned.
787    pub async fn compare_and_evolve_schema<K, V, T, D>(
788        &self,
789        shard_id: ShardId,
790        expected: SchemaId,
791        key_schema: &K::Schema,
792        val_schema: &V::Schema,
793        diagnostics: Diagnostics,
794    ) -> Result<CaESchema<K, V>, InvalidUsage<T>>
795    where
796        K: Debug + Codec,
797        V: Debug + Codec,
798        T: Timestamp + Lattice + Codec64 + Sync,
799        D: Monoid + Codec64 + Send + Sync,
800    {
801        let machine = self
802            .make_machine::<K, V, T, D>(shard_id, diagnostics)
803            .await?;
804        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
805        let (res, maintenance) = machine
806            .compare_and_evolve_schema(expected, key_schema, val_schema)
807            .await;
808        maintenance.start_performing(&machine, &gc);
809        Ok(res)
810    }
811
812    /// Check if the given shard is in a finalized state; ie. it can no longer be
813    /// read, any data that was written to it is no longer accessible, and we've
814    /// discarded references to that data from state.
815    pub async fn is_finalized<K, V, T, D>(
816        &self,
817        shard_id: ShardId,
818        diagnostics: Diagnostics,
819    ) -> Result<bool, InvalidUsage<T>>
820    where
821        K: Debug + Codec,
822        V: Debug + Codec,
823        T: Timestamp + Lattice + Codec64 + Sync,
824        D: Monoid + Codec64 + Send + Sync,
825    {
826        let machine = self
827            .make_machine::<K, V, T, D>(shard_id, diagnostics)
828            .await?;
829        Ok(machine.is_finalized())
830    }
831
832    /// If a shard is guaranteed to never be used again, finalize it to delete
833    /// the associated data and release any associated resources. (Except for a
834    /// little state in consensus we use to represent the tombstone.)
835    ///
836    /// The caller should ensure that both the `since` and `upper` of the shard
837    /// have been advanced to `[]`: ie. the shard is no longer writable or readable.
838    /// Otherwise an error is returned.
839    ///
840    /// Once `finalize_shard` has been called, the result of future operations on
841    /// the shard are not defined. They may return errors or succeed as a noop.
842    #[instrument(level = "debug", fields(shard = %shard_id))]
843    pub async fn finalize_shard<K, V, T, D>(
844        &self,
845        shard_id: ShardId,
846        diagnostics: Diagnostics,
847    ) -> Result<(), InvalidUsage<T>>
848    where
849        K: Debug + Codec,
850        V: Debug + Codec,
851        T: Timestamp + Lattice + Codec64 + Sync,
852        D: Monoid + Codec64 + Send + Sync,
853    {
854        let machine = self
855            .make_machine::<K, V, T, D>(shard_id, diagnostics)
856            .await?;
857
858        let maintenance = machine.become_tombstone().await?;
859        let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
860
861        let () = maintenance.perform(&machine, &gc).await;
862
863        Ok(())
864    }
865
866    /// Upgrade the state to the latest version. This should only be called once we will no longer
867    /// need to interoperate with older versions, like after a successful upgrade.
868    pub async fn upgrade_version<K, V, T, D>(
869        &self,
870        shard_id: ShardId,
871        diagnostics: Diagnostics,
872    ) -> Result<(), InvalidUsage<T>>
873    where
874        K: Debug + Codec,
875        V: Debug + Codec,
876        T: Timestamp + Lattice + Codec64 + Sync,
877        D: Monoid + Codec64 + Send + Sync,
878    {
879        let machine = self
880            .make_machine::<K, V, T, D>(shard_id, diagnostics)
881            .await?;
882
883        match machine.upgrade_version().await {
884            Ok(maintenance) => {
885                let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
886                let () = maintenance.perform(&machine, &gc).await;
887                Ok(())
888            }
889            Err(version) => Err(InvalidUsage::IncompatibleVersion { version }),
890        }
891    }
892
893    /// Returns the internal state of the shard for debugging and QA.
894    ///
895    /// We'll be thoughtful about making unnecessary changes, but the **output
896    /// of this method needs to be gated from users**, so that it's not subject
897    /// to our backward compatibility guarantees.
898    pub async fn inspect_shard<T: Timestamp + Lattice + Codec64>(
899        &self,
900        shard_id: &ShardId,
901    ) -> Result<impl serde::Serialize, anyhow::Error> {
902        let state_versions = StateVersions::new(
903            self.cfg.clone(),
904            Arc::clone(&self.consensus),
905            Arc::clone(&self.blob),
906            Arc::clone(&self.metrics),
907        );
908        // TODO: Don't fetch all live diffs. Feels like we should pull out a new
909        // method in StateVersions for fetching the latest version of State of a
910        // shard that might or might not exist.
911        let versions = state_versions.fetch_all_live_diffs(shard_id).await;
912        if versions.is_empty() {
913            return Err(anyhow::anyhow!("{} does not exist", shard_id));
914        }
915        let state = state_versions
916            .fetch_current_state::<T>(shard_id, versions)
917            .await;
918        let state = state.check_ts_codec(shard_id)?;
919        Ok(state)
920    }
921
922    /// Test helper for a [Self::open] call that is expected to succeed.
923    #[cfg(test)]
924    #[track_caller]
925    pub async fn expect_open<K, V, T, D>(
926        &self,
927        shard_id: ShardId,
928    ) -> (WriteHandle<K, V, T, D>, ReadHandle<K, V, T, D>)
929    where
930        K: Debug + Codec,
931        V: Debug + Codec,
932        T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
933        D: Monoid + Ord + Codec64 + Send + Sync,
934        K::Schema: Default,
935        V::Schema: Default,
936    {
937        self.open(
938            shard_id,
939            Arc::new(K::Schema::default()),
940            Arc::new(V::Schema::default()),
941            Diagnostics::for_tests(),
942            true,
943        )
944        .await
945        .expect("codec mismatch")
946    }
947
948    /// Return the metrics being used by this client.
949    ///
950    /// Only exposed for tests, persistcli, and benchmarks.
951    pub fn metrics(&self) -> &Arc<Metrics> {
952        &self.metrics
953    }
954}
955
956#[cfg(test)]
957mod tests {
958    use std::future::Future;
959    use std::pin::Pin;
960    use std::task::Context;
961    use std::time::Duration;
962
963    use differential_dataflow::consolidation::consolidate_updates;
964    use differential_dataflow::lattice::Lattice;
965    use futures_task::noop_waker;
966    use mz_dyncfg::ConfigUpdates;
967    use mz_ore::assert_ok;
968    use mz_persist::indexed::encoding::BlobTraceBatchPart;
969    use mz_persist::workload::DataGenerator;
970    use mz_persist_types::codec_impls::{StringSchema, VecU8Schema};
971    use mz_proto::protobuf_roundtrip;
972    use proptest::prelude::*;
973    use timely::order::PartialOrder;
974    use timely::progress::Antichain;
975
976    use crate::batch::BLOB_TARGET_SIZE;
977    use crate::cache::PersistClientCache;
978    use crate::cfg::BATCH_BUILDER_MAX_OUTSTANDING_PARTS;
979    use crate::critical::Opaque;
980    use crate::error::{CodecConcreteType, CodecMismatch, UpperMismatch};
981    use crate::internal::paths::BlobKey;
982    use crate::read::ListenEvent;
983
984    use super::*;
985
986    pub fn new_test_client_cache(dyncfgs: &ConfigUpdates) -> PersistClientCache {
987        // Configure an aggressively small blob_target_size so we get some
988        // amount of coverage of that in tests. Similarly, for max_outstanding.
989        let mut cache = PersistClientCache::new_no_metrics();
990        cache.cfg.set_config(&BLOB_TARGET_SIZE, 10);
991        cache
992            .cfg
993            .set_config(&BATCH_BUILDER_MAX_OUTSTANDING_PARTS, 1);
994        dyncfgs.apply(cache.cfg());
995
996        // Enable compaction in tests to ensure we get coverage.
997        cache.cfg.compaction_enabled = true;
998        cache
999    }
1000
1001    pub async fn new_test_client(dyncfgs: &ConfigUpdates) -> PersistClient {
1002        let cache = new_test_client_cache(dyncfgs);
1003        cache
1004            .open(PersistLocation::new_in_mem())
1005            .await
1006            .expect("client construction failed")
1007    }
1008
1009    pub fn all_ok<'a, K, V, T, D, I>(iter: I, as_of: T) -> Vec<((K, V), T, D)>
1010    where
1011        K: Ord + Clone + 'a,
1012        V: Ord + Clone + 'a,
1013        T: Timestamp + Lattice + Clone + 'a,
1014        D: Monoid + Clone + 'a,
1015        I: IntoIterator<Item = &'a ((K, V), T, D)>,
1016    {
1017        let as_of = Antichain::from_elem(as_of);
1018        let mut ret = iter
1019            .into_iter()
1020            .map(|((k, v), t, d)| {
1021                let mut t = t.clone();
1022                t.advance_by(as_of.borrow());
1023                ((k.clone(), v.clone()), t, d.clone())
1024            })
1025            .collect();
1026        consolidate_updates(&mut ret);
1027        ret
1028    }
1029
1030    pub async fn expect_fetch_part<K, V, T, D>(
1031        blob: &dyn Blob,
1032        key: &BlobKey,
1033        metrics: &Metrics,
1034        read_schemas: &Schemas<K, V>,
1035    ) -> (BlobTraceBatchPart<T>, Vec<((K, V), T, D)>)
1036    where
1037        K: Codec + Clone,
1038        V: Codec + Clone,
1039        T: Timestamp + Codec64,
1040        D: Codec64,
1041    {
1042        let value = blob
1043            .get(key)
1044            .await
1045            .expect("failed to fetch part")
1046            .expect("missing part");
1047        let mut part =
1048            BlobTraceBatchPart::decode(&value, &metrics.columnar).expect("failed to decode part");
1049        let structured = part
1050            .updates
1051            .into_part::<K, V>(&*read_schemas.key, &*read_schemas.val);
1052        let updates = structured
1053            .decode_iter::<K, V, T, D>(&*read_schemas.key, &*read_schemas.val)
1054            .expect("structured data")
1055            .collect();
1056        (part, updates)
1057    }
1058
1059    #[mz_persist_proc::test(tokio::test)]
1060    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1061    async fn sanity_check(dyncfgs: ConfigUpdates) {
1062        let data = [
1063            (("1".to_owned(), "one".to_owned()), 1, 1),
1064            (("2".to_owned(), "two".to_owned()), 2, 1),
1065            (("3".to_owned(), "three".to_owned()), 3, 1),
1066        ];
1067
1068        let (mut write, mut read) = new_test_client(&dyncfgs)
1069            .await
1070            .expect_open::<String, String, u64, i64>(ShardId::new())
1071            .await;
1072        assert_eq!(write.upper(), &Antichain::from_elem(u64::minimum()));
1073        assert_eq!(read.since(), &Antichain::from_elem(u64::minimum()));
1074
1075        // Write a [0,3) batch.
1076        write
1077            .expect_append(&data[..2], write.upper().clone(), vec![3])
1078            .await;
1079        assert_eq!(write.upper(), &Antichain::from_elem(3));
1080
1081        // Grab a snapshot and listener as_of 1. Snapshot should only have part of what we wrote.
1082        assert_eq!(
1083            read.expect_snapshot_and_fetch(1).await,
1084            all_ok(&data[..1], 1)
1085        );
1086
1087        let mut listen = read.clone("").await.expect_listen(1).await;
1088
1089        // Write a [3,4) batch.
1090        write
1091            .expect_append(&data[2..], write.upper().clone(), vec![4])
1092            .await;
1093        assert_eq!(write.upper(), &Antichain::from_elem(4));
1094
1095        // Listen should have part of the initial write plus the new one.
1096        assert_eq!(
1097            listen.read_until(&4).await,
1098            (all_ok(&data[1..], 1), Antichain::from_elem(4))
1099        );
1100
1101        // Downgrading the since is tracked locally (but otherwise is a no-op).
1102        read.downgrade_since(&Antichain::from_elem(2)).await;
1103        assert_eq!(read.since(), &Antichain::from_elem(2));
1104    }
1105
1106    // Sanity check that the open_reader and open_writer calls work.
1107    #[mz_persist_proc::test(tokio::test)]
1108    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1109    async fn open_reader_writer(dyncfgs: ConfigUpdates) {
1110        let data = vec![
1111            (("1".to_owned(), "one".to_owned()), 1, 1),
1112            (("2".to_owned(), "two".to_owned()), 2, 1),
1113            (("3".to_owned(), "three".to_owned()), 3, 1),
1114        ];
1115
1116        let shard_id = ShardId::new();
1117        let client = new_test_client(&dyncfgs).await;
1118        let mut write1 = client
1119            .open_writer::<String, String, u64, i64>(
1120                shard_id,
1121                Arc::new(StringSchema),
1122                Arc::new(StringSchema),
1123                Diagnostics::for_tests(),
1124            )
1125            .await
1126            .expect("codec mismatch");
1127        let mut read1 = client
1128            .open_leased_reader::<String, String, u64, i64>(
1129                shard_id,
1130                Arc::new(StringSchema),
1131                Arc::new(StringSchema),
1132                Diagnostics::for_tests(),
1133                true,
1134            )
1135            .await
1136            .expect("codec mismatch");
1137        let mut read2 = client
1138            .open_leased_reader::<String, String, u64, i64>(
1139                shard_id,
1140                Arc::new(StringSchema),
1141                Arc::new(StringSchema),
1142                Diagnostics::for_tests(),
1143                true,
1144            )
1145            .await
1146            .expect("codec mismatch");
1147        let mut write2 = client
1148            .open_writer::<String, String, u64, i64>(
1149                shard_id,
1150                Arc::new(StringSchema),
1151                Arc::new(StringSchema),
1152                Diagnostics::for_tests(),
1153            )
1154            .await
1155            .expect("codec mismatch");
1156
1157        write2.expect_compare_and_append(&data[..1], 0, 2).await;
1158        assert_eq!(
1159            read2.expect_snapshot_and_fetch(1).await,
1160            all_ok(&data[..1], 1)
1161        );
1162        write1.expect_compare_and_append(&data[1..], 2, 4).await;
1163        assert_eq!(read1.expect_snapshot_and_fetch(3).await, all_ok(&data, 3));
1164    }
1165
1166    #[mz_persist_proc::test(tokio::test)]
1167    #[cfg_attr(miri, ignore)] // too slow
1168    async fn invalid_usage(dyncfgs: ConfigUpdates) {
1169        let data = vec![
1170            (("1".to_owned(), "one".to_owned()), 1, 1),
1171            (("2".to_owned(), "two".to_owned()), 2, 1),
1172            (("3".to_owned(), "three".to_owned()), 3, 1),
1173        ];
1174
1175        let shard_id0 = "s00000000-0000-0000-0000-000000000000"
1176            .parse::<ShardId>()
1177            .expect("invalid shard id");
1178        let mut client = new_test_client(&dyncfgs).await;
1179
1180        let (mut write0, mut read0) = client
1181            .expect_open::<String, String, u64, i64>(shard_id0)
1182            .await;
1183
1184        write0.expect_compare_and_append(&data, 0, 4).await;
1185
1186        // InvalidUsage from PersistClient methods.
1187        {
1188            fn codecs(
1189                k: &str,
1190                v: &str,
1191                t: &str,
1192                d: &str,
1193            ) -> (String, String, String, String, Option<CodecConcreteType>) {
1194                (k.to_owned(), v.to_owned(), t.to_owned(), d.to_owned(), None)
1195            }
1196
1197            client.shared_states = Arc::new(StateCache::new_no_metrics());
1198            assert_eq!(
1199                client
1200                    .open::<Vec<u8>, String, u64, i64>(
1201                        shard_id0,
1202                        Arc::new(VecU8Schema),
1203                        Arc::new(StringSchema),
1204                        Diagnostics::for_tests(),
1205                        true,
1206                    )
1207                    .await
1208                    .unwrap_err(),
1209                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1210                    requested: codecs("Vec<u8>", "String", "u64", "i64"),
1211                    actual: codecs("String", "String", "u64", "i64"),
1212                }))
1213            );
1214            assert_eq!(
1215                client
1216                    .open::<String, Vec<u8>, u64, i64>(
1217                        shard_id0,
1218                        Arc::new(StringSchema),
1219                        Arc::new(VecU8Schema),
1220                        Diagnostics::for_tests(),
1221                        true,
1222                    )
1223                    .await
1224                    .unwrap_err(),
1225                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1226                    requested: codecs("String", "Vec<u8>", "u64", "i64"),
1227                    actual: codecs("String", "String", "u64", "i64"),
1228                }))
1229            );
1230            assert_eq!(
1231                client
1232                    .open::<String, String, i64, i64>(
1233                        shard_id0,
1234                        Arc::new(StringSchema),
1235                        Arc::new(StringSchema),
1236                        Diagnostics::for_tests(),
1237                        true,
1238                    )
1239                    .await
1240                    .unwrap_err(),
1241                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1242                    requested: codecs("String", "String", "i64", "i64"),
1243                    actual: codecs("String", "String", "u64", "i64"),
1244                }))
1245            );
1246            assert_eq!(
1247                client
1248                    .open::<String, String, u64, u64>(
1249                        shard_id0,
1250                        Arc::new(StringSchema),
1251                        Arc::new(StringSchema),
1252                        Diagnostics::for_tests(),
1253                        true,
1254                    )
1255                    .await
1256                    .unwrap_err(),
1257                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1258                    requested: codecs("String", "String", "u64", "u64"),
1259                    actual: codecs("String", "String", "u64", "i64"),
1260                }))
1261            );
1262
1263            // open_reader and open_writer end up using the same checks, so just
1264            // verify one type each to verify the plumbing instead of the full
1265            // set.
1266            assert_eq!(
1267                client
1268                    .open_leased_reader::<Vec<u8>, String, u64, i64>(
1269                        shard_id0,
1270                        Arc::new(VecU8Schema),
1271                        Arc::new(StringSchema),
1272                        Diagnostics::for_tests(),
1273                        true,
1274                    )
1275                    .await
1276                    .unwrap_err(),
1277                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1278                    requested: codecs("Vec<u8>", "String", "u64", "i64"),
1279                    actual: codecs("String", "String", "u64", "i64"),
1280                }))
1281            );
1282            assert_eq!(
1283                client
1284                    .open_writer::<Vec<u8>, String, u64, i64>(
1285                        shard_id0,
1286                        Arc::new(VecU8Schema),
1287                        Arc::new(StringSchema),
1288                        Diagnostics::for_tests(),
1289                    )
1290                    .await
1291                    .unwrap_err(),
1292                InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1293                    requested: codecs("Vec<u8>", "String", "u64", "i64"),
1294                    actual: codecs("String", "String", "u64", "i64"),
1295                }))
1296            );
1297        }
1298
1299        // InvalidUsage from ReadHandle methods.
1300        {
1301            let snap = read0
1302                .snapshot(Antichain::from_elem(3))
1303                .await
1304                .expect("cannot serve requested as_of");
1305
1306            let shard_id1 = "s11111111-1111-1111-1111-111111111111"
1307                .parse::<ShardId>()
1308                .expect("invalid shard id");
1309            let mut fetcher1 = client
1310                .create_batch_fetcher::<String, String, u64, i64>(
1311                    shard_id1,
1312                    Default::default(),
1313                    Default::default(),
1314                    false,
1315                    Diagnostics::for_tests(),
1316                )
1317                .await
1318                .unwrap();
1319            for part in snap {
1320                let (part, _lease) = part.into_exchangeable_part();
1321                let res = fetcher1.fetch_leased_part(part).await;
1322                assert_eq!(
1323                    res.unwrap_err(),
1324                    InvalidUsage::BatchNotFromThisShard {
1325                        batch_shard: shard_id0,
1326                        handle_shard: shard_id1,
1327                    }
1328                );
1329            }
1330        }
1331
1332        // InvalidUsage from WriteHandle methods.
1333        {
1334            let ts3 = &data[2];
1335            assert_eq!(ts3.1, 3);
1336            let ts3 = vec![ts3.clone()];
1337
1338            // WriteHandle::append also covers append_batch,
1339            // compare_and_append_batch, compare_and_append.
1340            assert_eq!(
1341                write0
1342                    .append(&ts3, Antichain::from_elem(4), Antichain::from_elem(5))
1343                    .await
1344                    .unwrap_err(),
1345                InvalidUsage::UpdateNotBeyondLower {
1346                    ts: 3,
1347                    lower: Antichain::from_elem(4),
1348                },
1349            );
1350            assert_eq!(
1351                write0
1352                    .append(&ts3, Antichain::from_elem(2), Antichain::from_elem(3))
1353                    .await
1354                    .unwrap_err(),
1355                InvalidUsage::UpdateBeyondUpper {
1356                    ts: 3,
1357                    expected_upper: Antichain::from_elem(3),
1358                },
1359            );
1360            // NB unlike the previous tests, this one has empty updates.
1361            assert_eq!(
1362                write0
1363                    .append(&data[..0], Antichain::from_elem(3), Antichain::from_elem(2))
1364                    .await
1365                    .unwrap_err(),
1366                InvalidUsage::InvalidBounds {
1367                    lower: Antichain::from_elem(3),
1368                    upper: Antichain::from_elem(2),
1369                },
1370            );
1371
1372            // Tests for the BatchBuilder.
1373            assert_eq!(
1374                write0
1375                    .builder(Antichain::from_elem(3))
1376                    .finish(Antichain::from_elem(2))
1377                    .await
1378                    .unwrap_err(),
1379                InvalidUsage::InvalidBounds {
1380                    lower: Antichain::from_elem(3),
1381                    upper: Antichain::from_elem(2)
1382                },
1383            );
1384            let batch = write0
1385                .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1386                .await
1387                .expect("invalid usage");
1388            assert_eq!(
1389                write0
1390                    .append_batch(batch, Antichain::from_elem(4), Antichain::from_elem(5))
1391                    .await
1392                    .unwrap_err(),
1393                InvalidUsage::InvalidBatchBounds {
1394                    batch_lower: Antichain::from_elem(3),
1395                    batch_upper: Antichain::from_elem(4),
1396                    append_lower: Antichain::from_elem(4),
1397                    append_upper: Antichain::from_elem(5),
1398                },
1399            );
1400            let batch = write0
1401                .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1402                .await
1403                .expect("invalid usage");
1404            assert_eq!(
1405                write0
1406                    .append_batch(batch, Antichain::from_elem(2), Antichain::from_elem(3))
1407                    .await
1408                    .unwrap_err(),
1409                InvalidUsage::InvalidBatchBounds {
1410                    batch_lower: Antichain::from_elem(3),
1411                    batch_upper: Antichain::from_elem(4),
1412                    append_lower: Antichain::from_elem(2),
1413                    append_upper: Antichain::from_elem(3),
1414                },
1415            );
1416            let batch = write0
1417                .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1418                .await
1419                .expect("invalid usage");
1420            // NB unlike the others, this one uses matches! because it's
1421            // non-deterministic (the key)
1422            assert!(matches!(
1423                write0
1424                    .append_batch(batch, Antichain::from_elem(3), Antichain::from_elem(3))
1425                    .await
1426                    .unwrap_err(),
1427                InvalidUsage::InvalidEmptyTimeInterval { .. }
1428            ));
1429        }
1430    }
1431
1432    #[mz_persist_proc::test(tokio::test)]
1433    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1434    async fn multiple_shards(dyncfgs: ConfigUpdates) {
1435        let data1 = [
1436            (("1".to_owned(), "one".to_owned()), 1, 1),
1437            (("2".to_owned(), "two".to_owned()), 2, 1),
1438        ];
1439
1440        let data2 = [(("1".to_owned(), ()), 1, 1), (("2".to_owned(), ()), 2, 1)];
1441
1442        let client = new_test_client(&dyncfgs).await;
1443
1444        let (mut write1, mut read1) = client
1445            .expect_open::<String, String, u64, i64>(ShardId::new())
1446            .await;
1447
1448        // Different types, so that checks would fail in case we were not separating these
1449        // collections internally.
1450        let (mut write2, mut read2) = client
1451            .expect_open::<String, (), u64, i64>(ShardId::new())
1452            .await;
1453
1454        write1
1455            .expect_compare_and_append(&data1[..], u64::minimum(), 3)
1456            .await;
1457
1458        write2
1459            .expect_compare_and_append(&data2[..], u64::minimum(), 3)
1460            .await;
1461
1462        assert_eq!(
1463            read1.expect_snapshot_and_fetch(2).await,
1464            all_ok(&data1[..], 2)
1465        );
1466
1467        assert_eq!(
1468            read2.expect_snapshot_and_fetch(2).await,
1469            all_ok(&data2[..], 2)
1470        );
1471    }
1472
1473    #[mz_persist_proc::test(tokio::test)]
1474    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1475    async fn fetch_upper(dyncfgs: ConfigUpdates) {
1476        let data = [
1477            (("1".to_owned(), "one".to_owned()), 1, 1),
1478            (("2".to_owned(), "two".to_owned()), 2, 1),
1479        ];
1480
1481        let client = new_test_client(&dyncfgs).await;
1482
1483        let shard_id = ShardId::new();
1484
1485        let (mut write1, _read1) = client
1486            .expect_open::<String, String, u64, i64>(shard_id)
1487            .await;
1488
1489        let (mut write2, _read2) = client
1490            .expect_open::<String, String, u64, i64>(shard_id)
1491            .await;
1492
1493        write1
1494            .expect_append(&data[..], write1.upper().clone(), vec![3])
1495            .await;
1496
1497        // The shard-global upper does advance, even if this writer didn't advance its local upper.
1498        assert_eq!(write2.fetch_recent_upper().await, &Antichain::from_elem(3));
1499
1500        // The writer-local upper should advance, even if it was another writer
1501        // that advanced the frontier.
1502        assert_eq!(write2.upper(), &Antichain::from_elem(3));
1503    }
1504
1505    #[mz_persist_proc::test(tokio::test)]
1506    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1507    async fn append_with_invalid_upper(dyncfgs: ConfigUpdates) {
1508        let data = [
1509            (("1".to_owned(), "one".to_owned()), 1, 1),
1510            (("2".to_owned(), "two".to_owned()), 2, 1),
1511        ];
1512
1513        let client = new_test_client(&dyncfgs).await;
1514
1515        let shard_id = ShardId::new();
1516
1517        let (mut write, _read) = client
1518            .expect_open::<String, String, u64, i64>(shard_id)
1519            .await;
1520
1521        write
1522            .expect_append(&data[..], write.upper().clone(), vec![3])
1523            .await;
1524
1525        let data = [
1526            (("5".to_owned(), "fünf".to_owned()), 5, 1),
1527            (("6".to_owned(), "sechs".to_owned()), 6, 1),
1528        ];
1529        let res = write
1530            .append(
1531                data.iter(),
1532                Antichain::from_elem(5),
1533                Antichain::from_elem(7),
1534            )
1535            .await;
1536        assert_eq!(
1537            res,
1538            Ok(Err(UpperMismatch {
1539                expected: Antichain::from_elem(5),
1540                current: Antichain::from_elem(3)
1541            }))
1542        );
1543
1544        // Writing with an outdated upper updates the write handle's upper to the correct upper.
1545        assert_eq!(write.upper(), &Antichain::from_elem(3));
1546    }
1547
1548    // Make sure that the API structs are Sync + Send, so that they can be used in async tasks.
1549    // NOTE: This is a compile-time only test. If it compiles, we're good.
1550    #[allow(unused)]
1551    async fn sync_send(dyncfgs: ConfigUpdates) {
1552        mz_ore::test::init_logging();
1553
1554        fn is_send_sync<T: Send + Sync>(_x: T) -> bool {
1555            true
1556        }
1557
1558        let client = new_test_client(&dyncfgs).await;
1559
1560        let (write, read) = client
1561            .expect_open::<String, String, u64, i64>(ShardId::new())
1562            .await;
1563
1564        assert!(is_send_sync(client));
1565        assert!(is_send_sync(write));
1566        assert!(is_send_sync(read));
1567    }
1568
1569    #[mz_persist_proc::test(tokio::test)]
1570    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1571    async fn compare_and_append(dyncfgs: ConfigUpdates) {
1572        let data = vec![
1573            (("1".to_owned(), "one".to_owned()), 1, 1),
1574            (("2".to_owned(), "two".to_owned()), 2, 1),
1575            (("3".to_owned(), "three".to_owned()), 3, 1),
1576        ];
1577
1578        let id = ShardId::new();
1579        let client = new_test_client(&dyncfgs).await;
1580        let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1581
1582        let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1583
1584        assert_eq!(write1.upper(), &Antichain::from_elem(u64::minimum()));
1585        assert_eq!(write2.upper(), &Antichain::from_elem(u64::minimum()));
1586        assert_eq!(read.since(), &Antichain::from_elem(u64::minimum()));
1587
1588        // Write a [0,3) batch.
1589        write1
1590            .expect_compare_and_append(&data[..2], u64::minimum(), 3)
1591            .await;
1592        assert_eq!(write1.upper(), &Antichain::from_elem(3));
1593
1594        assert_eq!(
1595            read.expect_snapshot_and_fetch(2).await,
1596            all_ok(&data[..2], 2)
1597        );
1598
1599        // Try and write with a wrong expected upper.
1600        let res = write2
1601            .compare_and_append(
1602                &data[..2],
1603                Antichain::from_elem(u64::minimum()),
1604                Antichain::from_elem(3),
1605            )
1606            .await;
1607        assert_eq!(
1608            res,
1609            Ok(Err(UpperMismatch {
1610                expected: Antichain::from_elem(u64::minimum()),
1611                current: Antichain::from_elem(3)
1612            }))
1613        );
1614
1615        // A failed write updates our local cache of the shard upper.
1616        assert_eq!(write2.upper(), &Antichain::from_elem(3));
1617
1618        // Try again with a good expected upper.
1619        write2.expect_compare_and_append(&data[2..], 3, 4).await;
1620
1621        assert_eq!(write2.upper(), &Antichain::from_elem(4));
1622
1623        assert_eq!(read.expect_snapshot_and_fetch(3).await, all_ok(&data, 3));
1624    }
1625
1626    #[mz_persist_proc::test(tokio::test)]
1627    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1628    async fn overlapping_append(dyncfgs: ConfigUpdates) {
1629        mz_ore::test::init_logging_default("info");
1630
1631        let data = vec![
1632            (("1".to_owned(), "one".to_owned()), 1, 1),
1633            (("2".to_owned(), "two".to_owned()), 2, 1),
1634            (("3".to_owned(), "three".to_owned()), 3, 1),
1635            (("4".to_owned(), "vier".to_owned()), 4, 1),
1636            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1637        ];
1638
1639        let id = ShardId::new();
1640        let client = new_test_client(&dyncfgs).await;
1641
1642        let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1643
1644        let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1645
1646        // Grab a listener before we do any writing
1647        let mut listen = read.clone("").await.expect_listen(0).await;
1648
1649        // Write a [0,3) batch.
1650        write1
1651            .expect_append(&data[..2], write1.upper().clone(), vec![3])
1652            .await;
1653        assert_eq!(write1.upper(), &Antichain::from_elem(3));
1654
1655        // Write a [0,5) batch with the second writer.
1656        write2
1657            .expect_append(&data[..4], write2.upper().clone(), vec![5])
1658            .await;
1659        assert_eq!(write2.upper(), &Antichain::from_elem(5));
1660
1661        // Write a [3,6) batch with the first writer.
1662        write1
1663            .expect_append(&data[2..5], write1.upper().clone(), vec![6])
1664            .await;
1665        assert_eq!(write1.upper(), &Antichain::from_elem(6));
1666
1667        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1668
1669        assert_eq!(
1670            listen.read_until(&6).await,
1671            (all_ok(&data[..], 1), Antichain::from_elem(6))
1672        );
1673    }
1674
1675    // Appends need to be contiguous for a shard, meaning the lower of an appended batch must not
1676    // be in advance of the current shard upper.
1677    #[mz_persist_proc::test(tokio::test)]
1678    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1679    async fn contiguous_append(dyncfgs: ConfigUpdates) {
1680        let data = vec![
1681            (("1".to_owned(), "one".to_owned()), 1, 1),
1682            (("2".to_owned(), "two".to_owned()), 2, 1),
1683            (("3".to_owned(), "three".to_owned()), 3, 1),
1684            (("4".to_owned(), "vier".to_owned()), 4, 1),
1685            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1686        ];
1687
1688        let id = ShardId::new();
1689        let client = new_test_client(&dyncfgs).await;
1690
1691        let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1692
1693        // Write a [0,3) batch.
1694        write
1695            .expect_append(&data[..2], write.upper().clone(), vec![3])
1696            .await;
1697        assert_eq!(write.upper(), &Antichain::from_elem(3));
1698
1699        // Appending a non-contiguous batch should fail.
1700        // Write a [5,6) batch with the second writer.
1701        let result = write
1702            .append(
1703                &data[4..5],
1704                Antichain::from_elem(5),
1705                Antichain::from_elem(6),
1706            )
1707            .await;
1708        assert_eq!(
1709            result,
1710            Ok(Err(UpperMismatch {
1711                expected: Antichain::from_elem(5),
1712                current: Antichain::from_elem(3)
1713            }))
1714        );
1715
1716        // Fixing the lower to make the write contiguous should make the append succeed.
1717        write.expect_append(&data[2..5], vec![3], vec![6]).await;
1718        assert_eq!(write.upper(), &Antichain::from_elem(6));
1719
1720        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1721    }
1722
1723    // Per-writer appends can be non-contiguous, as long as appends to the shard from all writers
1724    // combined are contiguous.
1725    #[mz_persist_proc::test(tokio::test)]
1726    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1727    async fn noncontiguous_append_per_writer(dyncfgs: ConfigUpdates) {
1728        let data = vec![
1729            (("1".to_owned(), "one".to_owned()), 1, 1),
1730            (("2".to_owned(), "two".to_owned()), 2, 1),
1731            (("3".to_owned(), "three".to_owned()), 3, 1),
1732            (("4".to_owned(), "vier".to_owned()), 4, 1),
1733            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1734        ];
1735
1736        let id = ShardId::new();
1737        let client = new_test_client(&dyncfgs).await;
1738
1739        let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1740
1741        let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1742
1743        // Write a [0,3) batch with writer 1.
1744        write1
1745            .expect_append(&data[..2], write1.upper().clone(), vec![3])
1746            .await;
1747        assert_eq!(write1.upper(), &Antichain::from_elem(3));
1748
1749        // Write a [3,5) batch with writer 2.
1750        write2.upper = Antichain::from_elem(3);
1751        write2
1752            .expect_append(&data[2..4], write2.upper().clone(), vec![5])
1753            .await;
1754        assert_eq!(write2.upper(), &Antichain::from_elem(5));
1755
1756        // Write a [5,6) batch with writer 1.
1757        write1.upper = Antichain::from_elem(5);
1758        write1
1759            .expect_append(&data[4..5], write1.upper().clone(), vec![6])
1760            .await;
1761        assert_eq!(write1.upper(), &Antichain::from_elem(6));
1762
1763        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1764    }
1765
1766    // Compare_and_appends need to be contiguous for a shard, meaning the lower of an appended
1767    // batch needs to match the current shard upper.
1768    #[mz_persist_proc::test(tokio::test)]
1769    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1770    async fn contiguous_compare_and_append(dyncfgs: ConfigUpdates) {
1771        let data = vec![
1772            (("1".to_owned(), "one".to_owned()), 1, 1),
1773            (("2".to_owned(), "two".to_owned()), 2, 1),
1774            (("3".to_owned(), "three".to_owned()), 3, 1),
1775            (("4".to_owned(), "vier".to_owned()), 4, 1),
1776            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1777        ];
1778
1779        let id = ShardId::new();
1780        let client = new_test_client(&dyncfgs).await;
1781
1782        let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1783
1784        // Write a [0,3) batch.
1785        write.expect_compare_and_append(&data[..2], 0, 3).await;
1786        assert_eq!(write.upper(), &Antichain::from_elem(3));
1787
1788        // Appending a non-contiguous batch should fail.
1789        // Write a [5,6) batch with the second writer.
1790        let result = write
1791            .compare_and_append(
1792                &data[4..5],
1793                Antichain::from_elem(5),
1794                Antichain::from_elem(6),
1795            )
1796            .await;
1797        assert_eq!(
1798            result,
1799            Ok(Err(UpperMismatch {
1800                expected: Antichain::from_elem(5),
1801                current: Antichain::from_elem(3)
1802            }))
1803        );
1804
1805        // Writing with the correct expected upper to make the write contiguous should make the
1806        // append succeed.
1807        write.expect_compare_and_append(&data[2..5], 3, 6).await;
1808        assert_eq!(write.upper(), &Antichain::from_elem(6));
1809
1810        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1811    }
1812
1813    // Per-writer compare_and_appends can be non-contiguous, as long as appends to the shard from
1814    // all writers combined are contiguous.
1815    #[mz_persist_proc::test(tokio::test)]
1816    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1817    async fn noncontiguous_compare_and_append_per_writer(dyncfgs: ConfigUpdates) {
1818        let data = vec![
1819            (("1".to_owned(), "one".to_owned()), 1, 1),
1820            (("2".to_owned(), "two".to_owned()), 2, 1),
1821            (("3".to_owned(), "three".to_owned()), 3, 1),
1822            (("4".to_owned(), "vier".to_owned()), 4, 1),
1823            (("5".to_owned(), "cinque".to_owned()), 5, 1),
1824        ];
1825
1826        let id = ShardId::new();
1827        let client = new_test_client(&dyncfgs).await;
1828
1829        let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1830
1831        let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1832
1833        // Write a [0,3) batch with writer 1.
1834        write1.expect_compare_and_append(&data[..2], 0, 3).await;
1835        assert_eq!(write1.upper(), &Antichain::from_elem(3));
1836
1837        // Write a [3,5) batch with writer 2.
1838        write2.expect_compare_and_append(&data[2..4], 3, 5).await;
1839        assert_eq!(write2.upper(), &Antichain::from_elem(5));
1840
1841        // Write a [5,6) batch with writer 1.
1842        write1.expect_compare_and_append(&data[4..5], 5, 6).await;
1843        assert_eq!(write1.upper(), &Antichain::from_elem(6));
1844
1845        assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1846    }
1847
1848    #[mz_ore::test]
1849    fn fmt_ids() {
1850        assert_eq!(
1851            format!("{}", LeasedReaderId([0u8; 16])),
1852            "r00000000-0000-0000-0000-000000000000"
1853        );
1854        assert_eq!(
1855            format!("{:?}", LeasedReaderId([0u8; 16])),
1856            "LeasedReaderId(00000000-0000-0000-0000-000000000000)"
1857        );
1858    }
1859
1860    #[mz_persist_proc::test(tokio::test(flavor = "multi_thread"))]
1861    #[cfg_attr(miri, ignore)] // error: unsupported operation: integer-to-pointer casts and `ptr::from_exposed_addr` are not supported with `-Zmiri-strict-provenance`
1862    async fn concurrency(dyncfgs: ConfigUpdates) {
1863        let data = DataGenerator::small();
1864
1865        const NUM_WRITERS: usize = 2;
1866        let id = ShardId::new();
1867        let client = new_test_client(&dyncfgs).await;
1868        let mut handles = Vec::<mz_ore::task::JoinHandle<()>>::new();
1869        for idx in 0..NUM_WRITERS {
1870            let (data, client) = (data.clone(), client.clone());
1871
1872            let (batch_tx, mut batch_rx) = tokio::sync::mpsc::channel(1);
1873
1874            let client1 = client.clone();
1875            let handle = mz_ore::task::spawn(|| format!("writer-{}", idx), async move {
1876                let (write, _) = client1.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1877                let mut current_upper = 0;
1878                for batch in data.batches() {
1879                    let new_upper = match batch.get(batch.len() - 1) {
1880                        Some((_, max_ts, _)) => u64::decode(max_ts) + 1,
1881                        None => continue,
1882                    };
1883                    // Because we (intentionally) call open inside the task,
1884                    // some other writer may have raced ahead and already
1885                    // appended some data before this one was registered. As a
1886                    // result, this writer may not be starting with an upper of
1887                    // the initial empty antichain. This is nice because it
1888                    // mimics how a real HA source would work, but it means we
1889                    // have to skip any batches that have already been committed
1890                    // (otherwise our new_upper would be before our upper).
1891                    //
1892                    // Note however, that unlike a real source, our
1893                    // DataGenerator-derived batches are guaranteed to be
1894                    // chunked along the same boundaries. This means we don't
1895                    // have to consider partial batches when generating the
1896                    // updates below.
1897                    if PartialOrder::less_equal(&Antichain::from_elem(new_upper), write.upper()) {
1898                        continue;
1899                    }
1900
1901                    let current_upper_chain = Antichain::from_elem(current_upper);
1902                    current_upper = new_upper;
1903                    let new_upper_chain = Antichain::from_elem(new_upper);
1904                    let mut builder = write.builder(current_upper_chain);
1905
1906                    for ((k, v), t, d) in batch.iter() {
1907                        builder
1908                            .add(&k.to_vec(), &v.to_vec(), &u64::decode(t), &i64::decode(d))
1909                            .await
1910                            .expect("invalid usage");
1911                    }
1912
1913                    let batch = builder
1914                        .finish(new_upper_chain)
1915                        .await
1916                        .expect("invalid usage");
1917
1918                    match batch_tx.send(batch).await {
1919                        Ok(_) => (),
1920                        Err(e) => panic!("send error: {}", e),
1921                    }
1922                }
1923            });
1924            handles.push(handle);
1925
1926            let handle = mz_ore::task::spawn(|| format!("appender-{}", idx), async move {
1927                let (mut write, _) = client.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1928
1929                while let Some(batch) = batch_rx.recv().await {
1930                    let lower = batch.lower().clone();
1931                    let upper = batch.upper().clone();
1932                    write
1933                        .append_batch(batch, lower, upper)
1934                        .await
1935                        .expect("invalid usage")
1936                        .expect("unexpected upper");
1937                }
1938            });
1939            handles.push(handle);
1940        }
1941
1942        for handle in handles {
1943            let () = handle.await;
1944        }
1945
1946        let expected = data.records().collect::<Vec<_>>();
1947        let max_ts = expected.last().map(|(_, t, _)| *t).unwrap_or_default();
1948        let (_, mut read) = client.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1949        assert_eq!(
1950            read.expect_snapshot_and_fetch(max_ts).await,
1951            all_ok(expected.iter(), max_ts)
1952        );
1953    }
1954
1955    // Regression test for database-issues#3523. Snapshot with as_of >= upper would
1956    // immediately return the data currently available instead of waiting for
1957    // upper to advance past as_of.
1958    #[mz_persist_proc::test(tokio::test)]
1959    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1960    async fn regression_blocking_reads(dyncfgs: ConfigUpdates) {
1961        let waker = noop_waker();
1962        let mut cx = Context::from_waker(&waker);
1963
1964        let data = [
1965            (("1".to_owned(), "one".to_owned()), 1, 1),
1966            (("2".to_owned(), "two".to_owned()), 2, 1),
1967            (("3".to_owned(), "three".to_owned()), 3, 1),
1968        ];
1969
1970        let id = ShardId::new();
1971        let client = new_test_client(&dyncfgs).await;
1972        let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1973
1974        // Grab a listener as_of (aka gt) 1, which is not yet closed out.
1975        let mut listen = read.clone("").await.expect_listen(1).await;
1976        let mut listen_next = Box::pin(listen.fetch_next());
1977        // Intentionally don't await the listen_next, but instead manually poke
1978        // it for a while and assert that it doesn't resolve yet. See below for
1979        // discussion of some alternative ways of writing this unit test.
1980        for _ in 0..100 {
1981            assert!(
1982                Pin::new(&mut listen_next).poll(&mut cx).is_pending(),
1983                "listen::next unexpectedly ready"
1984            );
1985        }
1986
1987        // Write a [0,3) batch.
1988        write
1989            .expect_compare_and_append(&data[..2], u64::minimum(), 3)
1990            .await;
1991
1992        // The initial listen_next call should now be able to return data at 2.
1993        // It doesn't get 1 because the as_of was 1 and listen is strictly gt.
1994        assert_eq!(
1995            listen_next.await,
1996            vec![
1997                ListenEvent::Updates(vec![(("2".to_owned(), "two".to_owned()), 2, 1)]),
1998                ListenEvent::Progress(Antichain::from_elem(3)),
1999            ]
2000        );
2001
2002        // Grab a snapshot as_of 3, which is not yet closed out. Intentionally
2003        // don't await the snap, but instead manually poke it for a while and
2004        // assert that it doesn't resolve yet.
2005        //
2006        // An alternative to this would be to run it in a task and poll the task
2007        // with some timeout, but this would introduce a fixed test execution
2008        // latency of the timeout in the happy case. Plus, it would be
2009        // non-deterministic.
2010        //
2011        // Another alternative (that's potentially quite interesting!) would be
2012        // to separate creating a snapshot immediately (which would fail if
2013        // as_of was >= upper) from a bit of logic that retries until that case
2014        // is ready.
2015        let mut snap = Box::pin(read.expect_snapshot_and_fetch(3));
2016        for _ in 0..100 {
2017            assert!(
2018                Pin::new(&mut snap).poll(&mut cx).is_pending(),
2019                "snapshot unexpectedly ready"
2020            );
2021        }
2022
2023        // Now add the data at 3 and also unblock the snapshot.
2024        write.expect_compare_and_append(&data[2..], 3, 4).await;
2025
2026        // Read the snapshot and check that it got all the appropriate data.
2027        assert_eq!(snap.await, all_ok(&data[..], 3));
2028    }
2029
2030    #[mz_persist_proc::test(tokio::test)]
2031    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
2032    async fn heartbeat_task_shutdown(dyncfgs: ConfigUpdates) {
2033        // Verify that the ReadHandle and WriteHandle background heartbeat tasks
2034        // shut down cleanly after the handle is expired.
2035        let mut cache = new_test_client_cache(&dyncfgs);
2036        cache
2037            .cfg
2038            .set_config(&READER_LEASE_DURATION, Duration::from_millis(1));
2039        cache.cfg.writer_lease_duration = Duration::from_millis(1);
2040        let (_write, mut read) = cache
2041            .open(PersistLocation::new_in_mem())
2042            .await
2043            .expect("client construction failed")
2044            .expect_open::<(), (), u64, i64>(ShardId::new())
2045            .await;
2046        let read_unexpired_state = read
2047            .unexpired_state
2048            .take()
2049            .expect("handle should have unexpired state");
2050        read.expire().await;
2051        read_unexpired_state.heartbeat_task.await
2052    }
2053
2054    /// Verify that shard finalization works with empty shards, shards that have
2055    /// an empty write up to the empty upper Antichain.
2056    #[mz_persist_proc::test(tokio::test)]
2057    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
2058    async fn finalize_empty_shard(dyncfgs: ConfigUpdates) {
2059        let persist_client = new_test_client(&dyncfgs).await;
2060
2061        let shard_id = ShardId::new();
2062        pub const CRITICAL_SINCE: CriticalReaderId =
2063            CriticalReaderId([0, 0, 0, 0, 17, 17, 34, 34, 51, 51, 68, 68, 68, 68, 68, 68]);
2064
2065        let (mut write, mut read) = persist_client
2066            .expect_open::<(), (), u64, i64>(shard_id)
2067            .await;
2068
2069        // Advance since and upper to empty, which is a pre-requisite for
2070        // finalization/tombstoning.
2071        let () = read.downgrade_since(&Antichain::new()).await;
2072        let () = write.advance_upper(&Antichain::new()).await;
2073
2074        let mut since_handle: SinceHandle<(), (), u64, i64> = persist_client
2075            .open_critical_since(
2076                shard_id,
2077                CRITICAL_SINCE,
2078                Opaque::encode(&0u64),
2079                Diagnostics::for_tests(),
2080            )
2081            .await
2082            .expect("invalid persist usage");
2083
2084        let epoch = since_handle.opaque().clone();
2085        let new_since = Antichain::new();
2086        let downgrade = since_handle
2087            .compare_and_downgrade_since(&epoch, (&epoch, &new_since))
2088            .await;
2089
2090        assert!(
2091            downgrade.is_ok(),
2092            "downgrade of critical handle must succeed"
2093        );
2094
2095        let finalize = persist_client
2096            .finalize_shard::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2097            .await;
2098
2099        assert_ok!(finalize, "finalization must succeed");
2100
2101        let is_finalized = persist_client
2102            .is_finalized::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2103            .await
2104            .expect("invalid persist usage");
2105        assert!(is_finalized, "shard must still be finalized");
2106    }
2107
2108    /// Verify that shard finalization works with shards that had some data
2109    /// written to them, plus then an empty batch to bring their upper to the
2110    /// empty Antichain.
2111    #[mz_persist_proc::test(tokio::test)]
2112    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
2113    async fn finalize_shard(dyncfgs: ConfigUpdates) {
2114        const DATA: &[(((), ()), u64, i64)] = &[(((), ()), 0, 1)];
2115        let persist_client = new_test_client(&dyncfgs).await;
2116
2117        let shard_id = ShardId::new();
2118        pub const CRITICAL_SINCE: CriticalReaderId =
2119            CriticalReaderId([0, 0, 0, 0, 17, 17, 34, 34, 51, 51, 68, 68, 68, 68, 68, 68]);
2120
2121        let (mut write, mut read) = persist_client
2122            .expect_open::<(), (), u64, i64>(shard_id)
2123            .await;
2124
2125        // Write some data.
2126        let () = write
2127            .compare_and_append(DATA, Antichain::from_elem(0), Antichain::from_elem(1))
2128            .await
2129            .expect("usage should be valid")
2130            .expect("upper should match");
2131
2132        // Advance since and upper to empty, which is a pre-requisite for
2133        // finalization/tombstoning.
2134        let () = read.downgrade_since(&Antichain::new()).await;
2135        let () = write.advance_upper(&Antichain::new()).await;
2136
2137        let mut since_handle: SinceHandle<(), (), u64, i64> = persist_client
2138            .open_critical_since(
2139                shard_id,
2140                CRITICAL_SINCE,
2141                Opaque::encode(&0u64),
2142                Diagnostics::for_tests(),
2143            )
2144            .await
2145            .expect("invalid persist usage");
2146
2147        let epoch = since_handle.opaque().clone();
2148        let new_since = Antichain::new();
2149        let downgrade = since_handle
2150            .compare_and_downgrade_since(&epoch, (&epoch, &new_since))
2151            .await;
2152
2153        assert!(
2154            downgrade.is_ok(),
2155            "downgrade of critical handle must succeed"
2156        );
2157
2158        let finalize = persist_client
2159            .finalize_shard::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2160            .await;
2161
2162        assert_ok!(finalize, "finalization must succeed");
2163
2164        let is_finalized = persist_client
2165            .is_finalized::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2166            .await
2167            .expect("invalid persist usage");
2168        assert!(is_finalized, "shard must still be finalized");
2169    }
2170
2171    proptest! {
2172        #![proptest_config(ProptestConfig::with_cases(4096))]
2173
2174        #[mz_ore::test]
2175        #[cfg_attr(miri, ignore)] // too slow
2176        fn shard_id_protobuf_roundtrip(expect in any::<ShardId>() ) {
2177            let actual = protobuf_roundtrip::<_, String>(&expect);
2178            assert_ok!(actual);
2179            assert_eq!(actual.unwrap(), expect);
2180        }
2181    }
2182}