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