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