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 () = 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 run_metas.push(run_meta.clone());
695 }
696 num_updates += batch.batch.len;
697 }
698
699 let mut flushed_inline_batch = if let Some((_, builder)) = inline_batch_builder.take() {
700 let mut finished = builder
701 .finish(desc.upper().clone())
702 .await
703 .expect("invalid usage");
704 let cfg = BatchBuilderConfig::new(&self.cfg, self.shard_id());
705 finished
706 .flush_to_blob(
707 &cfg,
708 &self.metrics.inline.backpressure,
709 &self.isolated_runtime,
710 &self.write_schemas,
711 )
712 .await;
713 Some(finished)
714 } else {
715 None
716 };
717
718 if let Some(batch) = &flushed_inline_batch {
719 for (run_meta, run) in batch.batch.runs() {
720 assert!(run.len() > 0);
721 let start_index = parts.len();
722 if start_index != 0 {
723 run_splits.push(start_index);
724 }
725 run_metas.push(run_meta.clone());
726 parts.extend(run.iter().cloned())
727 }
728 }
729
730 let mut combined_batch =
731 HollowBatch::new(desc.clone(), parts, num_updates, run_metas, run_splits);
732
733 match schema_id {
737 Some(schema_id) => {
738 ensure_batch_schema(&mut combined_batch, self.shard_id(), schema_id);
739 }
740 None => {
741 assert!(
742 self.fetch_recent_upper().await.is_empty(),
743 "fetching a schema id should only fail when the shard is tombstoned"
744 )
745 }
746 }
747
748 let res = self
749 .machine
750 .compare_and_append(&combined_batch, &self.writer_id, &self.debug_state)
751 .await;
752
753 match res {
754 CompareAndAppendRes::Success(_seqno, maintenance) => {
755 self.upper.clone_from(desc.upper());
756 for batch in batches.iter_mut() {
757 batch.mark_consumed();
758 }
759 if let Some(batch) = &mut flushed_inline_batch {
760 batch.mark_consumed();
761 }
762 break maintenance;
763 }
764 CompareAndAppendRes::InvalidUsage(invalid_usage) => {
765 if let Some(batch) = flushed_inline_batch.take() {
766 batch.delete().await;
767 }
768 return Err(invalid_usage);
769 }
770 CompareAndAppendRes::UpperMismatch(_seqno, current_upper) => {
771 if let Some(batch) = flushed_inline_batch.take() {
772 batch.delete().await;
773 }
774 self.upper.clone_from(¤t_upper);
777 return Ok(Err(UpperMismatch {
778 current: current_upper,
779 expected: expected_upper,
780 }));
781 }
782 CompareAndAppendRes::InlineBackpressure => {
783 assert_eq!(received_inline_backpressure, false);
786 received_inline_backpressure = true;
787 if COMBINE_INLINE_WRITES.get(&self.cfg) {
788 inline_batch_builder = Some((
789 self.machine.applier.schema_cache(),
790 self.builder(desc.lower().clone()),
791 ));
792 continue;
793 }
794
795 let cfg = BatchBuilderConfig::new(&self.cfg, self.shard_id());
796 let flush_batches = batches
799 .iter_mut()
800 .map(|batch| async {
801 batch
802 .flush_to_blob(
803 &cfg,
804 &self.metrics.inline.backpressure,
805 &self.isolated_runtime,
806 &self.write_schemas,
807 )
808 .await
809 })
810 .collect::<FuturesUnordered<_>>();
811 let () = flush_batches.collect::<()>().await;
812
813 for batch in batches.iter() {
814 assert_eq!(batch.batch.inline_bytes(), 0);
815 }
816
817 continue;
818 }
819 }
820 };
821
822 maintenance.start_performing(&self.machine, &self.gc, self.compact.as_ref());
823
824 Ok(Ok(()))
825 }
826
827 pub fn batch_from_transmittable_batch(&self, batch: ProtoBatch) -> Batch<K, V, T, D> {
830 let shard_id: ShardId = batch
831 .shard_id
832 .into_rust()
833 .expect("valid transmittable batch");
834 assert_eq!(shard_id, self.machine.shard_id());
835
836 let ret = Batch {
837 batch_delete_enabled: BATCH_DELETE_ENABLED.get(&self.cfg),
838 metrics: Arc::clone(&self.metrics),
839 shard_metrics: Arc::clone(&self.machine.applier.shard_metrics),
840 version: Version::parse(&batch.version).expect("valid transmittable batch"),
841 schemas: (batch.key_schema, batch.val_schema),
842 batch: batch
843 .batch
844 .into_rust_if_some("ProtoBatch::batch")
845 .expect("valid transmittable batch"),
846 blob: Arc::clone(&self.blob),
847 _phantom: std::marker::PhantomData,
848 };
849 assert_eq!(ret.shard_id(), self.machine.shard_id());
850 ret
851 }
852
853 pub fn builder(&self, lower: Antichain<T>) -> BatchBuilder<K, V, T, D> {
866 Self::builder_inner(
867 &self.cfg,
868 CompactConfig::new(&self.cfg, self.shard_id()),
869 Arc::clone(&self.metrics),
870 Arc::clone(&self.machine.applier.shard_metrics),
871 &self.metrics.user,
872 Arc::clone(&self.isolated_runtime),
873 Arc::clone(&self.blob),
874 self.shard_id(),
875 self.write_schemas.clone(),
876 lower,
877 )
878 }
879
880 pub(crate) fn builder_inner(
883 persist_cfg: &PersistConfig,
884 compact_cfg: CompactConfig,
885 metrics: Arc<Metrics>,
886 shard_metrics: Arc<ShardMetrics>,
887 user_batch_metrics: &BatchWriteMetrics,
888 isolated_runtime: Arc<IsolatedRuntime>,
889 blob: Arc<dyn Blob>,
890 shard_id: ShardId,
891 schemas: Schemas<K, V>,
892 lower: Antichain<T>,
893 ) -> BatchBuilder<K, V, T, D> {
894 let parts = if let Some(max_runs) = compact_cfg.batch.max_runs {
895 BatchParts::new_compacting::<K, V, D>(
896 compact_cfg,
897 Description::new(
898 lower.clone(),
899 Antichain::new(),
900 Antichain::from_elem(T::minimum()),
901 ),
902 max_runs,
903 Arc::clone(&metrics),
904 shard_metrics,
905 shard_id,
906 Arc::clone(&blob),
907 isolated_runtime,
908 user_batch_metrics,
909 schemas.clone(),
910 )
911 } else {
912 BatchParts::new_ordered::<D>(
913 compact_cfg.batch,
914 RunOrder::Unordered,
915 Arc::clone(&metrics),
916 shard_metrics,
917 shard_id,
918 Arc::clone(&blob),
919 isolated_runtime,
920 user_batch_metrics,
921 )
922 };
923 let builder = BatchBuilderInternal::new(
924 BatchBuilderConfig::new(persist_cfg, shard_id),
925 parts,
926 metrics,
927 schemas,
928 blob,
929 shard_id,
930 persist_cfg.build_version.clone(),
931 );
932 BatchBuilder::new(
933 builder,
934 Description::new(lower, Antichain::new(), Antichain::from_elem(T::minimum())),
935 )
936 }
937
938 #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
941 pub async fn batch<SB, KB, VB, TB, DB, I>(
942 &mut self,
943 updates: I,
944 lower: Antichain<T>,
945 upper: Antichain<T>,
946 ) -> Result<Batch<K, V, T, D>, InvalidUsage<T>>
947 where
948 SB: Borrow<((KB, VB), TB, DB)>,
949 KB: Borrow<K>,
950 VB: Borrow<V>,
951 TB: Borrow<T>,
952 DB: Borrow<D>,
953 I: IntoIterator<Item = SB>,
954 {
955 let iter = updates.into_iter();
956
957 let mut builder = self.builder(lower.clone());
958
959 for update in iter {
960 let ((k, v), t, d) = update.borrow();
961 let (k, v, t, d) = (k.borrow(), v.borrow(), t.borrow(), d.borrow());
962 match builder.add(k, v, t, d).await {
963 Ok(Added::Record | Added::RecordAndParts) => (),
964 Err(invalid_usage) => return Err(invalid_usage),
965 }
966 }
967
968 builder.finish(upper.clone()).await
969 }
970
971 pub async fn wait_for_upper_past(&mut self, frontier: &Antichain<T>) {
973 let mut watch = self.machine.applier.watch();
974 self.machine
975 .wait_for_upper_past(
976 frontier,
977 &mut watch,
978 None,
979 &self.metrics.retries.next_listen_batch, next_listen_batch_retry_params(&self.cfg),
981 )
982 .await;
983 let upper = self.machine.applier.clone_upper();
984 if PartialOrder::less_than(&self.upper, &upper) {
985 self.upper.clone_from(&upper);
986 }
987 assert!(PartialOrder::less_than(frontier, &self.upper));
988 }
989
990 #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
999 pub async fn expire(mut self) {
1000 let Some(expire_fn) = self.expire_fn.take() else {
1001 return;
1002 };
1003 expire_fn.0().await;
1004 }
1005
1006 fn expire_fn(
1007 machine: Machine<K, V, T, D>,
1008 gc: GarbageCollector<K, V, T, D>,
1009 writer_id: WriterId,
1010 ) -> ExpireFn {
1011 ExpireFn(Box::new(move || {
1012 Box::pin(async move {
1013 let (_, maintenance) = machine.expire_writer(&writer_id).await;
1014 maintenance.start_performing(&machine, &gc);
1015 })
1016 }))
1017 }
1018
1019 #[cfg(test)]
1021 #[track_caller]
1022 pub async fn expect_append<L, U>(&mut self, updates: &[((K, V), T, D)], lower: L, new_upper: U)
1023 where
1024 L: Into<Antichain<T>>,
1025 U: Into<Antichain<T>>,
1026 D: Send + Sync,
1027 {
1028 self.append(updates.iter(), lower.into(), new_upper.into())
1029 .await
1030 .expect("invalid usage")
1031 .expect("unexpected upper");
1032 }
1033
1034 #[cfg(test)]
1037 #[track_caller]
1038 pub async fn expect_compare_and_append(
1039 &mut self,
1040 updates: &[((K, V), T, D)],
1041 expected_upper: T,
1042 new_upper: T,
1043 ) where
1044 D: Send + Sync,
1045 {
1046 self.compare_and_append(
1047 updates.iter().map(|((k, v), t, d)| ((k, v), t, d)),
1048 Antichain::from_elem(expected_upper),
1049 Antichain::from_elem(new_upper),
1050 )
1051 .await
1052 .expect("invalid usage")
1053 .expect("unexpected upper")
1054 }
1055
1056 #[cfg(test)]
1059 #[track_caller]
1060 pub async fn expect_compare_and_append_batch(
1061 &mut self,
1062 batches: &mut [&mut Batch<K, V, T, D>],
1063 expected_upper: T,
1064 new_upper: T,
1065 ) {
1066 self.compare_and_append_batch(
1067 batches,
1068 Antichain::from_elem(expected_upper),
1069 Antichain::from_elem(new_upper),
1070 true,
1071 )
1072 .await
1073 .expect("invalid usage")
1074 .expect("unexpected upper")
1075 }
1076
1077 #[cfg(test)]
1079 #[track_caller]
1080 pub async fn expect_batch(
1081 &mut self,
1082 updates: &[((K, V), T, D)],
1083 lower: T,
1084 upper: T,
1085 ) -> Batch<K, V, T, D> {
1086 self.batch(
1087 updates.iter(),
1088 Antichain::from_elem(lower),
1089 Antichain::from_elem(upper),
1090 )
1091 .await
1092 .expect("invalid usage")
1093 }
1094}
1095
1096impl<K: Codec, V: Codec, T, D> Drop for WriteHandle<K, V, T, D> {
1097 fn drop(&mut self) {
1098 let Some(expire_fn) = self.expire_fn.take() else {
1099 return;
1100 };
1101 let handle = match Handle::try_current() {
1102 Ok(x) => x,
1103 Err(_) => {
1104 warn!(
1105 "WriteHandle {} dropped without being explicitly expired, falling back to lease timeout",
1106 self.writer_id
1107 );
1108 return;
1109 }
1110 };
1111 let expire_span = debug_span!("drop::expire");
1117 handle.spawn_named(
1118 || format!("WriteHandle::expire ({})", self.writer_id),
1119 expire_fn.0().instrument(expire_span),
1120 );
1121 }
1122}
1123
1124fn ensure_batch_schema<T>(batch: &mut HollowBatch<T>, shard_id: ShardId, schema_id: SchemaId)
1129where
1130 T: Timestamp + Lattice + Codec64,
1131{
1132 let ensure = |id: &mut Option<SchemaId>| match id {
1133 Some(id) => assert_eq!(*id, schema_id, "schema ID mismatch; shard={shard_id}"),
1134 None => *id = Some(schema_id),
1135 };
1136
1137 for run_meta in &mut batch.run_meta {
1138 ensure(&mut run_meta.schema);
1139 }
1140 for part in &mut batch.parts {
1141 match part {
1142 RunPart::Single(BatchPart::Hollow(part)) => ensure(&mut part.schema_id),
1143 RunPart::Single(BatchPart::Inline { schema_id, .. }) => ensure(schema_id),
1144 RunPart::Many(_hollow_run_ref) => {
1145 }
1149 }
1150 }
1151}
1152
1153#[cfg(test)]
1154mod tests {
1155 use std::str::FromStr;
1156 use std::sync::mpsc;
1157
1158 use differential_dataflow::consolidation::consolidate_updates;
1159 use futures_util::FutureExt;
1160 use mz_dyncfg::ConfigUpdates;
1161 use mz_ore::collections::CollectionExt;
1162 use mz_ore::task;
1163 use serde_json::json;
1164
1165 use crate::cache::PersistClientCache;
1166 use crate::tests::{all_ok, new_test_client};
1167 use crate::{PersistLocation, ShardId};
1168
1169 use super::*;
1170
1171 #[mz_persist_proc::test(tokio::test)]
1172 #[cfg_attr(miri, ignore)] async fn empty_batches(dyncfgs: ConfigUpdates) {
1174 let data = [
1175 (("1".to_owned(), "one".to_owned()), 1, 1),
1176 (("2".to_owned(), "two".to_owned()), 2, 1),
1177 (("3".to_owned(), "three".to_owned()), 3, 1),
1178 ];
1179
1180 let (mut write, _) = new_test_client(&dyncfgs)
1181 .await
1182 .expect_open::<String, String, u64, i64>(ShardId::new())
1183 .await;
1184 let blob = Arc::clone(&write.blob);
1185
1186 let mut upper = 3;
1188 write.expect_append(&data[..2], vec![0], vec![upper]).await;
1189
1190 let mut count_before = 0;
1192 blob.list_keys_and_metadata("", &mut |_| {
1193 count_before += 1;
1194 })
1195 .await
1196 .expect("list_keys failed");
1197 for _ in 0..5 {
1198 let new_upper = upper + 1;
1199 write.expect_compare_and_append(&[], upper, new_upper).await;
1200 upper = new_upper;
1201 }
1202 let mut count_after = 0;
1203 blob.list_keys_and_metadata("", &mut |_| {
1204 count_after += 1;
1205 })
1206 .await
1207 .expect("list_keys failed");
1208 assert_eq!(count_after, count_before);
1209 }
1210
1211 #[mz_persist_proc::test(tokio::test)]
1212 #[cfg_attr(miri, ignore)] async fn compare_and_append_batch_multi(dyncfgs: ConfigUpdates) {
1214 let data0 = vec![
1215 (("1".to_owned(), "one".to_owned()), 1, 1),
1216 (("2".to_owned(), "two".to_owned()), 2, 1),
1217 (("4".to_owned(), "four".to_owned()), 4, 1),
1218 ];
1219 let data1 = vec![
1220 (("1".to_owned(), "one".to_owned()), 1, 1),
1221 (("2".to_owned(), "two".to_owned()), 2, 1),
1222 (("3".to_owned(), "three".to_owned()), 3, 1),
1223 ];
1224
1225 let (mut write, mut read) = new_test_client(&dyncfgs)
1226 .await
1227 .expect_open::<String, String, u64, i64>(ShardId::new())
1228 .await;
1229
1230 let mut batch0 = write.expect_batch(&data0, 0, 5).await;
1231 let mut batch1 = write.expect_batch(&data1, 0, 4).await;
1232
1233 write
1234 .expect_compare_and_append_batch(&mut [&mut batch0, &mut batch1], 0, 4)
1235 .await;
1236
1237 let batch = write
1238 .machine
1239 .unleased_snapshot(&Antichain::from_elem(3))
1240 .await
1241 .expect("just wrote this")
1242 .into_element();
1243
1244 assert!(batch.runs().count() >= 2);
1245
1246 let expected = vec![
1247 (("1".to_owned(), "one".to_owned()), 1, 2),
1248 (("2".to_owned(), "two".to_owned()), 2, 2),
1249 (("3".to_owned(), "three".to_owned()), 3, 1),
1250 ];
1251 let mut actual = read.expect_snapshot_and_fetch(3).await;
1252 consolidate_updates(&mut actual);
1253 assert_eq!(actual, all_ok(&expected, 3));
1254 }
1255
1256 #[mz_ore::test]
1257 fn writer_id_human_readable_serde() {
1258 #[derive(Debug, Serialize, Deserialize)]
1259 struct Container {
1260 writer_id: WriterId,
1261 }
1262
1263 let id = WriterId::from_str("w00000000-1234-5678-0000-000000000000").expect("valid id");
1265 assert_eq!(
1266 id,
1267 serde_json::from_value(serde_json::to_value(id.clone()).expect("serializable"))
1268 .expect("deserializable")
1269 );
1270
1271 assert_eq!(
1273 id,
1274 serde_json::from_str("\"w00000000-1234-5678-0000-000000000000\"")
1275 .expect("deserializable")
1276 );
1277
1278 let json = json!({ "writer_id": id });
1280 assert_eq!(
1281 "{\"writer_id\":\"w00000000-1234-5678-0000-000000000000\"}",
1282 &json.to_string()
1283 );
1284 let container: Container = serde_json::from_value(json).expect("deserializable");
1285 assert_eq!(container.writer_id, id);
1286 }
1287
1288 #[mz_persist_proc::test(tokio::test)]
1289 #[cfg_attr(miri, ignore)] async fn hollow_batch_roundtrip(dyncfgs: ConfigUpdates) {
1291 let data = vec![
1292 (("1".to_owned(), "one".to_owned()), 1, 1),
1293 (("2".to_owned(), "two".to_owned()), 2, 1),
1294 (("3".to_owned(), "three".to_owned()), 3, 1),
1295 ];
1296
1297 let (mut write, mut read) = new_test_client(&dyncfgs)
1298 .await
1299 .expect_open::<String, String, u64, i64>(ShardId::new())
1300 .await;
1301
1302 let batch = write.expect_batch(&data, 0, 4).await;
1307 let hollow_batch = batch.into_transmittable_batch();
1308 let mut rehydrated_batch = write.batch_from_transmittable_batch(hollow_batch);
1309
1310 write
1311 .expect_compare_and_append_batch(&mut [&mut rehydrated_batch], 0, 4)
1312 .await;
1313
1314 let expected = vec![
1315 (("1".to_owned(), "one".to_owned()), 1, 1),
1316 (("2".to_owned(), "two".to_owned()), 2, 1),
1317 (("3".to_owned(), "three".to_owned()), 3, 1),
1318 ];
1319 let mut actual = read.expect_snapshot_and_fetch(3).await;
1320 consolidate_updates(&mut actual);
1321 assert_eq!(actual, all_ok(&expected, 3));
1322 }
1323
1324 #[mz_persist_proc::test(tokio::test)]
1325 #[cfg_attr(miri, ignore)] async fn wait_for_upper_past(dyncfgs: ConfigUpdates) {
1327 let client = new_test_client(&dyncfgs).await;
1328 let (mut write, _) = client.expect_open::<(), (), u64, i64>(ShardId::new()).await;
1329 let five = Antichain::from_elem(5);
1330
1331 assert_eq!(write.wait_for_upper_past(&five).now_or_never(), None);
1333
1334 write
1336 .expect_compare_and_append(&[(((), ()), 1, 1)], 0, 5)
1337 .await;
1338 assert_eq!(write.wait_for_upper_past(&five).now_or_never(), None);
1339
1340 write
1342 .expect_compare_and_append(&[(((), ()), 5, 1)], 5, 7)
1343 .await;
1344 assert_eq!(write.wait_for_upper_past(&five).now_or_never(), Some(()));
1345 assert_eq!(write.upper(), &Antichain::from_elem(7));
1346
1347 assert_eq!(
1350 write
1351 .wait_for_upper_past(&Antichain::from_elem(2))
1352 .now_or_never(),
1353 Some(())
1354 );
1355 assert_eq!(write.upper(), &Antichain::from_elem(7));
1356 }
1357
1358 #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
1359 #[cfg_attr(miri, ignore)] async fn fetch_recent_upper_linearized() {
1361 type Timestamp = u64;
1362 let max_upper = 1000;
1363
1364 let shard_id = ShardId::new();
1365 let mut clients = PersistClientCache::new_no_metrics();
1366 let upper_writer_client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
1367 let (mut upper_writer, _) = upper_writer_client
1368 .expect_open::<(), (), Timestamp, i64>(shard_id)
1369 .await;
1370 clients.clear_state_cache();
1373 let upper_reader_client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
1374 let (mut upper_reader, _) = upper_reader_client
1375 .expect_open::<(), (), Timestamp, i64>(shard_id)
1376 .await;
1377 let (tx, rx) = mpsc::channel();
1378
1379 let task = task::spawn(|| "upper-reader", async move {
1380 let mut upper = Timestamp::MIN;
1381
1382 while upper < max_upper {
1383 while let Ok(new_upper) = rx.try_recv() {
1384 upper = new_upper;
1385 }
1386
1387 let recent_upper = upper_reader
1388 .fetch_recent_upper()
1389 .await
1390 .as_option()
1391 .cloned()
1392 .expect("u64 is totally ordered and the shard is not finalized");
1393 assert!(
1394 recent_upper >= upper,
1395 "recent upper {recent_upper:?} is less than known upper {upper:?}"
1396 );
1397 }
1398 });
1399
1400 for upper in Timestamp::MIN..max_upper {
1401 let next_upper = upper + 1;
1402 upper_writer
1403 .expect_compare_and_append(&[], upper, next_upper)
1404 .await;
1405 tx.send(next_upper).expect("send failed");
1406 }
1407
1408 task.await;
1409 }
1410}