1use std::borrow::Borrow;
13use std::fmt::Debug;
14use std::sync::Arc;
15
16use differential_dataflow::difference::Monoid;
17use differential_dataflow::lattice::Lattice;
18use differential_dataflow::trace::Description;
19use futures::StreamExt;
20use futures::stream::FuturesUnordered;
21use mz_dyncfg::{Config, ParameterScope};
22use mz_ore::task::RuntimeExt;
23use mz_ore::{instrument, soft_panic_or_log};
24use mz_persist::location::Blob;
25use mz_persist_types::schema::SchemaId;
26use mz_persist_types::{Codec, Codec64};
27use mz_proto::{IntoRustIfSome, ProtoType};
28use proptest_derive::Arbitrary;
29use semver::Version;
30use serde::{Deserialize, Serialize};
31use timely::PartialOrder;
32use timely::order::TotalOrder;
33use timely::progress::{Antichain, Timestamp};
34use tokio::runtime::Handle;
35use tracing::{Instrument, debug_span, error, info, warn};
36use uuid::Uuid;
37
38use crate::batch::{
39 Added, BATCH_DELETE_ENABLED, Batch, BatchBuilder, BatchBuilderConfig, BatchBuilderInternal,
40 BatchParts, ProtoBatch, validate_truncate_batch,
41};
42use crate::error::{InvalidUsage, UpperMismatch};
43use crate::fetch::{
44 EncodedPart, FetchBatchFilter, FetchedPart, PartDecodeFormat, VALIDATE_PART_BOUNDS_ON_READ,
45};
46use crate::internal::compact::{CompactConfig, Compactor};
47use crate::internal::encoding::{Schemas, assert_code_can_read_data};
48use crate::internal::machine::{
49 CompareAndAppendRes, ExpireFn, Machine, next_listen_batch_retry_params,
50};
51use crate::internal::metrics::{BatchWriteMetrics, Metrics, ShardMetrics};
52use crate::internal::state::{BatchPart, HandleDebugState, HollowBatch, RunOrder, RunPart};
53use crate::read::ReadHandle;
54use crate::schema::PartMigration;
55use crate::{GarbageCollector, IsolatedRuntime, PersistConfig, ShardId, parse_id};
56
57pub(crate) const COMBINE_INLINE_WRITES: Config<bool> = Config::new(
58 "persist_write_combine_inline_writes",
59 true,
60 "If set, re-encode inline writes if they don't fit into the batch metadata limits.",
61 ParameterScope::Environment,
62);
63
64pub(crate) const VALIDATE_PART_BOUNDS_ON_WRITE: Config<bool> = Config::new(
65 "persist_validate_part_bounds_on_write",
66 false,
67 "Validate the part lower <= the batch lower and the part upper <= batch upper,\
68 for the batch being appended.",
69 ParameterScope::Environment,
70);
71
72#[derive(
74 Arbitrary,
75 Clone,
76 PartialEq,
77 Eq,
78 PartialOrd,
79 Ord,
80 Hash,
81 Serialize,
82 Deserialize
83)]
84#[serde(try_from = "String", into = "String")]
85pub struct WriterId(pub(crate) [u8; 16]);
86
87impl std::fmt::Display for WriterId {
88 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89 write!(f, "w{}", Uuid::from_bytes(self.0))
90 }
91}
92
93impl std::fmt::Debug for WriterId {
94 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95 write!(f, "WriterId({})", Uuid::from_bytes(self.0))
96 }
97}
98
99impl std::str::FromStr for WriterId {
100 type Err = String;
101
102 fn from_str(s: &str) -> Result<Self, Self::Err> {
103 parse_id("w", "WriterId", s).map(WriterId)
104 }
105}
106
107impl From<WriterId> for String {
108 fn from(writer_id: WriterId) -> Self {
109 writer_id.to_string()
110 }
111}
112
113impl TryFrom<String> for WriterId {
114 type Error = String;
115
116 fn try_from(s: String) -> Result<Self, Self::Error> {
117 s.parse()
118 }
119}
120
121impl WriterId {
122 pub(crate) fn new() -> Self {
123 WriterId(*Uuid::new_v4().as_bytes())
124 }
125}
126
127#[derive(Debug)]
143pub struct WriteHandle<K: Codec, V: Codec, T, D> {
144 pub(crate) cfg: PersistConfig,
145 pub(crate) metrics: Arc<Metrics>,
146 pub(crate) machine: Machine<K, V, T, D>,
147 pub(crate) gc: GarbageCollector<K, V, T, D>,
148 pub(crate) compact: Option<Compactor<K, V, T, D>>,
149 pub(crate) blob: Arc<dyn Blob>,
150 pub(crate) isolated_runtime: Arc<IsolatedRuntime>,
151 pub(crate) writer_id: WriterId,
152 pub(crate) debug_state: HandleDebugState,
153 pub(crate) write_schemas: Schemas<K, V>,
154
155 pub(crate) upper: Antichain<T>,
156 expire_fn: Option<ExpireFn>,
157}
158
159impl<K, V, T, D> WriteHandle<K, V, T, D>
160where
161 K: Debug + Codec,
162 V: Debug + Codec,
163 T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
164 D: Monoid + Ord + Codec64 + Send + Sync,
165{
166 pub(crate) fn new(
167 cfg: PersistConfig,
168 metrics: Arc<Metrics>,
169 machine: Machine<K, V, T, D>,
170 gc: GarbageCollector<K, V, T, D>,
171 blob: Arc<dyn Blob>,
172 writer_id: WriterId,
173 purpose: &str,
174 write_schemas: Schemas<K, V>,
175 ) -> Self {
176 let isolated_runtime = Arc::clone(&machine.isolated_runtime);
177 let compact = cfg
178 .compaction_enabled
179 .then(|| Compactor::new(cfg.clone(), Arc::clone(&metrics), gc.clone()));
180 let debug_state = HandleDebugState {
181 hostname: cfg.hostname.to_owned(),
182 purpose: purpose.to_owned(),
183 };
184 let upper = machine.applier.clone_upper();
185 let expire_fn = Self::expire_fn(machine.clone(), gc.clone(), writer_id.clone());
186 WriteHandle {
187 cfg,
188 metrics,
189 machine,
190 gc,
191 compact,
192 blob,
193 isolated_runtime,
194 writer_id,
195 debug_state,
196 write_schemas,
197 upper,
198 expire_fn: Some(expire_fn),
199 }
200 }
201
202 pub fn from_read(read: &ReadHandle<K, V, T, D>, purpose: &str) -> Self {
205 Self::new(
206 read.cfg.clone(),
207 Arc::clone(&read.metrics),
208 read.machine.clone(),
209 read.gc.clone(),
210 Arc::clone(&read.blob),
211 WriterId::new(),
212 purpose,
213 read.read_schemas.clone(),
214 )
215 }
216
217 pub fn validate_part_bounds_on_write(&self) -> bool {
220 VALIDATE_PART_BOUNDS_ON_WRITE.get(&self.cfg) || VALIDATE_PART_BOUNDS_ON_READ.get(&self.cfg)
223 }
224
225 pub fn shard_id(&self) -> ShardId {
227 self.machine.shard_id()
228 }
229
230 pub fn schema_id(&self) -> Option<SchemaId> {
232 self.write_schemas.id
233 }
234
235 pub async fn try_register_schema(&mut self) -> Option<SchemaId> {
241 let Schemas { id, key, val } = &self.write_schemas;
242
243 if let Some(id) = id {
244 return Some(*id);
245 }
246
247 let (schema_id, maintenance) = self.machine.register_schema(key, val).await;
248 maintenance.start_performing(&self.machine, &self.gc);
249
250 self.write_schemas.id = schema_id;
251 schema_id
252 }
253
254 pub fn upper(&self) -> &Antichain<T> {
261 &self.upper
262 }
263
264 pub fn shared_upper(&self) -> Antichain<T> {
270 self.machine.applier.clone_upper()
271 }
272
273 #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
279 pub async fn fetch_recent_upper(&mut self) -> &Antichain<T> {
280 self.machine
283 .applier
284 .fetch_upper(|current_upper| self.upper.clone_from(current_upper))
285 .await;
286 &self.upper
287 }
288
289 pub async fn advance_upper(&mut self, target: &Antichain<T>) {
297 let mut lower = self.shared_upper().clone();
300
301 while !PartialOrder::less_equal(target, &lower) {
302 let since = Antichain::from_elem(T::minimum());
303 let desc = Description::new(lower.clone(), target.clone(), since);
304 let batch = HollowBatch::empty(desc);
305
306 let res = self
307 .machine
308 .compare_and_append(&batch, &self.writer_id, &self.debug_state)
309 .await;
310
311 use CompareAndAppendRes::*;
312 let new_upper = match res {
313 Success(_seq_no, maintenance) => {
314 maintenance.start_performing(&self.machine, &self.gc, self.compact.as_ref());
315 batch.desc.upper().clone()
316 }
317 UpperMismatch(_seq_no, actual_upper) => actual_upper,
318 InvalidUsage(_invalid_usage) => unreachable!("batch bounds checked above"),
319 InlineBackpressure => unreachable!("batch was empty"),
320 };
321
322 self.upper.clone_from(&new_upper);
323 lower = new_upper;
324 }
325 }
326
327 #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
357 pub async fn append<SB, KB, VB, TB, DB, I>(
358 &mut self,
359 updates: I,
360 lower: Antichain<T>,
361 upper: Antichain<T>,
362 ) -> Result<Result<(), UpperMismatch<T>>, InvalidUsage<T>>
363 where
364 SB: Borrow<((KB, VB), TB, DB)>,
365 KB: Borrow<K>,
366 VB: Borrow<V>,
367 TB: Borrow<T>,
368 DB: Borrow<D>,
369 I: IntoIterator<Item = SB>,
370 D: Send + Sync,
371 {
372 let batch = self.batch(updates, lower.clone(), upper.clone()).await?;
373 self.append_batch(batch, lower, upper).await
374 }
375
376 #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
405 pub async fn compare_and_append<SB, KB, VB, TB, DB, I>(
406 &mut self,
407 updates: I,
408 expected_upper: Antichain<T>,
409 new_upper: Antichain<T>,
410 ) -> Result<Result<(), UpperMismatch<T>>, InvalidUsage<T>>
411 where
412 SB: Borrow<((KB, VB), TB, DB)>,
413 KB: Borrow<K>,
414 VB: Borrow<V>,
415 TB: Borrow<T>,
416 DB: Borrow<D>,
417 I: IntoIterator<Item = SB>,
418 D: Send + Sync,
419 {
420 let mut batch = self
421 .batch(updates, expected_upper.clone(), new_upper.clone())
422 .await?;
423 match self
424 .compare_and_append_batch(&mut [&mut batch], expected_upper, new_upper, true)
425 .await
426 {
427 ok @ Ok(Ok(())) => ok,
428 err => {
429 batch.delete().await;
434 err
435 }
436 }
437 }
438
439 #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
465 pub async fn append_batch(
466 &mut self,
467 mut batch: Batch<K, V, T, D>,
468 mut lower: Antichain<T>,
469 upper: Antichain<T>,
470 ) -> Result<Result<(), UpperMismatch<T>>, InvalidUsage<T>>
471 where
472 D: Send + Sync,
473 {
474 loop {
475 let res = self
476 .compare_and_append_batch(&mut [&mut batch], lower.clone(), upper.clone(), true)
477 .await?;
478 match res {
479 Ok(()) => {
480 self.upper = upper;
481 return Ok(Ok(()));
482 }
483 Err(mismatch) => {
484 if PartialOrder::less_than(&mismatch.current, &lower) {
486 self.upper.clone_from(&mismatch.current);
487
488 batch.delete().await;
489
490 return Ok(Err(mismatch));
491 } else if PartialOrder::less_than(&mismatch.current, &upper) {
492 lower = mismatch.current;
499 } else {
500 self.upper = mismatch.current;
502
503 batch.delete().await;
507
508 return Ok(Ok(()));
509 }
510 }
511 }
512 }
513 }
514
515 #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
549 pub async fn compare_and_append_batch(
550 &mut self,
551 batches: &mut [&mut Batch<K, V, T, D>],
552 expected_upper: Antichain<T>,
553 new_upper: Antichain<T>,
554 validate_part_bounds_on_write: bool,
555 ) -> Result<Result<(), UpperMismatch<T>>, InvalidUsage<T>>
556 where
557 D: Send + Sync,
558 {
559 let schema_id = self.try_register_schema().await;
563
564 for batch in batches.iter() {
565 if self.machine.shard_id() != batch.shard_id() {
566 return Err(InvalidUsage::BatchNotFromThisShard {
567 batch_shard: batch.shard_id(),
568 handle_shard: self.machine.shard_id(),
569 });
570 }
571 assert_code_can_read_data(&self.cfg.build_version, &batch.version);
572 if self.cfg.build_version > batch.version {
573 info!(
574 shard_id =? self.machine.shard_id(),
575 batch_version =? batch.version,
576 writer_version =? self.cfg.build_version,
577 "Appending batch from the past. This is fine but should be rare. \
578 TODO: Error on very old versions once the leaked blob detector exists."
579 )
580 }
581 fn assert_schema<A: Codec>(writer_schema: &A::Schema, batch_schema: &bytes::Bytes) {
582 if batch_schema.is_empty() {
583 return;
585 }
586 let batch_schema: A::Schema = A::decode_schema(batch_schema);
587 if *writer_schema != batch_schema {
588 error!(
589 ?writer_schema,
590 ?batch_schema,
591 "writer and batch schemas should be identical"
592 );
593 soft_panic_or_log!("writer and batch schemas should be identical");
594 }
595 }
596 assert_schema::<K>(&*self.write_schemas.key, &batch.schemas.0);
597 assert_schema::<V>(&*self.write_schemas.val, &batch.schemas.1);
598 }
599
600 let lower = expected_upper.clone();
601 let upper = new_upper;
602 let since = Antichain::from_elem(T::minimum());
603 let desc = Description::new(lower, upper, since);
604
605 let mut received_inline_backpressure = false;
606 let mut inline_batch_builder: Option<(_, BatchBuilder<K, V, T, D>)> = None;
613 let maintenance = loop {
614 let any_batch_rewrite = batches
615 .iter()
616 .any(|x| x.batch.parts.iter().any(|x| x.ts_rewrite().is_some()));
617 let (mut parts, mut num_updates, mut run_splits, mut run_metas) =
618 (vec![], 0, vec![], vec![]);
619 let mut key_storage = None;
620 let mut val_storage = None;
621 for batch in batches.iter() {
622 let bounds_truncated = validate_truncate_batch(
623 &batch.batch,
624 &desc,
625 any_batch_rewrite,
626 validate_part_bounds_on_write,
627 )?;
628 for (run_meta, run) in batch.batch.runs() {
629 let start_index = parts.len();
630 for part in run {
631 if let (
632 RunPart::Single(
633 batch_part @ BatchPart::Inline {
634 updates,
635 ts_rewrite,
636 schema_id: _,
637 deprecated_schema_id: _,
638 },
639 ),
640 Some((schema_cache, builder)),
641 ) = (part, &mut inline_batch_builder)
642 {
643 let schema_migration = PartMigration::new(
644 batch_part,
645 self.write_schemas.clone(),
646 schema_cache,
647 )
648 .await
649 .expect("schemas for inline user part");
650
651 let encoded_part = EncodedPart::from_inline(
652 &crate::fetch::FetchConfig::from_persist_config(&self.cfg),
653 &*self.metrics,
654 self.metrics.read.compaction.clone(),
655 desc.clone(),
656 updates,
657 ts_rewrite.as_ref(),
658 );
659 let mut fetched_part = FetchedPart::new(
660 Arc::clone(&self.metrics),
661 encoded_part,
662 schema_migration,
663 FetchBatchFilter::Compaction {
664 since: desc.since().clone(),
665 },
666 false,
667 PartDecodeFormat::Arrow,
668 None,
669 );
670
671 while let Some(((k, v), t, d)) =
672 fetched_part.next_with_storage(&mut key_storage, &mut val_storage)
673 {
674 builder
675 .add(&k, &v, &t, &d)
676 .await
677 .expect("re-encoding just-decoded data");
678 }
679 } else {
680 parts.push(part.clone())
681 }
682 }
683
684 let end_index = parts.len();
685
686 if start_index == end_index {
687 continue;
688 }
689
690 if start_index != 0 {
692 run_splits.push(start_index);
693 }
694 let mut run_meta = run_meta.clone();
695 if bounds_truncated {
696 run_meta.set_bounds_truncated();
697 }
698 run_metas.push(run_meta);
699 }
700 num_updates += batch.batch.len;
701 }
702
703 let mut flushed_inline_batch = if let Some((_, builder)) = inline_batch_builder.take() {
704 let mut finished = builder
705 .finish(desc.upper().clone())
706 .await
707 .expect("invalid usage");
708 let cfg = BatchBuilderConfig::new(&self.cfg, self.shard_id());
709 finished
710 .flush_to_blob(
711 &cfg,
712 &self.metrics.inline.backpressure,
713 &self.isolated_runtime,
714 &self.write_schemas,
715 )
716 .await;
717 Some(finished)
718 } else {
719 None
720 };
721
722 if let Some(batch) = &flushed_inline_batch {
723 for (run_meta, run) in batch.batch.runs() {
724 assert!(run.len() > 0);
725 let start_index = parts.len();
726 if start_index != 0 {
727 run_splits.push(start_index);
728 }
729 run_metas.push(run_meta.clone());
730 parts.extend(run.iter().cloned())
731 }
732 }
733
734 let mut combined_batch =
735 HollowBatch::new(desc.clone(), parts, num_updates, run_metas, run_splits);
736
737 match schema_id {
741 Some(schema_id) => {
742 ensure_batch_schema(&mut combined_batch, self.shard_id(), schema_id);
743 }
744 None => {
745 assert!(
746 self.fetch_recent_upper().await.is_empty(),
747 "fetching a schema id should only fail when the shard is tombstoned"
748 )
749 }
750 }
751
752 let res = self
753 .machine
754 .compare_and_append(&combined_batch, &self.writer_id, &self.debug_state)
755 .await;
756
757 match res {
758 CompareAndAppendRes::Success(_seqno, maintenance) => {
759 self.upper.clone_from(desc.upper());
760 for batch in batches.iter_mut() {
761 batch.mark_consumed();
762 }
763 if let Some(batch) = &mut flushed_inline_batch {
764 batch.mark_consumed();
765 }
766 break maintenance;
767 }
768 CompareAndAppendRes::InvalidUsage(invalid_usage) => {
769 if let Some(batch) = flushed_inline_batch.take() {
770 batch.delete().await;
771 }
772 return Err(invalid_usage);
773 }
774 CompareAndAppendRes::UpperMismatch(_seqno, current_upper) => {
775 if let Some(batch) = flushed_inline_batch.take() {
776 batch.delete().await;
777 }
778 self.upper.clone_from(¤t_upper);
781 return Ok(Err(UpperMismatch {
782 current: current_upper,
783 expected: expected_upper,
784 }));
785 }
786 CompareAndAppendRes::InlineBackpressure => {
787 assert_eq!(received_inline_backpressure, false);
790 received_inline_backpressure = true;
791 if COMBINE_INLINE_WRITES.get(&self.cfg) {
792 inline_batch_builder = Some((
793 self.machine.applier.schema_cache(),
794 self.builder(desc.lower().clone()),
795 ));
796 continue;
797 }
798
799 let cfg = BatchBuilderConfig::new(&self.cfg, self.shard_id());
800 let flush_batches = batches
803 .iter_mut()
804 .map(|batch| async {
805 batch
806 .flush_to_blob(
807 &cfg,
808 &self.metrics.inline.backpressure,
809 &self.isolated_runtime,
810 &self.write_schemas,
811 )
812 .await
813 })
814 .collect::<FuturesUnordered<_>>();
815 let () = flush_batches.collect::<()>().await;
816
817 for batch in batches.iter() {
818 assert_eq!(batch.batch.inline_bytes(), 0);
819 }
820
821 continue;
822 }
823 }
824 };
825
826 maintenance.start_performing(&self.machine, &self.gc, self.compact.as_ref());
827
828 Ok(Ok(()))
829 }
830
831 pub fn batch_from_transmittable_batch(&self, batch: ProtoBatch) -> Batch<K, V, T, D> {
834 let shard_id: ShardId = batch
835 .shard_id
836 .into_rust()
837 .expect("valid transmittable batch");
838 assert_eq!(shard_id, self.machine.shard_id());
839
840 let ret = Batch {
841 batch_delete_enabled: BATCH_DELETE_ENABLED.get(&self.cfg),
842 metrics: Arc::clone(&self.metrics),
843 shard_metrics: Arc::clone(&self.machine.applier.shard_metrics),
844 version: Version::parse(&batch.version).expect("valid transmittable batch"),
845 schemas: (batch.key_schema, batch.val_schema),
846 batch: batch
847 .batch
848 .into_rust_if_some("ProtoBatch::batch")
849 .expect("valid transmittable batch"),
850 blob: Arc::clone(&self.blob),
851 _phantom: std::marker::PhantomData,
852 };
853 assert_eq!(ret.shard_id(), self.machine.shard_id());
854 ret
855 }
856
857 pub fn builder(&self, lower: Antichain<T>) -> BatchBuilder<K, V, T, D> {
870 Self::builder_inner(
871 &self.cfg,
872 CompactConfig::new(&self.cfg, self.shard_id()),
873 Arc::clone(&self.metrics),
874 Arc::clone(&self.machine.applier.shard_metrics),
875 &self.metrics.user,
876 Arc::clone(&self.isolated_runtime),
877 Arc::clone(&self.blob),
878 self.shard_id(),
879 self.write_schemas.clone(),
880 lower,
881 )
882 }
883
884 pub(crate) fn builder_inner(
887 persist_cfg: &PersistConfig,
888 compact_cfg: CompactConfig,
889 metrics: Arc<Metrics>,
890 shard_metrics: Arc<ShardMetrics>,
891 user_batch_metrics: &BatchWriteMetrics,
892 isolated_runtime: Arc<IsolatedRuntime>,
893 blob: Arc<dyn Blob>,
894 shard_id: ShardId,
895 schemas: Schemas<K, V>,
896 lower: Antichain<T>,
897 ) -> BatchBuilder<K, V, T, D> {
898 let parts = if let Some(max_runs) = compact_cfg.batch.max_runs {
899 BatchParts::new_compacting::<K, V, D>(
900 compact_cfg,
901 Description::new(
902 lower.clone(),
903 Antichain::new(),
904 Antichain::from_elem(T::minimum()),
905 ),
906 max_runs,
907 Arc::clone(&metrics),
908 shard_metrics,
909 shard_id,
910 Arc::clone(&blob),
911 isolated_runtime,
912 user_batch_metrics,
913 schemas.clone(),
914 )
915 } else {
916 BatchParts::new_ordered::<D>(
917 compact_cfg.batch,
918 RunOrder::Unordered,
919 Arc::clone(&metrics),
920 shard_metrics,
921 shard_id,
922 Arc::clone(&blob),
923 isolated_runtime,
924 user_batch_metrics,
925 )
926 };
927 let builder = BatchBuilderInternal::new(
928 BatchBuilderConfig::new(persist_cfg, shard_id),
929 parts,
930 metrics,
931 schemas,
932 blob,
933 shard_id,
934 persist_cfg.build_version.clone(),
935 );
936 BatchBuilder::new(
937 builder,
938 Description::new(lower, Antichain::new(), Antichain::from_elem(T::minimum())),
939 )
940 }
941
942 #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
945 pub async fn batch<SB, KB, VB, TB, DB, I>(
946 &mut self,
947 updates: I,
948 lower: Antichain<T>,
949 upper: Antichain<T>,
950 ) -> Result<Batch<K, V, T, D>, InvalidUsage<T>>
951 where
952 SB: Borrow<((KB, VB), TB, DB)>,
953 KB: Borrow<K>,
954 VB: Borrow<V>,
955 TB: Borrow<T>,
956 DB: Borrow<D>,
957 I: IntoIterator<Item = SB>,
958 {
959 let iter = updates.into_iter();
960
961 let mut builder = self.builder(lower.clone());
962
963 for update in iter {
964 let ((k, v), t, d) = update.borrow();
965 let (k, v, t, d) = (k.borrow(), v.borrow(), t.borrow(), d.borrow());
966 match builder.add(k, v, t, d).await {
967 Ok(Added::Record | Added::RecordAndParts) => (),
968 Err(invalid_usage) => return Err(invalid_usage),
969 }
970 }
971
972 builder.finish(upper.clone()).await
973 }
974
975 pub async fn wait_for_upper_past(&mut self, frontier: &Antichain<T>) {
977 let mut watch = self.machine.applier.watch();
978 self.machine
979 .wait_for_upper_past(
980 frontier,
981 &mut watch,
982 None,
983 &self.metrics.retries.next_listen_batch, next_listen_batch_retry_params(&self.cfg),
985 )
986 .await;
987 let upper = self.machine.applier.clone_upper();
988 if PartialOrder::less_than(&self.upper, &upper) {
989 self.upper.clone_from(&upper);
990 }
991 assert!(PartialOrder::less_than(frontier, &self.upper));
992 }
993
994 #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
1003 pub async fn expire(mut self) {
1004 let Some(expire_fn) = self.expire_fn.take() else {
1005 return;
1006 };
1007 expire_fn.0().await;
1008 }
1009
1010 fn expire_fn(
1011 machine: Machine<K, V, T, D>,
1012 gc: GarbageCollector<K, V, T, D>,
1013 writer_id: WriterId,
1014 ) -> ExpireFn {
1015 ExpireFn(Box::new(move || {
1016 Box::pin(async move {
1017 let (_, maintenance) = machine.expire_writer(&writer_id).await;
1018 maintenance.start_performing(&machine, &gc);
1019 })
1020 }))
1021 }
1022
1023 #[cfg(test)]
1025 #[track_caller]
1026 pub async fn expect_append<L, U>(&mut self, updates: &[((K, V), T, D)], lower: L, new_upper: U)
1027 where
1028 L: Into<Antichain<T>>,
1029 U: Into<Antichain<T>>,
1030 D: Send + Sync,
1031 {
1032 self.append(updates.iter(), lower.into(), new_upper.into())
1033 .await
1034 .expect("invalid usage")
1035 .expect("unexpected upper");
1036 }
1037
1038 #[cfg(test)]
1041 #[track_caller]
1042 pub async fn expect_compare_and_append(
1043 &mut self,
1044 updates: &[((K, V), T, D)],
1045 expected_upper: T,
1046 new_upper: T,
1047 ) where
1048 D: Send + Sync,
1049 {
1050 self.compare_and_append(
1051 updates.iter().map(|((k, v), t, d)| ((k, v), t, d)),
1052 Antichain::from_elem(expected_upper),
1053 Antichain::from_elem(new_upper),
1054 )
1055 .await
1056 .expect("invalid usage")
1057 .expect("unexpected upper")
1058 }
1059
1060 #[cfg(test)]
1063 #[track_caller]
1064 pub async fn expect_compare_and_append_batch(
1065 &mut self,
1066 batches: &mut [&mut Batch<K, V, T, D>],
1067 expected_upper: T,
1068 new_upper: T,
1069 ) {
1070 self.compare_and_append_batch(
1071 batches,
1072 Antichain::from_elem(expected_upper),
1073 Antichain::from_elem(new_upper),
1074 true,
1075 )
1076 .await
1077 .expect("invalid usage")
1078 .expect("unexpected upper")
1079 }
1080
1081 #[cfg(test)]
1083 #[track_caller]
1084 pub async fn expect_batch(
1085 &mut self,
1086 updates: &[((K, V), T, D)],
1087 lower: T,
1088 upper: T,
1089 ) -> Batch<K, V, T, D> {
1090 self.batch(
1091 updates.iter(),
1092 Antichain::from_elem(lower),
1093 Antichain::from_elem(upper),
1094 )
1095 .await
1096 .expect("invalid usage")
1097 }
1098}
1099
1100impl<K: Codec, V: Codec, T, D> Drop for WriteHandle<K, V, T, D> {
1101 fn drop(&mut self) {
1102 let Some(expire_fn) = self.expire_fn.take() else {
1103 return;
1104 };
1105 let handle = match Handle::try_current() {
1106 Ok(x) => x,
1107 Err(_) => {
1108 warn!(
1109 "WriteHandle {} dropped without being explicitly expired, falling back to lease timeout",
1110 self.writer_id
1111 );
1112 return;
1113 }
1114 };
1115 let expire_span = debug_span!("drop::expire");
1121 handle.spawn_named(
1122 || format!("WriteHandle::expire ({})", self.writer_id),
1123 expire_fn.0().instrument(expire_span),
1124 );
1125 }
1126}
1127
1128fn ensure_batch_schema<T>(batch: &mut HollowBatch<T>, shard_id: ShardId, schema_id: SchemaId)
1133where
1134 T: Timestamp + Lattice + Codec64,
1135{
1136 let ensure = |id: &mut Option<SchemaId>| match id {
1137 Some(id) => assert_eq!(*id, schema_id, "schema ID mismatch; shard={shard_id}"),
1138 None => *id = Some(schema_id),
1139 };
1140
1141 for run_meta in &mut batch.run_meta {
1142 ensure(&mut run_meta.schema);
1143 }
1144 for part in &mut batch.parts {
1145 match part {
1146 RunPart::Single(BatchPart::Hollow(part)) => ensure(&mut part.schema_id),
1147 RunPart::Single(BatchPart::Inline { schema_id, .. }) => ensure(schema_id),
1148 RunPart::Many(_hollow_run_ref) => {
1149 }
1153 }
1154 }
1155}
1156
1157#[cfg(test)]
1158mod tests {
1159 use std::str::FromStr;
1160 use std::sync::mpsc;
1161
1162 use differential_dataflow::consolidation::consolidate_updates;
1163 use futures_util::FutureExt;
1164 use mz_dyncfg::ConfigUpdates;
1165 use mz_ore::collections::CollectionExt;
1166 use mz_ore::task;
1167 use serde_json::json;
1168
1169 use crate::cache::PersistClientCache;
1170 use crate::tests::{all_ok, new_test_client};
1171 use crate::{PersistLocation, ShardId};
1172
1173 use super::*;
1174
1175 #[mz_persist_proc::test(tokio::test)]
1176 #[cfg_attr(miri, ignore)] async fn empty_batches(dyncfgs: ConfigUpdates) {
1178 let data = [
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 (mut write, _) = new_test_client(&dyncfgs)
1185 .await
1186 .expect_open::<String, String, u64, i64>(ShardId::new())
1187 .await;
1188 let blob = Arc::clone(&write.blob);
1189
1190 let mut upper = 3;
1192 write.expect_append(&data[..2], vec![0], vec![upper]).await;
1193
1194 let mut count_before = 0;
1196 blob.list_keys_and_metadata("", &mut |_| {
1197 count_before += 1;
1198 })
1199 .await
1200 .expect("list_keys failed");
1201 for _ in 0..5 {
1202 let new_upper = upper + 1;
1203 write.expect_compare_and_append(&[], upper, new_upper).await;
1204 upper = new_upper;
1205 }
1206 let mut count_after = 0;
1207 blob.list_keys_and_metadata("", &mut |_| {
1208 count_after += 1;
1209 })
1210 .await
1211 .expect("list_keys failed");
1212 assert_eq!(count_after, count_before);
1213 }
1214
1215 #[mz_persist_proc::test(tokio::test)]
1216 #[cfg_attr(miri, ignore)] async fn compare_and_append_batch_multi(dyncfgs: ConfigUpdates) {
1218 let data0 = vec![
1219 (("1".to_owned(), "one".to_owned()), 1, 1),
1220 (("2".to_owned(), "two".to_owned()), 2, 1),
1221 (("4".to_owned(), "four".to_owned()), 4, 1),
1222 ];
1223 let data1 = vec![
1224 (("1".to_owned(), "one".to_owned()), 1, 1),
1225 (("2".to_owned(), "two".to_owned()), 2, 1),
1226 (("3".to_owned(), "three".to_owned()), 3, 1),
1227 ];
1228
1229 let (mut write, mut read) = new_test_client(&dyncfgs)
1230 .await
1231 .expect_open::<String, String, u64, i64>(ShardId::new())
1232 .await;
1233
1234 let mut batch0 = write.expect_batch(&data0, 0, 5).await;
1235 let mut batch1 = write.expect_batch(&data1, 0, 4).await;
1236
1237 write
1238 .expect_compare_and_append_batch(&mut [&mut batch0, &mut batch1], 0, 4)
1239 .await;
1240
1241 let batch = write
1242 .machine
1243 .unleased_snapshot(&Antichain::from_elem(3))
1244 .await
1245 .expect("just wrote this")
1246 .into_element();
1247
1248 assert!(batch.runs().count() >= 2);
1249
1250 let expected = vec![
1251 (("1".to_owned(), "one".to_owned()), 1, 2),
1252 (("2".to_owned(), "two".to_owned()), 2, 2),
1253 (("3".to_owned(), "three".to_owned()), 3, 1),
1254 ];
1255 let mut actual = read.expect_snapshot_and_fetch(3).await;
1256 consolidate_updates(&mut actual);
1257 assert_eq!(actual, all_ok(&expected, 3));
1258 }
1259
1260 #[mz_ore::test]
1261 fn writer_id_human_readable_serde() {
1262 #[derive(Debug, Serialize, Deserialize)]
1263 struct Container {
1264 writer_id: WriterId,
1265 }
1266
1267 let id = WriterId::from_str("w00000000-1234-5678-0000-000000000000").expect("valid id");
1269 assert_eq!(
1270 id,
1271 serde_json::from_value(serde_json::to_value(id.clone()).expect("serializable"))
1272 .expect("deserializable")
1273 );
1274
1275 assert_eq!(
1277 id,
1278 serde_json::from_str("\"w00000000-1234-5678-0000-000000000000\"")
1279 .expect("deserializable")
1280 );
1281
1282 let json = json!({ "writer_id": id });
1284 assert_eq!(
1285 "{\"writer_id\":\"w00000000-1234-5678-0000-000000000000\"}",
1286 &json.to_string()
1287 );
1288 let container: Container = serde_json::from_value(json).expect("deserializable");
1289 assert_eq!(container.writer_id, id);
1290 }
1291
1292 #[mz_persist_proc::test(tokio::test)]
1293 #[cfg_attr(miri, ignore)] async fn hollow_batch_roundtrip(dyncfgs: ConfigUpdates) {
1295 let data = vec![
1296 (("1".to_owned(), "one".to_owned()), 1, 1),
1297 (("2".to_owned(), "two".to_owned()), 2, 1),
1298 (("3".to_owned(), "three".to_owned()), 3, 1),
1299 ];
1300
1301 let (mut write, mut read) = new_test_client(&dyncfgs)
1302 .await
1303 .expect_open::<String, String, u64, i64>(ShardId::new())
1304 .await;
1305
1306 let batch = write.expect_batch(&data, 0, 4).await;
1311 let hollow_batch = batch.into_transmittable_batch();
1312 let mut rehydrated_batch = write.batch_from_transmittable_batch(hollow_batch);
1313
1314 write
1315 .expect_compare_and_append_batch(&mut [&mut rehydrated_batch], 0, 4)
1316 .await;
1317
1318 let expected = vec![
1319 (("1".to_owned(), "one".to_owned()), 1, 1),
1320 (("2".to_owned(), "two".to_owned()), 2, 1),
1321 (("3".to_owned(), "three".to_owned()), 3, 1),
1322 ];
1323 let mut actual = read.expect_snapshot_and_fetch(3).await;
1324 consolidate_updates(&mut actual);
1325 assert_eq!(actual, all_ok(&expected, 3));
1326 }
1327
1328 #[mz_persist_proc::test(tokio::test)]
1329 #[cfg_attr(miri, ignore)] async fn wait_for_upper_past(dyncfgs: ConfigUpdates) {
1331 let client = new_test_client(&dyncfgs).await;
1332 let (mut write, _) = client.expect_open::<(), (), u64, i64>(ShardId::new()).await;
1333 let five = Antichain::from_elem(5);
1334
1335 assert_eq!(write.wait_for_upper_past(&five).now_or_never(), None);
1337
1338 write
1340 .expect_compare_and_append(&[(((), ()), 1, 1)], 0, 5)
1341 .await;
1342 assert_eq!(write.wait_for_upper_past(&five).now_or_never(), None);
1343
1344 write
1346 .expect_compare_and_append(&[(((), ()), 5, 1)], 5, 7)
1347 .await;
1348 assert_eq!(write.wait_for_upper_past(&five).now_or_never(), Some(()));
1349 assert_eq!(write.upper(), &Antichain::from_elem(7));
1350
1351 assert_eq!(
1354 write
1355 .wait_for_upper_past(&Antichain::from_elem(2))
1356 .now_or_never(),
1357 Some(())
1358 );
1359 assert_eq!(write.upper(), &Antichain::from_elem(7));
1360 }
1361
1362 #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
1363 #[cfg_attr(miri, ignore)] async fn fetch_recent_upper_linearized() {
1365 type Timestamp = u64;
1366 let max_upper = 1000;
1367
1368 let shard_id = ShardId::new();
1369 let mut clients = PersistClientCache::new_no_metrics();
1370 let upper_writer_client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
1371 let (mut upper_writer, _) = upper_writer_client
1372 .expect_open::<(), (), Timestamp, i64>(shard_id)
1373 .await;
1374 clients.clear_state_cache();
1377 let upper_reader_client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
1378 let (mut upper_reader, _) = upper_reader_client
1379 .expect_open::<(), (), Timestamp, i64>(shard_id)
1380 .await;
1381 let (tx, rx) = mpsc::channel();
1382
1383 let task = task::spawn(|| "upper-reader", async move {
1384 let mut upper = Timestamp::MIN;
1385
1386 while upper < max_upper {
1387 while let Ok(new_upper) = rx.try_recv() {
1388 upper = new_upper;
1389 }
1390
1391 let recent_upper = upper_reader
1392 .fetch_recent_upper()
1393 .await
1394 .as_option()
1395 .cloned()
1396 .expect("u64 is totally ordered and the shard is not finalized");
1397 assert!(
1398 recent_upper >= upper,
1399 "recent upper {recent_upper:?} is less than known upper {upper:?}"
1400 );
1401 }
1402 });
1403
1404 for upper in Timestamp::MIN..max_upper {
1405 let next_upper = upper + 1;
1406 upper_writer
1407 .expect_compare_and_append(&[], upper, next_upper)
1408 .await;
1409 tx.send(next_upper).expect("send failed");
1410 }
1411
1412 task.await;
1413 }
1414}