1#![cfg_attr(
16 not(feature = "fuzzing"),
17 warn(missing_docs, missing_debug_implementations)
18)]
19#![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 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 pub use crate::internal::metrics::{
81 Metrics, SinkMetrics, SinkWorkerMetrics, UpdateDelta, encode_ts_metric,
82 };
83}
84pub mod operators {
85 use mz_dyncfg::{Config, ParameterScope};
88
89 pub mod shard_source;
90
91 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#[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
119mod 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
143pub const BUILD_INFO: BuildInfo = build_info!();
145
146pub use mz_persist_types::{PersistLocation, ShardId};
148
149pub use crate::internal::encoding::Schemas;
150
151#[derive(Clone, Debug)]
154pub struct Diagnostics {
155 pub shard_name: String,
157 pub handle_purpose: String,
159}
160
161impl Diagnostics {
162 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 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#[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 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 Ok(PersistClient {
228 cfg,
229 blob,
230 consensus,
231 metrics,
232 isolated_runtime,
233 shared_states,
234 pubsub_sender,
235 })
236 }
237
238 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 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 #[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 #[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 #[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 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 #[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 #[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 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 #[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 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 #[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 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 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 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 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 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 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 #[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 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 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 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 #[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 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 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 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)] 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
1077 .expect_append(&data[..2], write.upper().clone(), vec![3])
1078 .await;
1079 assert_eq!(write.upper(), &Antichain::from_elem(3));
1080
1081 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
1091 .expect_append(&data[2..], write.upper().clone(), vec![4])
1092 .await;
1093 assert_eq!(write.upper(), &Antichain::from_elem(4));
1094
1095 assert_eq!(
1097 listen.read_until(&4).await,
1098 (all_ok(&data[1..], 1), Antichain::from_elem(4))
1099 );
1100
1101 read.downgrade_since(&Antichain::from_elem(2)).await;
1103 assert_eq!(read.since(), &Antichain::from_elem(2));
1104 }
1105
1106 #[mz_persist_proc::test(tokio::test)]
1108 #[cfg_attr(miri, ignore)] 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)] 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 {
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 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 {
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 {
1334 let ts3 = &data[2];
1335 assert_eq!(ts3.1, 3);
1336 let ts3 = vec![ts3.clone()];
1337
1338 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 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 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 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)] 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 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)] 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 assert_eq!(write2.fetch_recent_upper().await, &Antichain::from_elem(3));
1499
1500 assert_eq!(write2.upper(), &Antichain::from_elem(3));
1503 }
1504
1505 #[mz_persist_proc::test(tokio::test)]
1506 #[cfg_attr(miri, ignore)] 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 assert_eq!(write.upper(), &Antichain::from_elem(3));
1546 }
1547
1548 #[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)] 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 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 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 assert_eq!(write2.upper(), &Antichain::from_elem(3));
1617
1618 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)] 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 let mut listen = read.clone("").await.expect_listen(0).await;
1648
1649 write1
1651 .expect_append(&data[..2], write1.upper().clone(), vec![3])
1652 .await;
1653 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1654
1655 write2
1657 .expect_append(&data[..4], write2.upper().clone(), vec![5])
1658 .await;
1659 assert_eq!(write2.upper(), &Antichain::from_elem(5));
1660
1661 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 #[mz_persist_proc::test(tokio::test)]
1678 #[cfg_attr(miri, ignore)] 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
1695 .expect_append(&data[..2], write.upper().clone(), vec![3])
1696 .await;
1697 assert_eq!(write.upper(), &Antichain::from_elem(3));
1698
1699 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 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 #[mz_persist_proc::test(tokio::test)]
1726 #[cfg_attr(miri, ignore)] 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 write1
1745 .expect_append(&data[..2], write1.upper().clone(), vec![3])
1746 .await;
1747 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1748
1749 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 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 #[mz_persist_proc::test(tokio::test)]
1769 #[cfg_attr(miri, ignore)] 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.expect_compare_and_append(&data[..2], 0, 3).await;
1786 assert_eq!(write.upper(), &Antichain::from_elem(3));
1787
1788 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 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 #[mz_persist_proc::test(tokio::test)]
1816 #[cfg_attr(miri, ignore)] 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 write1.expect_compare_and_append(&data[..2], 0, 3).await;
1835 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1836
1837 write2.expect_compare_and_append(&data[2..4], 3, 5).await;
1839 assert_eq!(write2.upper(), &Antichain::from_elem(5));
1840
1841 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)] 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 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 #[mz_persist_proc::test(tokio::test)]
1959 #[cfg_attr(miri, ignore)] 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 let mut listen = read.clone("").await.expect_listen(1).await;
1976 let mut listen_next = Box::pin(listen.fetch_next());
1977 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
1989 .expect_compare_and_append(&data[..2], u64::minimum(), 3)
1990 .await;
1991
1992 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 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 write.expect_compare_and_append(&data[2..], 3, 4).await;
2025
2026 assert_eq!(snap.await, all_ok(&data[..], 3));
2028 }
2029
2030 #[mz_persist_proc::test(tokio::test)]
2031 #[cfg_attr(miri, ignore)] async fn heartbeat_task_shutdown(dyncfgs: ConfigUpdates) {
2033 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 #[mz_persist_proc::test(tokio::test)]
2057 #[cfg_attr(miri, ignore)] 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 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 #[mz_persist_proc::test(tokio::test)]
2112 #[cfg_attr(miri, ignore)] 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 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 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)] 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}