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;
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 );
99}
100pub mod read;
101pub mod rpc;
102pub mod schema;
103pub mod stats;
104pub mod usage;
105pub mod write;
106
107#[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
118mod 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
142pub const BUILD_INFO: BuildInfo = build_info!();
144
145pub use mz_persist_types::{PersistLocation, ShardId};
147
148pub use crate::internal::encoding::Schemas;
149
150#[derive(Clone, Debug)]
153pub struct Diagnostics {
154 pub shard_name: String,
156 pub handle_purpose: String,
158}
159
160impl Diagnostics {
161 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 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#[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 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 Ok(PersistClient {
227 cfg,
228 blob,
229 consensus,
230 metrics,
231 isolated_runtime,
232 shared_states,
233 pubsub_sender,
234 })
235 }
236
237 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 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 #[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 #[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 #[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 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 #[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 #[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 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 #[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 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 #[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 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 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 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 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 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 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 #[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 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 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 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 #[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 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 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 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)] 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
1078 .expect_append(&data[..2], write.upper().clone(), vec![3])
1079 .await;
1080 assert_eq!(write.upper(), &Antichain::from_elem(3));
1081
1082 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
1092 .expect_append(&data[2..], write.upper().clone(), vec![4])
1093 .await;
1094 assert_eq!(write.upper(), &Antichain::from_elem(4));
1095
1096 assert_eq!(
1098 listen.read_until(&4).await,
1099 (all_ok(&data[1..], 1), Antichain::from_elem(4))
1100 );
1101
1102 read.downgrade_since(&Antichain::from_elem(2)).await;
1104 assert_eq!(read.since(), &Antichain::from_elem(2));
1105 }
1106
1107 #[mz_persist_proc::test(tokio::test)]
1109 #[cfg_attr(miri, ignore)] 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)] 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 {
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 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 {
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 {
1335 let ts3 = &data[2];
1336 assert_eq!(ts3.1, 3);
1337 let ts3 = vec![ts3.clone()];
1338
1339 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 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 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 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)] 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 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)] 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 assert_eq!(write2.fetch_recent_upper().await, &Antichain::from_elem(3));
1500
1501 assert_eq!(write2.upper(), &Antichain::from_elem(3));
1504 }
1505
1506 #[mz_persist_proc::test(tokio::test)]
1507 #[cfg_attr(miri, ignore)] 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 assert_eq!(write.upper(), &Antichain::from_elem(3));
1547 }
1548
1549 #[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)] 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 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 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 assert_eq!(write2.upper(), &Antichain::from_elem(3));
1618
1619 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)] 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 let mut listen = read.clone("").await.expect_listen(0).await;
1649
1650 write1
1652 .expect_append(&data[..2], write1.upper().clone(), vec![3])
1653 .await;
1654 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1655
1656 write2
1658 .expect_append(&data[..4], write2.upper().clone(), vec![5])
1659 .await;
1660 assert_eq!(write2.upper(), &Antichain::from_elem(5));
1661
1662 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 #[mz_persist_proc::test(tokio::test)]
1679 #[cfg_attr(miri, ignore)] 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
1696 .expect_append(&data[..2], write.upper().clone(), vec![3])
1697 .await;
1698 assert_eq!(write.upper(), &Antichain::from_elem(3));
1699
1700 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 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 #[mz_persist_proc::test(tokio::test)]
1727 #[cfg_attr(miri, ignore)] 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 write1
1746 .expect_append(&data[..2], write1.upper().clone(), vec![3])
1747 .await;
1748 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1749
1750 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 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 #[mz_persist_proc::test(tokio::test)]
1770 #[cfg_attr(miri, ignore)] 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.expect_compare_and_append(&data[..2], 0, 3).await;
1787 assert_eq!(write.upper(), &Antichain::from_elem(3));
1788
1789 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 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 #[mz_persist_proc::test(tokio::test)]
1817 #[cfg_attr(miri, ignore)] 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 write1.expect_compare_and_append(&data[..2], 0, 3).await;
1836 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1837
1838 write2.expect_compare_and_append(&data[2..4], 3, 5).await;
1840 assert_eq!(write2.upper(), &Antichain::from_elem(5));
1841
1842 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)] 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 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 #[mz_persist_proc::test(tokio::test)]
1960 #[cfg_attr(miri, ignore)] 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 let mut listen = read.clone("").await.expect_listen(1).await;
1977 let mut listen_next = Box::pin(listen.fetch_next());
1978 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
1990 .expect_compare_and_append(&data[..2], u64::minimum(), 3)
1991 .await;
1992
1993 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 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 write.expect_compare_and_append(&data[2..], 3, 4).await;
2026
2027 assert_eq!(snap.await, all_ok(&data[..], 3));
2029 }
2030
2031 #[mz_persist_proc::test(tokio::test)]
2032 #[cfg_attr(miri, ignore)] async fn heartbeat_task_shutdown(dyncfgs: ConfigUpdates) {
2034 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 #[mz_persist_proc::test(tokio::test)]
2058 #[cfg_attr(miri, ignore)] 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 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 #[mz_persist_proc::test(tokio::test)]
2113 #[cfg_attr(miri, ignore)] 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 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 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)] 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}