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 BackpressureMetrics, 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))]
557 pub fn batch_builder<K, V, T, D>(
558 &self,
559 shard_id: ShardId,
560 shard_name: &str,
561 write_schemas: Schemas<K, V>,
562 lower: Antichain<T>,
563 max_runs: Option<usize>,
564 ) -> BatchBuilder<K, V, T, D>
565 where
566 K: Debug + Codec,
567 V: Debug + Codec,
568 T: Timestamp + Lattice + Codec64 + TotalOrder + Sync,
569 D: Monoid + Ord + Codec64 + Send + Sync,
570 {
571 let mut compact_cfg = CompactConfig::new(&self.cfg, shard_id);
572 if let Some(max_runs) = max_runs {
573 compact_cfg.batch.max_runs = Some(max_runs);
574 }
575 WriteHandle::builder_inner(
576 &self.cfg,
577 compact_cfg,
578 Arc::clone(&self.metrics),
579 self.metrics.shards.shard(&shard_id, shard_name),
580 &self.metrics.user,
581 Arc::clone(&self.isolated_runtime),
582 Arc::clone(&self.blob),
583 shard_id,
584 write_schemas,
585 lower,
586 )
587 }
588
589 pub fn batch_from_transmittable_batch<K, V, T, D>(
598 &self,
599 shard_id: &ShardId,
600 batch: ProtoBatch,
601 ) -> Batch<K, V, T, D>
602 where
603 K: Debug + Codec,
604 V: Debug + Codec,
605 T: Timestamp + Lattice + Codec64 + Sync,
606 D: Monoid + Ord + Codec64 + Send + Sync,
607 {
608 let batch_shard_id: ShardId = batch
609 .shard_id
610 .into_rust()
611 .expect("valid transmittable batch");
612 assert_eq!(&batch_shard_id, shard_id);
613
614 let shard_metrics = self.metrics.shards.shard(shard_id, "peek_stash");
615
616 let ret = Batch {
617 batch_delete_enabled: BATCH_DELETE_ENABLED.get(&self.cfg),
618 metrics: Arc::clone(&self.metrics),
619 shard_metrics,
620 version: Version::parse(&batch.version).expect("valid transmittable batch"),
621 schemas: (batch.key_schema, batch.val_schema),
622 batch: batch
623 .batch
624 .into_rust_if_some("ProtoBatch::batch")
625 .expect("valid transmittable batch"),
626 blob: Arc::clone(&self.blob),
627 _phantom: std::marker::PhantomData,
628 };
629
630 assert_eq!(&ret.shard_id(), shard_id);
631 ret
632 }
633
634 #[allow(clippy::unused_async)]
649 pub async fn read_batches_consolidated<K, V, T, D>(
650 &mut self,
651 shard_id: ShardId,
652 as_of: Antichain<T>,
653 read_schemas: Schemas<K, V>,
654 batches: Vec<Batch<K, V, T, D>>,
655 should_fetch_part: impl for<'a> Fn(Option<&'a LazyPartStats>) -> bool,
656 memory_budget_bytes: usize,
657 ) -> Result<Cursor<K, V, T, D, Vec<Batch<K, V, T, D>>>, Since<T>>
658 where
659 K: Debug + Codec + Ord,
660 V: Debug + Codec + Ord,
661 T: Timestamp + Lattice + Codec64 + TotalOrder + Sync,
662 D: Monoid + Ord + Codec64 + Send + Sync,
663 {
664 let shard_metrics = self.metrics.shards.shard(&shard_id, "peek_stash");
665
666 let hollow_batches = batches.iter().map(|b| b.batch.clone()).collect_vec();
667
668 ReadHandle::read_batches_consolidated(
669 &self.cfg,
670 Arc::clone(&self.metrics),
671 shard_metrics,
672 self.metrics.read.snapshot.clone(),
673 Arc::clone(&self.blob),
674 shard_id,
675 as_of,
676 read_schemas,
677 &hollow_batches,
678 batches,
679 should_fetch_part,
680 memory_budget_bytes,
681 )
682 }
683
684 pub async fn get_schema<K, V, T, D>(
686 &self,
687 shard_id: ShardId,
688 schema_id: SchemaId,
689 diagnostics: Diagnostics,
690 ) -> Result<Option<(K::Schema, V::Schema)>, InvalidUsage<T>>
691 where
692 K: Debug + Codec,
693 V: Debug + Codec,
694 T: Timestamp + Lattice + Codec64 + Sync,
695 D: Monoid + Codec64 + Send + Sync,
696 {
697 let machine = self
698 .make_machine::<K, V, T, D>(shard_id, diagnostics)
699 .await?;
700 Ok(machine.get_schema(schema_id))
701 }
702
703 pub async fn latest_schema<K, V, T, D>(
705 &self,
706 shard_id: ShardId,
707 diagnostics: Diagnostics,
708 ) -> Result<Option<(SchemaId, K::Schema, V::Schema)>, InvalidUsage<T>>
709 where
710 K: Debug + Codec,
711 V: Debug + Codec,
712 T: Timestamp + Lattice + Codec64 + Sync,
713 D: Monoid + Codec64 + Send + Sync,
714 {
715 let machine = self
716 .make_machine::<K, V, T, D>(shard_id, diagnostics)
717 .await?;
718 Ok(machine.latest_schema())
719 }
720
721 pub async fn recent_upper<K, V, T, D>(
732 &self,
733 shard_id: ShardId,
734 diagnostics: Diagnostics,
735 ) -> Result<Antichain<T>, InvalidUsage<T>>
736 where
737 K: Debug + Codec,
738 V: Debug + Codec,
739 T: Timestamp + Lattice + Codec64 + Sync,
740 D: Monoid + Codec64 + Send + Sync,
741 {
742 let machine = self
743 .make_machine::<K, V, T, D>(shard_id, diagnostics)
744 .await?;
745 Ok(machine.applier.fetch_upper(|upper| upper.clone()).await)
746 }
747
748 pub async fn register_schema<K, V, T, D>(
760 &self,
761 shard_id: ShardId,
762 key_schema: &K::Schema,
763 val_schema: &V::Schema,
764 diagnostics: Diagnostics,
765 ) -> Result<Option<SchemaId>, InvalidUsage<T>>
766 where
767 K: Debug + Codec,
768 V: Debug + Codec,
769 T: Timestamp + Lattice + Codec64 + Sync,
770 D: Monoid + Codec64 + Send + Sync,
771 {
772 let machine = self
773 .make_machine::<K, V, T, D>(shard_id, diagnostics)
774 .await?;
775 let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
776
777 let (schema_id, maintenance) = machine.register_schema(key_schema, val_schema).await;
778 maintenance.start_performing(&machine, &gc);
779
780 Ok(schema_id)
781 }
782
783 pub async fn compare_and_evolve_schema<K, V, T, D>(
797 &self,
798 shard_id: ShardId,
799 expected: SchemaId,
800 key_schema: &K::Schema,
801 val_schema: &V::Schema,
802 diagnostics: Diagnostics,
803 ) -> Result<CaESchema<K, V>, InvalidUsage<T>>
804 where
805 K: Debug + Codec,
806 V: Debug + Codec,
807 T: Timestamp + Lattice + Codec64 + Sync,
808 D: Monoid + Codec64 + Send + Sync,
809 {
810 let machine = self
811 .make_machine::<K, V, T, D>(shard_id, diagnostics)
812 .await?;
813 let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
814 let (res, maintenance) = machine
815 .compare_and_evolve_schema(expected, key_schema, val_schema)
816 .await;
817 maintenance.start_performing(&machine, &gc);
818 Ok(res)
819 }
820
821 pub async fn is_finalized<K, V, T, D>(
825 &self,
826 shard_id: ShardId,
827 diagnostics: Diagnostics,
828 ) -> Result<bool, InvalidUsage<T>>
829 where
830 K: Debug + Codec,
831 V: Debug + Codec,
832 T: Timestamp + Lattice + Codec64 + Sync,
833 D: Monoid + Codec64 + Send + Sync,
834 {
835 let machine = self
836 .make_machine::<K, V, T, D>(shard_id, diagnostics)
837 .await?;
838 Ok(machine.is_finalized())
839 }
840
841 #[instrument(level = "debug", fields(shard = %shard_id))]
852 pub async fn finalize_shard<K, V, T, D>(
853 &self,
854 shard_id: ShardId,
855 diagnostics: Diagnostics,
856 ) -> Result<(), InvalidUsage<T>>
857 where
858 K: Debug + Codec,
859 V: Debug + Codec,
860 T: Timestamp + Lattice + Codec64 + Sync,
861 D: Monoid + Codec64 + Send + Sync,
862 {
863 let machine = self
864 .make_machine::<K, V, T, D>(shard_id, diagnostics)
865 .await?;
866
867 let maintenance = machine.become_tombstone().await?;
868 let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
869
870 let () = maintenance.perform(&machine, &gc).await;
871
872 Ok(())
873 }
874
875 pub async fn upgrade_version<K, V, T, D>(
878 &self,
879 shard_id: ShardId,
880 diagnostics: Diagnostics,
881 ) -> Result<(), InvalidUsage<T>>
882 where
883 K: Debug + Codec,
884 V: Debug + Codec,
885 T: Timestamp + Lattice + Codec64 + Sync,
886 D: Monoid + Codec64 + Send + Sync,
887 {
888 let machine = self
889 .make_machine::<K, V, T, D>(shard_id, diagnostics)
890 .await?;
891
892 match machine.upgrade_version().await {
893 Ok(maintenance) => {
894 let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
895 let () = maintenance.perform(&machine, &gc).await;
896 Ok(())
897 }
898 Err(version) => Err(InvalidUsage::IncompatibleVersion { version }),
899 }
900 }
901
902 pub async fn inspect_shard<T: Timestamp + Lattice + Codec64>(
908 &self,
909 shard_id: &ShardId,
910 ) -> Result<impl serde::Serialize, anyhow::Error> {
911 let state_versions = StateVersions::new(
912 self.cfg.clone(),
913 Arc::clone(&self.consensus),
914 Arc::clone(&self.blob),
915 Arc::clone(&self.metrics),
916 );
917 let versions = state_versions.fetch_all_live_diffs(shard_id).await;
921 if versions.is_empty() {
922 return Err(anyhow::anyhow!("{} does not exist", shard_id));
923 }
924 let state = state_versions
925 .fetch_current_state::<T>(shard_id, versions)
926 .await;
927 let state = state.check_ts_codec(shard_id)?;
928 Ok(state)
929 }
930
931 #[cfg(test)]
933 #[track_caller]
934 pub async fn expect_open<K, V, T, D>(
935 &self,
936 shard_id: ShardId,
937 ) -> (WriteHandle<K, V, T, D>, ReadHandle<K, V, T, D>)
938 where
939 K: Debug + Codec,
940 V: Debug + Codec,
941 T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
942 D: Monoid + Ord + Codec64 + Send + Sync,
943 K::Schema: Default,
944 V::Schema: Default,
945 {
946 self.open(
947 shard_id,
948 Arc::new(K::Schema::default()),
949 Arc::new(V::Schema::default()),
950 Diagnostics::for_tests(),
951 true,
952 )
953 .await
954 .expect("codec mismatch")
955 }
956
957 pub fn metrics(&self) -> &Arc<Metrics> {
961 &self.metrics
962 }
963}
964
965#[cfg(test)]
966mod tests {
967 use std::future::Future;
968 use std::pin::Pin;
969 use std::task::Context;
970 use std::time::Duration;
971
972 use differential_dataflow::consolidation::consolidate_updates;
973 use differential_dataflow::lattice::Lattice;
974 use futures_task::noop_waker;
975 use mz_dyncfg::ConfigUpdates;
976 use mz_ore::assert_ok;
977 use mz_persist::indexed::encoding::BlobTraceBatchPart;
978 use mz_persist::workload::DataGenerator;
979 use mz_persist_types::codec_impls::{StringSchema, VecU8Schema};
980 use mz_proto::protobuf_roundtrip;
981 use proptest::prelude::*;
982 use timely::order::PartialOrder;
983 use timely::progress::Antichain;
984
985 use crate::batch::BLOB_TARGET_SIZE;
986 use crate::cache::PersistClientCache;
987 use crate::cfg::BATCH_BUILDER_MAX_OUTSTANDING_PARTS;
988 use crate::critical::Opaque;
989 use crate::error::{CodecConcreteType, CodecMismatch, UpperMismatch};
990 use crate::internal::paths::BlobKey;
991 use crate::read::ListenEvent;
992
993 use super::*;
994
995 pub fn new_test_client_cache(dyncfgs: &ConfigUpdates) -> PersistClientCache {
996 let mut cache = PersistClientCache::new_no_metrics();
999 cache.cfg.set_config(&BLOB_TARGET_SIZE, 10);
1000 cache
1001 .cfg
1002 .set_config(&BATCH_BUILDER_MAX_OUTSTANDING_PARTS, 1);
1003 dyncfgs.apply(cache.cfg());
1004
1005 cache.cfg.compaction_enabled = true;
1007 cache
1008 }
1009
1010 pub async fn new_test_client(dyncfgs: &ConfigUpdates) -> PersistClient {
1011 let cache = new_test_client_cache(dyncfgs);
1012 cache
1013 .open(PersistLocation::new_in_mem())
1014 .await
1015 .expect("client construction failed")
1016 }
1017
1018 pub fn all_ok<'a, K, V, T, D, I>(iter: I, as_of: T) -> Vec<((K, V), T, D)>
1019 where
1020 K: Ord + Clone + 'a,
1021 V: Ord + Clone + 'a,
1022 T: Timestamp + Lattice + Clone + 'a,
1023 D: Monoid + Clone + 'a,
1024 I: IntoIterator<Item = &'a ((K, V), T, D)>,
1025 {
1026 let as_of = Antichain::from_elem(as_of);
1027 let mut ret = iter
1028 .into_iter()
1029 .map(|((k, v), t, d)| {
1030 let mut t = t.clone();
1031 t.advance_by(as_of.borrow());
1032 ((k.clone(), v.clone()), t, d.clone())
1033 })
1034 .collect();
1035 consolidate_updates(&mut ret);
1036 ret
1037 }
1038
1039 pub async fn expect_fetch_part<K, V, T, D>(
1040 blob: &dyn Blob,
1041 key: &BlobKey,
1042 metrics: &Metrics,
1043 read_schemas: &Schemas<K, V>,
1044 ) -> (BlobTraceBatchPart<T>, Vec<((K, V), T, D)>)
1045 where
1046 K: Codec + Clone,
1047 V: Codec + Clone,
1048 T: Timestamp + Codec64,
1049 D: Codec64,
1050 {
1051 let value = blob
1052 .get(key)
1053 .await
1054 .expect("failed to fetch part")
1055 .expect("missing part");
1056 let mut part =
1057 BlobTraceBatchPart::decode(&value, &metrics.columnar).expect("failed to decode part");
1058 let structured = part
1059 .updates
1060 .into_part::<K, V>(&*read_schemas.key, &*read_schemas.val);
1061 let updates = structured
1062 .decode_iter::<K, V, T, D>(&*read_schemas.key, &*read_schemas.val)
1063 .expect("structured data")
1064 .collect();
1065 (part, updates)
1066 }
1067
1068 #[mz_persist_proc::test(tokio::test)]
1069 #[cfg_attr(miri, ignore)] async fn sanity_check(dyncfgs: ConfigUpdates) {
1071 let data = [
1072 (("1".to_owned(), "one".to_owned()), 1, 1),
1073 (("2".to_owned(), "two".to_owned()), 2, 1),
1074 (("3".to_owned(), "three".to_owned()), 3, 1),
1075 ];
1076
1077 let (mut write, mut read) = new_test_client(&dyncfgs)
1078 .await
1079 .expect_open::<String, String, u64, i64>(ShardId::new())
1080 .await;
1081 assert_eq!(write.upper(), &Antichain::from_elem(u64::minimum()));
1082 assert_eq!(read.since(), &Antichain::from_elem(u64::minimum()));
1083
1084 write
1086 .expect_append(&data[..2], write.upper().clone(), vec![3])
1087 .await;
1088 assert_eq!(write.upper(), &Antichain::from_elem(3));
1089
1090 assert_eq!(
1092 read.expect_snapshot_and_fetch(1).await,
1093 all_ok(&data[..1], 1)
1094 );
1095
1096 let mut listen = read.clone("").await.expect_listen(1).await;
1097
1098 write
1100 .expect_append(&data[2..], write.upper().clone(), vec![4])
1101 .await;
1102 assert_eq!(write.upper(), &Antichain::from_elem(4));
1103
1104 assert_eq!(
1106 listen.read_until(&4).await,
1107 (all_ok(&data[1..], 1), Antichain::from_elem(4))
1108 );
1109
1110 read.downgrade_since(&Antichain::from_elem(2)).await;
1112 assert_eq!(read.since(), &Antichain::from_elem(2));
1113 }
1114
1115 #[mz_persist_proc::test(tokio::test)]
1117 #[cfg_attr(miri, ignore)] async fn open_reader_writer(dyncfgs: ConfigUpdates) {
1119 let data = vec![
1120 (("1".to_owned(), "one".to_owned()), 1, 1),
1121 (("2".to_owned(), "two".to_owned()), 2, 1),
1122 (("3".to_owned(), "three".to_owned()), 3, 1),
1123 ];
1124
1125 let shard_id = ShardId::new();
1126 let client = new_test_client(&dyncfgs).await;
1127 let mut write1 = client
1128 .open_writer::<String, String, u64, i64>(
1129 shard_id,
1130 Arc::new(StringSchema),
1131 Arc::new(StringSchema),
1132 Diagnostics::for_tests(),
1133 )
1134 .await
1135 .expect("codec mismatch");
1136 let mut read1 = client
1137 .open_leased_reader::<String, String, u64, i64>(
1138 shard_id,
1139 Arc::new(StringSchema),
1140 Arc::new(StringSchema),
1141 Diagnostics::for_tests(),
1142 true,
1143 )
1144 .await
1145 .expect("codec mismatch");
1146 let mut read2 = client
1147 .open_leased_reader::<String, String, u64, i64>(
1148 shard_id,
1149 Arc::new(StringSchema),
1150 Arc::new(StringSchema),
1151 Diagnostics::for_tests(),
1152 true,
1153 )
1154 .await
1155 .expect("codec mismatch");
1156 let mut write2 = client
1157 .open_writer::<String, String, u64, i64>(
1158 shard_id,
1159 Arc::new(StringSchema),
1160 Arc::new(StringSchema),
1161 Diagnostics::for_tests(),
1162 )
1163 .await
1164 .expect("codec mismatch");
1165
1166 write2.expect_compare_and_append(&data[..1], 0, 2).await;
1167 assert_eq!(
1168 read2.expect_snapshot_and_fetch(1).await,
1169 all_ok(&data[..1], 1)
1170 );
1171 write1.expect_compare_and_append(&data[1..], 2, 4).await;
1172 assert_eq!(read1.expect_snapshot_and_fetch(3).await, all_ok(&data, 3));
1173 }
1174
1175 #[mz_persist_proc::test(tokio::test)]
1176 #[cfg_attr(miri, ignore)] async fn invalid_usage(dyncfgs: ConfigUpdates) {
1178 let data = vec![
1179 (("1".to_owned(), "one".to_owned()), 1, 1),
1180 (("2".to_owned(), "two".to_owned()), 2, 1),
1181 (("3".to_owned(), "three".to_owned()), 3, 1),
1182 ];
1183
1184 let shard_id0 = "s00000000-0000-0000-0000-000000000000"
1185 .parse::<ShardId>()
1186 .expect("invalid shard id");
1187 let mut client = new_test_client(&dyncfgs).await;
1188
1189 let (mut write0, mut read0) = client
1190 .expect_open::<String, String, u64, i64>(shard_id0)
1191 .await;
1192
1193 write0.expect_compare_and_append(&data, 0, 4).await;
1194
1195 {
1197 fn codecs(
1198 k: &str,
1199 v: &str,
1200 t: &str,
1201 d: &str,
1202 ) -> (String, String, String, String, Option<CodecConcreteType>) {
1203 (k.to_owned(), v.to_owned(), t.to_owned(), d.to_owned(), None)
1204 }
1205
1206 client.shared_states = Arc::new(StateCache::new_no_metrics());
1207 assert_eq!(
1208 client
1209 .open::<Vec<u8>, String, u64, i64>(
1210 shard_id0,
1211 Arc::new(VecU8Schema),
1212 Arc::new(StringSchema),
1213 Diagnostics::for_tests(),
1214 true,
1215 )
1216 .await
1217 .unwrap_err(),
1218 InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1219 requested: codecs("Vec<u8>", "String", "u64", "i64"),
1220 actual: codecs("String", "String", "u64", "i64"),
1221 }))
1222 );
1223 assert_eq!(
1224 client
1225 .open::<String, Vec<u8>, u64, i64>(
1226 shard_id0,
1227 Arc::new(StringSchema),
1228 Arc::new(VecU8Schema),
1229 Diagnostics::for_tests(),
1230 true,
1231 )
1232 .await
1233 .unwrap_err(),
1234 InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1235 requested: codecs("String", "Vec<u8>", "u64", "i64"),
1236 actual: codecs("String", "String", "u64", "i64"),
1237 }))
1238 );
1239 assert_eq!(
1240 client
1241 .open::<String, String, i64, i64>(
1242 shard_id0,
1243 Arc::new(StringSchema),
1244 Arc::new(StringSchema),
1245 Diagnostics::for_tests(),
1246 true,
1247 )
1248 .await
1249 .unwrap_err(),
1250 InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1251 requested: codecs("String", "String", "i64", "i64"),
1252 actual: codecs("String", "String", "u64", "i64"),
1253 }))
1254 );
1255 assert_eq!(
1256 client
1257 .open::<String, String, u64, u64>(
1258 shard_id0,
1259 Arc::new(StringSchema),
1260 Arc::new(StringSchema),
1261 Diagnostics::for_tests(),
1262 true,
1263 )
1264 .await
1265 .unwrap_err(),
1266 InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1267 requested: codecs("String", "String", "u64", "u64"),
1268 actual: codecs("String", "String", "u64", "i64"),
1269 }))
1270 );
1271
1272 assert_eq!(
1276 client
1277 .open_leased_reader::<Vec<u8>, String, u64, i64>(
1278 shard_id0,
1279 Arc::new(VecU8Schema),
1280 Arc::new(StringSchema),
1281 Diagnostics::for_tests(),
1282 true,
1283 )
1284 .await
1285 .unwrap_err(),
1286 InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1287 requested: codecs("Vec<u8>", "String", "u64", "i64"),
1288 actual: codecs("String", "String", "u64", "i64"),
1289 }))
1290 );
1291 assert_eq!(
1292 client
1293 .open_writer::<Vec<u8>, String, u64, i64>(
1294 shard_id0,
1295 Arc::new(VecU8Schema),
1296 Arc::new(StringSchema),
1297 Diagnostics::for_tests(),
1298 )
1299 .await
1300 .unwrap_err(),
1301 InvalidUsage::CodecMismatch(Box::new(CodecMismatch {
1302 requested: codecs("Vec<u8>", "String", "u64", "i64"),
1303 actual: codecs("String", "String", "u64", "i64"),
1304 }))
1305 );
1306 }
1307
1308 {
1310 let snap = read0
1311 .snapshot(Antichain::from_elem(3))
1312 .await
1313 .expect("cannot serve requested as_of");
1314
1315 let shard_id1 = "s11111111-1111-1111-1111-111111111111"
1316 .parse::<ShardId>()
1317 .expect("invalid shard id");
1318 let mut fetcher1 = client
1319 .create_batch_fetcher::<String, String, u64, i64>(
1320 shard_id1,
1321 Default::default(),
1322 Default::default(),
1323 false,
1324 Diagnostics::for_tests(),
1325 )
1326 .await
1327 .unwrap();
1328 for part in snap {
1329 let (part, _lease) = part.into_exchangeable_part();
1330 let res = fetcher1.fetch_leased_part(part).await;
1331 assert_eq!(
1332 res.unwrap_err(),
1333 InvalidUsage::BatchNotFromThisShard {
1334 batch_shard: shard_id0,
1335 handle_shard: shard_id1,
1336 }
1337 );
1338 }
1339 }
1340
1341 {
1343 let ts3 = &data[2];
1344 assert_eq!(ts3.1, 3);
1345 let ts3 = vec![ts3.clone()];
1346
1347 assert_eq!(
1350 write0
1351 .append(&ts3, Antichain::from_elem(4), Antichain::from_elem(5))
1352 .await
1353 .unwrap_err(),
1354 InvalidUsage::UpdateNotBeyondLower {
1355 ts: 3,
1356 lower: Antichain::from_elem(4),
1357 },
1358 );
1359 assert_eq!(
1360 write0
1361 .append(&ts3, Antichain::from_elem(2), Antichain::from_elem(3))
1362 .await
1363 .unwrap_err(),
1364 InvalidUsage::UpdateBeyondUpper {
1365 ts: 3,
1366 expected_upper: Antichain::from_elem(3),
1367 },
1368 );
1369 assert_eq!(
1371 write0
1372 .append(&data[..0], Antichain::from_elem(3), Antichain::from_elem(2))
1373 .await
1374 .unwrap_err(),
1375 InvalidUsage::InvalidBounds {
1376 lower: Antichain::from_elem(3),
1377 upper: Antichain::from_elem(2),
1378 },
1379 );
1380
1381 assert_eq!(
1383 write0
1384 .builder(Antichain::from_elem(3))
1385 .finish(Antichain::from_elem(2))
1386 .await
1387 .unwrap_err(),
1388 InvalidUsage::InvalidBounds {
1389 lower: Antichain::from_elem(3),
1390 upper: Antichain::from_elem(2)
1391 },
1392 );
1393 let batch = write0
1394 .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1395 .await
1396 .expect("invalid usage");
1397 assert_eq!(
1398 write0
1399 .append_batch(batch, Antichain::from_elem(4), Antichain::from_elem(5))
1400 .await
1401 .unwrap_err(),
1402 InvalidUsage::InvalidBatchBounds {
1403 batch_lower: Antichain::from_elem(3),
1404 batch_upper: Antichain::from_elem(4),
1405 append_lower: Antichain::from_elem(4),
1406 append_upper: Antichain::from_elem(5),
1407 },
1408 );
1409 let batch = write0
1410 .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1411 .await
1412 .expect("invalid usage");
1413 assert_eq!(
1414 write0
1415 .append_batch(batch, Antichain::from_elem(2), Antichain::from_elem(3))
1416 .await
1417 .unwrap_err(),
1418 InvalidUsage::InvalidBatchBounds {
1419 batch_lower: Antichain::from_elem(3),
1420 batch_upper: Antichain::from_elem(4),
1421 append_lower: Antichain::from_elem(2),
1422 append_upper: Antichain::from_elem(3),
1423 },
1424 );
1425 let batch = write0
1426 .batch(&ts3, Antichain::from_elem(3), Antichain::from_elem(4))
1427 .await
1428 .expect("invalid usage");
1429 assert!(matches!(
1432 write0
1433 .append_batch(batch, Antichain::from_elem(3), Antichain::from_elem(3))
1434 .await
1435 .unwrap_err(),
1436 InvalidUsage::InvalidEmptyTimeInterval { .. }
1437 ));
1438 }
1439 }
1440
1441 #[mz_persist_proc::test(tokio::test)]
1442 #[cfg_attr(miri, ignore)] async fn multiple_shards(dyncfgs: ConfigUpdates) {
1444 let data1 = [
1445 (("1".to_owned(), "one".to_owned()), 1, 1),
1446 (("2".to_owned(), "two".to_owned()), 2, 1),
1447 ];
1448
1449 let data2 = [(("1".to_owned(), ()), 1, 1), (("2".to_owned(), ()), 2, 1)];
1450
1451 let client = new_test_client(&dyncfgs).await;
1452
1453 let (mut write1, mut read1) = client
1454 .expect_open::<String, String, u64, i64>(ShardId::new())
1455 .await;
1456
1457 let (mut write2, mut read2) = client
1460 .expect_open::<String, (), u64, i64>(ShardId::new())
1461 .await;
1462
1463 write1
1464 .expect_compare_and_append(&data1[..], u64::minimum(), 3)
1465 .await;
1466
1467 write2
1468 .expect_compare_and_append(&data2[..], u64::minimum(), 3)
1469 .await;
1470
1471 assert_eq!(
1472 read1.expect_snapshot_and_fetch(2).await,
1473 all_ok(&data1[..], 2)
1474 );
1475
1476 assert_eq!(
1477 read2.expect_snapshot_and_fetch(2).await,
1478 all_ok(&data2[..], 2)
1479 );
1480 }
1481
1482 #[mz_persist_proc::test(tokio::test)]
1483 #[cfg_attr(miri, ignore)] async fn fetch_upper(dyncfgs: ConfigUpdates) {
1485 let data = [
1486 (("1".to_owned(), "one".to_owned()), 1, 1),
1487 (("2".to_owned(), "two".to_owned()), 2, 1),
1488 ];
1489
1490 let client = new_test_client(&dyncfgs).await;
1491
1492 let shard_id = ShardId::new();
1493
1494 let (mut write1, _read1) = client
1495 .expect_open::<String, String, u64, i64>(shard_id)
1496 .await;
1497
1498 let (mut write2, _read2) = client
1499 .expect_open::<String, String, u64, i64>(shard_id)
1500 .await;
1501
1502 write1
1503 .expect_append(&data[..], write1.upper().clone(), vec![3])
1504 .await;
1505
1506 assert_eq!(write2.fetch_recent_upper().await, &Antichain::from_elem(3));
1508
1509 assert_eq!(write2.upper(), &Antichain::from_elem(3));
1512 }
1513
1514 #[mz_persist_proc::test(tokio::test)]
1515 #[cfg_attr(miri, ignore)] async fn append_with_invalid_upper(dyncfgs: ConfigUpdates) {
1517 let data = [
1518 (("1".to_owned(), "one".to_owned()), 1, 1),
1519 (("2".to_owned(), "two".to_owned()), 2, 1),
1520 ];
1521
1522 let client = new_test_client(&dyncfgs).await;
1523
1524 let shard_id = ShardId::new();
1525
1526 let (mut write, _read) = client
1527 .expect_open::<String, String, u64, i64>(shard_id)
1528 .await;
1529
1530 write
1531 .expect_append(&data[..], write.upper().clone(), vec![3])
1532 .await;
1533
1534 let data = [
1535 (("5".to_owned(), "fünf".to_owned()), 5, 1),
1536 (("6".to_owned(), "sechs".to_owned()), 6, 1),
1537 ];
1538 let res = write
1539 .append(
1540 data.iter(),
1541 Antichain::from_elem(5),
1542 Antichain::from_elem(7),
1543 )
1544 .await;
1545 assert_eq!(
1546 res,
1547 Ok(Err(UpperMismatch {
1548 expected: Antichain::from_elem(5),
1549 current: Antichain::from_elem(3)
1550 }))
1551 );
1552
1553 assert_eq!(write.upper(), &Antichain::from_elem(3));
1555 }
1556
1557 #[allow(unused)]
1560 async fn sync_send(dyncfgs: ConfigUpdates) {
1561 mz_ore::test::init_logging();
1562
1563 fn is_send_sync<T: Send + Sync>(_x: T) -> bool {
1564 true
1565 }
1566
1567 let client = new_test_client(&dyncfgs).await;
1568
1569 let (write, read) = client
1570 .expect_open::<String, String, u64, i64>(ShardId::new())
1571 .await;
1572
1573 assert!(is_send_sync(client));
1574 assert!(is_send_sync(write));
1575 assert!(is_send_sync(read));
1576 }
1577
1578 #[mz_persist_proc::test(tokio::test)]
1579 #[cfg_attr(miri, ignore)] async fn compare_and_append(dyncfgs: ConfigUpdates) {
1581 let data = vec![
1582 (("1".to_owned(), "one".to_owned()), 1, 1),
1583 (("2".to_owned(), "two".to_owned()), 2, 1),
1584 (("3".to_owned(), "three".to_owned()), 3, 1),
1585 ];
1586
1587 let id = ShardId::new();
1588 let client = new_test_client(&dyncfgs).await;
1589 let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1590
1591 let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1592
1593 assert_eq!(write1.upper(), &Antichain::from_elem(u64::minimum()));
1594 assert_eq!(write2.upper(), &Antichain::from_elem(u64::minimum()));
1595 assert_eq!(read.since(), &Antichain::from_elem(u64::minimum()));
1596
1597 write1
1599 .expect_compare_and_append(&data[..2], u64::minimum(), 3)
1600 .await;
1601 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1602
1603 assert_eq!(
1604 read.expect_snapshot_and_fetch(2).await,
1605 all_ok(&data[..2], 2)
1606 );
1607
1608 let res = write2
1610 .compare_and_append(
1611 &data[..2],
1612 Antichain::from_elem(u64::minimum()),
1613 Antichain::from_elem(3),
1614 )
1615 .await;
1616 assert_eq!(
1617 res,
1618 Ok(Err(UpperMismatch {
1619 expected: Antichain::from_elem(u64::minimum()),
1620 current: Antichain::from_elem(3)
1621 }))
1622 );
1623
1624 assert_eq!(write2.upper(), &Antichain::from_elem(3));
1626
1627 write2.expect_compare_and_append(&data[2..], 3, 4).await;
1629
1630 assert_eq!(write2.upper(), &Antichain::from_elem(4));
1631
1632 assert_eq!(read.expect_snapshot_and_fetch(3).await, all_ok(&data, 3));
1633 }
1634
1635 #[mz_persist_proc::test(tokio::test)]
1636 #[cfg_attr(miri, ignore)] async fn overlapping_append(dyncfgs: ConfigUpdates) {
1638 mz_ore::test::init_logging_default("info");
1639
1640 let data = vec![
1641 (("1".to_owned(), "one".to_owned()), 1, 1),
1642 (("2".to_owned(), "two".to_owned()), 2, 1),
1643 (("3".to_owned(), "three".to_owned()), 3, 1),
1644 (("4".to_owned(), "vier".to_owned()), 4, 1),
1645 (("5".to_owned(), "cinque".to_owned()), 5, 1),
1646 ];
1647
1648 let id = ShardId::new();
1649 let client = new_test_client(&dyncfgs).await;
1650
1651 let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1652
1653 let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1654
1655 let mut listen = read.clone("").await.expect_listen(0).await;
1657
1658 write1
1660 .expect_append(&data[..2], write1.upper().clone(), vec![3])
1661 .await;
1662 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1663
1664 write2
1666 .expect_append(&data[..4], write2.upper().clone(), vec![5])
1667 .await;
1668 assert_eq!(write2.upper(), &Antichain::from_elem(5));
1669
1670 write1
1672 .expect_append(&data[2..5], write1.upper().clone(), vec![6])
1673 .await;
1674 assert_eq!(write1.upper(), &Antichain::from_elem(6));
1675
1676 assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1677
1678 assert_eq!(
1679 listen.read_until(&6).await,
1680 (all_ok(&data[..], 1), Antichain::from_elem(6))
1681 );
1682 }
1683
1684 #[mz_persist_proc::test(tokio::test)]
1687 #[cfg_attr(miri, ignore)] async fn contiguous_append(dyncfgs: ConfigUpdates) {
1689 let data = vec![
1690 (("1".to_owned(), "one".to_owned()), 1, 1),
1691 (("2".to_owned(), "two".to_owned()), 2, 1),
1692 (("3".to_owned(), "three".to_owned()), 3, 1),
1693 (("4".to_owned(), "vier".to_owned()), 4, 1),
1694 (("5".to_owned(), "cinque".to_owned()), 5, 1),
1695 ];
1696
1697 let id = ShardId::new();
1698 let client = new_test_client(&dyncfgs).await;
1699
1700 let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1701
1702 write
1704 .expect_append(&data[..2], write.upper().clone(), vec![3])
1705 .await;
1706 assert_eq!(write.upper(), &Antichain::from_elem(3));
1707
1708 let result = write
1711 .append(
1712 &data[4..5],
1713 Antichain::from_elem(5),
1714 Antichain::from_elem(6),
1715 )
1716 .await;
1717 assert_eq!(
1718 result,
1719 Ok(Err(UpperMismatch {
1720 expected: Antichain::from_elem(5),
1721 current: Antichain::from_elem(3)
1722 }))
1723 );
1724
1725 write.expect_append(&data[2..5], vec![3], vec![6]).await;
1727 assert_eq!(write.upper(), &Antichain::from_elem(6));
1728
1729 assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1730 }
1731
1732 #[mz_persist_proc::test(tokio::test)]
1735 #[cfg_attr(miri, ignore)] async fn noncontiguous_append_per_writer(dyncfgs: ConfigUpdates) {
1737 let data = vec![
1738 (("1".to_owned(), "one".to_owned()), 1, 1),
1739 (("2".to_owned(), "two".to_owned()), 2, 1),
1740 (("3".to_owned(), "three".to_owned()), 3, 1),
1741 (("4".to_owned(), "vier".to_owned()), 4, 1),
1742 (("5".to_owned(), "cinque".to_owned()), 5, 1),
1743 ];
1744
1745 let id = ShardId::new();
1746 let client = new_test_client(&dyncfgs).await;
1747
1748 let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1749
1750 let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1751
1752 write1
1754 .expect_append(&data[..2], write1.upper().clone(), vec![3])
1755 .await;
1756 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1757
1758 write2.upper = Antichain::from_elem(3);
1760 write2
1761 .expect_append(&data[2..4], write2.upper().clone(), vec![5])
1762 .await;
1763 assert_eq!(write2.upper(), &Antichain::from_elem(5));
1764
1765 write1.upper = Antichain::from_elem(5);
1767 write1
1768 .expect_append(&data[4..5], write1.upper().clone(), vec![6])
1769 .await;
1770 assert_eq!(write1.upper(), &Antichain::from_elem(6));
1771
1772 assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1773 }
1774
1775 #[mz_persist_proc::test(tokio::test)]
1778 #[cfg_attr(miri, ignore)] async fn contiguous_compare_and_append(dyncfgs: ConfigUpdates) {
1780 let data = vec![
1781 (("1".to_owned(), "one".to_owned()), 1, 1),
1782 (("2".to_owned(), "two".to_owned()), 2, 1),
1783 (("3".to_owned(), "three".to_owned()), 3, 1),
1784 (("4".to_owned(), "vier".to_owned()), 4, 1),
1785 (("5".to_owned(), "cinque".to_owned()), 5, 1),
1786 ];
1787
1788 let id = ShardId::new();
1789 let client = new_test_client(&dyncfgs).await;
1790
1791 let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1792
1793 write.expect_compare_and_append(&data[..2], 0, 3).await;
1795 assert_eq!(write.upper(), &Antichain::from_elem(3));
1796
1797 let result = write
1800 .compare_and_append(
1801 &data[4..5],
1802 Antichain::from_elem(5),
1803 Antichain::from_elem(6),
1804 )
1805 .await;
1806 assert_eq!(
1807 result,
1808 Ok(Err(UpperMismatch {
1809 expected: Antichain::from_elem(5),
1810 current: Antichain::from_elem(3)
1811 }))
1812 );
1813
1814 write.expect_compare_and_append(&data[2..5], 3, 6).await;
1817 assert_eq!(write.upper(), &Antichain::from_elem(6));
1818
1819 assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1820 }
1821
1822 #[mz_persist_proc::test(tokio::test)]
1825 #[cfg_attr(miri, ignore)] async fn noncontiguous_compare_and_append_per_writer(dyncfgs: ConfigUpdates) {
1827 let data = vec![
1828 (("1".to_owned(), "one".to_owned()), 1, 1),
1829 (("2".to_owned(), "two".to_owned()), 2, 1),
1830 (("3".to_owned(), "three".to_owned()), 3, 1),
1831 (("4".to_owned(), "vier".to_owned()), 4, 1),
1832 (("5".to_owned(), "cinque".to_owned()), 5, 1),
1833 ];
1834
1835 let id = ShardId::new();
1836 let client = new_test_client(&dyncfgs).await;
1837
1838 let (mut write1, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1839
1840 let (mut write2, _read) = client.expect_open::<String, String, u64, i64>(id).await;
1841
1842 write1.expect_compare_and_append(&data[..2], 0, 3).await;
1844 assert_eq!(write1.upper(), &Antichain::from_elem(3));
1845
1846 write2.expect_compare_and_append(&data[2..4], 3, 5).await;
1848 assert_eq!(write2.upper(), &Antichain::from_elem(5));
1849
1850 write1.expect_compare_and_append(&data[4..5], 5, 6).await;
1852 assert_eq!(write1.upper(), &Antichain::from_elem(6));
1853
1854 assert_eq!(read.expect_snapshot_and_fetch(5).await, all_ok(&data, 5));
1855 }
1856
1857 #[mz_ore::test]
1858 fn fmt_ids() {
1859 assert_eq!(
1860 format!("{}", LeasedReaderId([0u8; 16])),
1861 "r00000000-0000-0000-0000-000000000000"
1862 );
1863 assert_eq!(
1864 format!("{:?}", LeasedReaderId([0u8; 16])),
1865 "LeasedReaderId(00000000-0000-0000-0000-000000000000)"
1866 );
1867 }
1868
1869 #[mz_persist_proc::test(tokio::test(flavor = "multi_thread"))]
1870 #[cfg_attr(miri, ignore)] async fn concurrency(dyncfgs: ConfigUpdates) {
1872 let data = DataGenerator::small();
1873
1874 const NUM_WRITERS: usize = 2;
1875 let id = ShardId::new();
1876 let client = new_test_client(&dyncfgs).await;
1877 let mut handles = Vec::<mz_ore::task::JoinHandle<()>>::new();
1878 for idx in 0..NUM_WRITERS {
1879 let (data, client) = (data.clone(), client.clone());
1880
1881 let (batch_tx, mut batch_rx) = tokio::sync::mpsc::channel(1);
1882
1883 let client1 = client.clone();
1884 let handle = mz_ore::task::spawn(|| format!("writer-{}", idx), async move {
1885 let (write, _) = client1.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1886 let mut current_upper = 0;
1887 for batch in data.batches() {
1888 let new_upper = match batch.get(batch.len() - 1) {
1889 Some((_, max_ts, _)) => u64::decode(max_ts) + 1,
1890 None => continue,
1891 };
1892 if PartialOrder::less_equal(&Antichain::from_elem(new_upper), write.upper()) {
1907 continue;
1908 }
1909
1910 let current_upper_chain = Antichain::from_elem(current_upper);
1911 current_upper = new_upper;
1912 let new_upper_chain = Antichain::from_elem(new_upper);
1913 let mut builder = write.builder(current_upper_chain);
1914
1915 for ((k, v), t, d) in batch.iter() {
1916 builder
1917 .add(&k.to_vec(), &v.to_vec(), &u64::decode(t), &i64::decode(d))
1918 .await
1919 .expect("invalid usage");
1920 }
1921
1922 let batch = builder
1923 .finish(new_upper_chain)
1924 .await
1925 .expect("invalid usage");
1926
1927 match batch_tx.send(batch).await {
1928 Ok(_) => (),
1929 Err(e) => panic!("send error: {}", e),
1930 }
1931 }
1932 });
1933 handles.push(handle);
1934
1935 let handle = mz_ore::task::spawn(|| format!("appender-{}", idx), async move {
1936 let (mut write, _) = client.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1937
1938 while let Some(batch) = batch_rx.recv().await {
1939 let lower = batch.lower().clone();
1940 let upper = batch.upper().clone();
1941 write
1942 .append_batch(batch, lower, upper)
1943 .await
1944 .expect("invalid usage")
1945 .expect("unexpected upper");
1946 }
1947 });
1948 handles.push(handle);
1949 }
1950
1951 for handle in handles {
1952 let () = handle.await;
1953 }
1954
1955 let expected = data.records().collect::<Vec<_>>();
1956 let max_ts = expected.last().map(|(_, t, _)| *t).unwrap_or_default();
1957 let (_, mut read) = client.expect_open::<Vec<u8>, Vec<u8>, u64, i64>(id).await;
1958 assert_eq!(
1959 read.expect_snapshot_and_fetch(max_ts).await,
1960 all_ok(expected.iter(), max_ts)
1961 );
1962 }
1963
1964 #[mz_persist_proc::test(tokio::test)]
1968 #[cfg_attr(miri, ignore)] async fn regression_blocking_reads(dyncfgs: ConfigUpdates) {
1970 let waker = noop_waker();
1971 let mut cx = Context::from_waker(&waker);
1972
1973 let data = [
1974 (("1".to_owned(), "one".to_owned()), 1, 1),
1975 (("2".to_owned(), "two".to_owned()), 2, 1),
1976 (("3".to_owned(), "three".to_owned()), 3, 1),
1977 ];
1978
1979 let id = ShardId::new();
1980 let client = new_test_client(&dyncfgs).await;
1981 let (mut write, mut read) = client.expect_open::<String, String, u64, i64>(id).await;
1982
1983 let mut listen = read.clone("").await.expect_listen(1).await;
1985 let mut listen_next = Box::pin(listen.fetch_next());
1986 for _ in 0..100 {
1990 assert!(
1991 Pin::new(&mut listen_next).poll(&mut cx).is_pending(),
1992 "listen::next unexpectedly ready"
1993 );
1994 }
1995
1996 write
1998 .expect_compare_and_append(&data[..2], u64::minimum(), 3)
1999 .await;
2000
2001 assert_eq!(
2004 listen_next.await,
2005 vec![
2006 ListenEvent::Updates(vec![(("2".to_owned(), "two".to_owned()), 2, 1)]),
2007 ListenEvent::Progress(Antichain::from_elem(3)),
2008 ]
2009 );
2010
2011 let mut snap = Box::pin(read.expect_snapshot_and_fetch(3));
2025 for _ in 0..100 {
2026 assert!(
2027 Pin::new(&mut snap).poll(&mut cx).is_pending(),
2028 "snapshot unexpectedly ready"
2029 );
2030 }
2031
2032 write.expect_compare_and_append(&data[2..], 3, 4).await;
2034
2035 assert_eq!(snap.await, all_ok(&data[..], 3));
2037 }
2038
2039 #[mz_persist_proc::test(tokio::test)]
2040 #[cfg_attr(miri, ignore)] async fn heartbeat_task_shutdown(dyncfgs: ConfigUpdates) {
2042 let mut cache = new_test_client_cache(&dyncfgs);
2045 cache
2046 .cfg
2047 .set_config(&READER_LEASE_DURATION, Duration::from_millis(1));
2048 cache.cfg.writer_lease_duration = Duration::from_millis(1);
2049 let (_write, mut read) = cache
2050 .open(PersistLocation::new_in_mem())
2051 .await
2052 .expect("client construction failed")
2053 .expect_open::<(), (), u64, i64>(ShardId::new())
2054 .await;
2055 let read_unexpired_state = read
2056 .unexpired_state
2057 .take()
2058 .expect("handle should have unexpired state");
2059 read.expire().await;
2060 read_unexpired_state.heartbeat_task.await
2061 }
2062
2063 #[mz_persist_proc::test(tokio::test)]
2066 #[cfg_attr(miri, ignore)] async fn finalize_empty_shard(dyncfgs: ConfigUpdates) {
2068 let persist_client = new_test_client(&dyncfgs).await;
2069
2070 let shard_id = ShardId::new();
2071 pub const CRITICAL_SINCE: CriticalReaderId =
2072 CriticalReaderId([0, 0, 0, 0, 17, 17, 34, 34, 51, 51, 68, 68, 68, 68, 68, 68]);
2073
2074 let (mut write, mut read) = persist_client
2075 .expect_open::<(), (), u64, i64>(shard_id)
2076 .await;
2077
2078 let () = read.downgrade_since(&Antichain::new()).await;
2081 let () = write.advance_upper(&Antichain::new()).await;
2082
2083 let mut since_handle: SinceHandle<(), (), u64, i64> = persist_client
2084 .open_critical_since(
2085 shard_id,
2086 CRITICAL_SINCE,
2087 Opaque::encode(&0u64),
2088 Diagnostics::for_tests(),
2089 )
2090 .await
2091 .expect("invalid persist usage");
2092
2093 let epoch = since_handle.opaque().clone();
2094 let new_since = Antichain::new();
2095 let downgrade = since_handle
2096 .compare_and_downgrade_since(&epoch, (&epoch, &new_since))
2097 .await;
2098
2099 assert!(
2100 downgrade.is_ok(),
2101 "downgrade of critical handle must succeed"
2102 );
2103
2104 let finalize = persist_client
2105 .finalize_shard::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2106 .await;
2107
2108 assert_ok!(finalize, "finalization must succeed");
2109
2110 let is_finalized = persist_client
2111 .is_finalized::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2112 .await
2113 .expect("invalid persist usage");
2114 assert!(is_finalized, "shard must still be finalized");
2115 }
2116
2117 #[mz_persist_proc::test(tokio::test)]
2121 #[cfg_attr(miri, ignore)] async fn finalize_shard(dyncfgs: ConfigUpdates) {
2123 const DATA: &[(((), ()), u64, i64)] = &[(((), ()), 0, 1)];
2124 let persist_client = new_test_client(&dyncfgs).await;
2125
2126 let shard_id = ShardId::new();
2127 pub const CRITICAL_SINCE: CriticalReaderId =
2128 CriticalReaderId([0, 0, 0, 0, 17, 17, 34, 34, 51, 51, 68, 68, 68, 68, 68, 68]);
2129
2130 let (mut write, mut read) = persist_client
2131 .expect_open::<(), (), u64, i64>(shard_id)
2132 .await;
2133
2134 let () = write
2136 .compare_and_append(DATA, Antichain::from_elem(0), Antichain::from_elem(1))
2137 .await
2138 .expect("usage should be valid")
2139 .expect("upper should match");
2140
2141 let () = read.downgrade_since(&Antichain::new()).await;
2144 let () = write.advance_upper(&Antichain::new()).await;
2145
2146 let mut since_handle: SinceHandle<(), (), u64, i64> = persist_client
2147 .open_critical_since(
2148 shard_id,
2149 CRITICAL_SINCE,
2150 Opaque::encode(&0u64),
2151 Diagnostics::for_tests(),
2152 )
2153 .await
2154 .expect("invalid persist usage");
2155
2156 let epoch = since_handle.opaque().clone();
2157 let new_since = Antichain::new();
2158 let downgrade = since_handle
2159 .compare_and_downgrade_since(&epoch, (&epoch, &new_since))
2160 .await;
2161
2162 assert!(
2163 downgrade.is_ok(),
2164 "downgrade of critical handle must succeed"
2165 );
2166
2167 let finalize = persist_client
2168 .finalize_shard::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2169 .await;
2170
2171 assert_ok!(finalize, "finalization must succeed");
2172
2173 let is_finalized = persist_client
2174 .is_finalized::<(), (), u64, i64>(shard_id, Diagnostics::for_tests())
2175 .await
2176 .expect("invalid persist usage");
2177 assert!(is_finalized, "shard must still be finalized");
2178 }
2179
2180 proptest! {
2181 #![proptest_config(ProptestConfig::with_cases(4096))]
2182
2183 #[mz_ore::test]
2184 #[cfg_attr(miri, ignore)] fn shard_id_protobuf_roundtrip(expect in any::<ShardId>() ) {
2186 let actual = protobuf_roundtrip::<_, String>(&expect);
2187 assert_ok!(actual);
2188 assert_eq!(actual.unwrap(), expect);
2189 }
2190 }
2191}