1use std::collections::BTreeMap;
13use std::fmt::Debug;
14use std::ops::ControlFlow::{self, Break, Continue};
15use std::sync::Arc;
16use std::time::Instant;
17
18use crate::cache::{LockingTypedState, StateCache};
19use crate::error::{CodecMismatch, InvalidUsage};
20use crate::internal::gc::GcReq;
21use crate::internal::maintenance::RoutineMaintenance;
22use crate::internal::metrics::{CmdMetrics, Metrics, ShardMetrics};
23use crate::internal::paths::{PartialRollupKey, RollupId};
24use crate::internal::state::{
25 ActiveGc, ActiveRollup, EncodedSchemas, ExpiryMetrics, GC_FALLBACK_THRESHOLD_MS,
26 GC_MAX_VERSIONS, GC_MIN_VERSIONS, GC_USE_ACTIVE_GC, GcConfig, HollowBatch, LeasedReaderState,
27 ROLLUP_FALLBACK_THRESHOLD_MS, ROLLUP_THRESHOLD, ROLLUP_USE_ACTIVE_ROLLUP, Since, SnapshotErr,
28 StateCollections, TypedState,
29};
30use crate::internal::state_diff::StateDiff;
31use crate::internal::state_versions::{EncodedRollup, StateVersions};
32use crate::internal::trace::FueledMergeReq;
33use crate::internal::watch::StateWatch;
34use crate::read::LeasedReaderId;
35use crate::rpc::{PUBSUB_PUSH_DIFF_ENABLED, PubSubSender};
36use crate::schema::SchemaCache;
37use crate::{Diagnostics, PersistConfig, ShardId, cfg};
38use differential_dataflow::difference::Monoid;
39use differential_dataflow::lattice::Lattice;
40use mz_ore::cast::CastFrom;
41use mz_ore::soft_assert_or_log;
42use mz_persist::location::{CaSResult, Indeterminate, SeqNo, VersionedData};
43use mz_persist_types::schema::SchemaId;
44use mz_persist_types::{Codec, Codec64};
45use timely::progress::{Antichain, Timestamp};
46use tracing::debug;
47
48#[derive(Debug)]
53pub struct Applier<K, V, T, D> {
54 pub(crate) cfg: PersistConfig,
55 pub(crate) metrics: Arc<Metrics>,
56 pub(crate) shard_metrics: Arc<ShardMetrics>,
57 pub(crate) state_versions: Arc<StateVersions>,
58 pubsub_sender: Arc<dyn PubSubSender>,
59 pub(crate) shard_id: ShardId,
60
61 state: Arc<LockingTypedState<K, V, T, D>>,
70}
71
72impl<K, V, T: Clone, D> Clone for Applier<K, V, T, D> {
74 fn clone(&self) -> Self {
75 Self {
76 cfg: self.cfg.clone(),
77 metrics: Arc::clone(&self.metrics),
78 shard_metrics: Arc::clone(&self.shard_metrics),
79 state_versions: Arc::clone(&self.state_versions),
80 pubsub_sender: Arc::clone(&self.pubsub_sender),
81 shard_id: self.shard_id,
82 state: Arc::clone(&self.state),
83 }
84 }
85}
86
87impl<K, V, T, D> Applier<K, V, T, D>
88where
89 K: Debug + Codec,
90 V: Debug + Codec,
91 T: Timestamp + Lattice + Codec64 + Sync,
92 D: Monoid + Codec64,
93{
94 pub async fn new(
95 cfg: PersistConfig,
96 shard_id: ShardId,
97 metrics: Arc<Metrics>,
98 state_versions: Arc<StateVersions>,
99 shared_states: Arc<StateCache>,
100 pubsub_sender: Arc<dyn PubSubSender>,
101 diagnostics: Diagnostics,
102 ) -> Result<Self, Box<CodecMismatch>> {
103 let shard_metrics = metrics.shards.shard(&shard_id, &diagnostics.shard_name);
104 let state = shared_states
105 .get::<K, V, T, D, _, _>(
106 shard_id,
107 || {
108 metrics.cmds.init_state.run_cmd(&shard_metrics, || {
109 state_versions.maybe_init_shard(&shard_metrics)
110 })
111 },
112 &diagnostics,
113 )
114 .await?;
115 let ret = Applier {
116 cfg,
117 metrics,
118 shard_metrics,
119 state_versions,
120 pubsub_sender,
121 shard_id,
122 state,
123 };
124 Ok(ret)
125 }
126
127 pub fn watch(&self) -> StateWatch<K, V, T, D> {
129 StateWatch::new(Arc::clone(&self.state), Arc::clone(&self.metrics))
130 }
131
132 pub async fn fetch_upper<R, F: FnMut(&Antichain<T>) -> R>(&self, mut f: F) -> R {
134 self.metrics.cmds.fetch_upper_count.inc();
135 self.fetch_and_update_state(None).await;
136 self.upper(|_seqno, upper| f(upper))
137 }
138
139 pub fn clone_upper(&self) -> Antichain<T> {
146 self.upper(|_seqno, upper| upper.clone())
147 }
148
149 pub(crate) fn upper<R, F: FnMut(SeqNo, &Antichain<T>) -> R>(&self, mut f: F) -> R {
150 self.state
151 .read_lock(&self.metrics.locks.applier_read_cacheable, move |state| {
152 f(state.seqno, state.upper())
153 })
154 }
155
156 pub(crate) fn schemas<R>(
157 &self,
158 mut f: impl FnMut(SeqNo, &BTreeMap<SchemaId, EncodedSchemas>) -> R,
159 ) -> R {
160 self.state
161 .read_lock(&self.metrics.locks.applier_read_cacheable, move |state| {
162 f(state.seqno, &state.collections.schemas)
163 })
164 }
165
166 pub(crate) fn schema_cache(&self) -> SchemaCache<K, V, T, D> {
167 SchemaCache::new(self.state.schema_cache(), self.clone())
168 }
169
170 #[cfg(test)]
177 pub fn since(&self) -> Antichain<T> {
178 self.state
179 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
180 state.since().clone()
181 })
182 }
183
184 pub fn reader_lease(&self, id: LeasedReaderId) -> Option<LeasedReaderState<T>> {
191 self.state
192 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
193 state.state.collections.leased_readers.get(&id).cloned()
194 })
195 }
196
197 pub fn seqno(&self) -> SeqNo {
204 self.state
205 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
206 state.seqno
207 })
208 }
209
210 pub fn seqno_since(&self) -> SeqNo {
217 self.state
218 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
219 state.seqno_since()
220 })
221 }
222
223 pub fn is_finalized(&self) -> bool {
230 self.state
231 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
232 state.collections.is_tombstone() && state.collections.is_single_empty_batch()
233 })
234 }
235
236 pub(crate) fn is_tombstone(&self) -> bool {
238 self.state
239 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
240 state.collections.is_tombstone()
241 })
242 }
243
244 pub fn get_schema(&self, schema_id: SchemaId) -> Option<(K::Schema, V::Schema)> {
246 self.state
247 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
248 let x = state.collections.schemas.get(&schema_id)?;
249 Some((K::decode_schema(&x.key), V::decode_schema(&x.val)))
250 })
251 }
252
253 pub fn latest_schema(&self) -> Option<(SchemaId, K::Schema, V::Schema)> {
255 self.state
256 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
257 let (id, x) = state.collections.schemas.last_key_value()?;
258 Some((*id, K::decode_schema(&x.key), V::decode_schema(&x.val)))
259 })
260 }
261
262 pub fn find_schema(&self, key_schema: &K::Schema, val_schema: &V::Schema) -> Option<SchemaId> {
264 self.state
265 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
266 let mut schemas = state.collections.schemas.iter().rev();
269 schemas
270 .find(|(_, x)| {
271 K::decode_schema(&x.key) == *key_schema
272 && V::decode_schema(&x.val) == *val_schema
273 })
274 .map(|(id, _)| *id)
275 })
276 }
277
278 pub fn check_since_upper_both_empty(&self) -> Result<(), InvalidUsage<T>> {
284 self.state
285 .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
286 if state.since().is_empty() && state.upper().is_empty() {
287 Ok(())
288 } else {
289 Err(InvalidUsage::FinalizationError {
290 since: state.since().clone(),
291 upper: state.upper().clone(),
292 })
293 }
294 })
295 }
296
297 pub fn rollups_lte_seqno(&self, seqno: SeqNo) -> Vec<(SeqNo, PartialRollupKey)> {
303 self.state
304 .read_lock(&self.metrics.locks.applier_read_noncacheable, |state| {
305 state
306 .collections
307 .rollups
308 .range(..=seqno)
309 .map(|(seqno, rollup)| (*seqno, rollup.key.clone()))
310 .collect::<Vec<(SeqNo, PartialRollupKey)>>()
311 })
312 }
313
314 pub fn all_fueled_merge_reqs(&self) -> Vec<FueledMergeReq<T>> {
315 self.state
316 .read_lock(&self.metrics.locks.applier_read_noncacheable, |state| {
317 state
318 .collections
319 .trace
320 .fueled_merge_reqs_before_ms(u64::MAX, None)
321 .collect()
322 })
323 }
324
325 pub fn snapshot(&self, as_of: &Antichain<T>) -> Result<Vec<HollowBatch<T>>, SnapshotErr<T>> {
326 self.state
327 .read_lock(&self.metrics.locks.applier_read_noncacheable, |state| {
328 state.snapshot(as_of)
329 })
330 }
331
332 pub fn all_batches(&self) -> Vec<HollowBatch<T>> {
333 self.state
334 .read_lock(&self.metrics.locks.applier_read_noncacheable, |state| {
335 state.state.collections.trace.batches().cloned().collect()
336 })
337 }
338
339 pub fn verify_listen(&self, as_of: &Antichain<T>) -> Result<(), Since<T>> {
340 self.state
341 .read_lock(&self.metrics.locks.applier_read_noncacheable, |state| {
342 state.verify_listen(as_of)
343 })
344 }
345
346 pub fn next_listen_batch(&self, frontier: &Antichain<T>) -> Result<HollowBatch<T>, SeqNo> {
347 self.state
348 .read_lock(&self.metrics.locks.applier_read_noncacheable, |state| {
349 state.next_listen_batch(frontier)
350 })
351 }
352
353 pub async fn write_rollup_for_state(&self) -> Option<EncodedRollup> {
354 let state = self
355 .state
356 .read_lock(&self.metrics.locks.applier_read_noncacheable, |state| {
357 state.clone_for_rollup()
358 });
359
360 self.state_versions
361 .write_rollup_for_state(self.shard_metrics.as_ref(), state, &RollupId::new())
362 .await
363 }
364
365 pub async fn apply_unbatched_cmd<
366 R,
367 E,
368 WorkFn: FnMut(SeqNo, &PersistConfig, &mut StateCollections<T>) -> ControlFlow<E, R>,
369 >(
370 &self,
371 cmd: &CmdMetrics,
372 mut work_fn: WorkFn,
373 ) -> Result<(SeqNo, Result<R, E>, RoutineMaintenance), Indeterminate> {
374 loop {
375 cmd.started.inc();
376 let now = Instant::now();
377 let ret = Self::apply_unbatched_cmd_locked(
378 &self.state,
379 cmd,
380 &mut work_fn,
381 &self.cfg,
382 &self.metrics,
383 &self.shard_metrics,
384 &self.state_versions,
385 )
386 .await;
387 cmd.seconds.inc_by(now.elapsed().as_secs_f64());
388
389 match ret {
390 ApplyCmdResult::Committed((diff, new_state, res, maintenance)) => {
391 cmd.succeeded.inc();
392 self.shard_metrics.cmd_succeeded.inc();
393 self.update_state(new_state);
394 if PUBSUB_PUSH_DIFF_ENABLED.get(&self.cfg) {
395 self.pubsub_sender.push_diff(&self.shard_id, &diff);
396 }
397 return Ok((diff.seqno, Ok(res), maintenance));
398 }
399 ApplyCmdResult::SkippedStateTransition((seqno, err, maintenance)) => {
400 cmd.succeeded.inc();
401 self.shard_metrics.cmd_succeeded.inc();
402 return Ok((seqno, Err(err), maintenance));
403 }
404 ApplyCmdResult::Indeterminate(err) => {
405 cmd.failed.inc();
406 return Err(err);
407 }
408 ApplyCmdResult::ExpectationMismatch(seqno) => {
409 cmd.cas_mismatch.inc();
410 self.fetch_and_update_state(Some(seqno)).await;
411 }
412 }
413 }
414 }
415
416 #[allow(clippy::needless_pass_by_ref_mut)]
418 async fn apply_unbatched_cmd_locked<
419 R,
420 E,
421 WorkFn: FnMut(SeqNo, &PersistConfig, &mut StateCollections<T>) -> ControlFlow<E, R>,
422 >(
423 state: &LockingTypedState<K, V, T, D>,
424 cmd: &CmdMetrics,
425 work_fn: &mut WorkFn,
426 cfg: &PersistConfig,
427 metrics: &Metrics,
428 shard_metrics: &ShardMetrics,
429 state_versions: &StateVersions,
430 ) -> ApplyCmdResult<K, V, T, D, R, E> {
431 let _permit_opt = state.lease_for_update().await;
436
437 let computed_next_state = state
438 .read_lock(&metrics.locks.applier_read_noncacheable, |state| {
439 Self::compute_next_state_locked(state, work_fn, metrics, cmd, cfg)
440 });
441
442 let next_state = match computed_next_state {
443 Ok(x) => x,
444 Err((seqno, err)) => {
445 return ApplyCmdResult::SkippedStateTransition((
446 seqno,
447 err,
448 RoutineMaintenance::default(),
449 ));
450 }
451 };
452
453 let NextState {
454 diff,
455 state,
456 expiry_metrics,
457 garbage_collection,
458 write_rollup,
459 work_ret,
460 } = next_state;
461
462 {
463 let build_version = &cfg.build_version;
464 let state_version = &state.state.collections.version;
465 soft_assert_or_log!(
466 cfg::code_can_write_data(build_version, state_version),
467 "current version {build_version} does not support state format {state_version}"
468 );
469 }
470
471 let cas_res = state_versions
477 .try_compare_and_set_current(&cmd.name, shard_metrics, &state, &diff)
478 .await;
479
480 match cas_res {
481 Ok((CaSResult::Committed, diff)) => {
482 metrics
483 .lease
484 .timeout_read
485 .inc_by(u64::cast_from(expiry_metrics.readers_expired));
486
487 metrics
488 .state
489 .writer_removed
490 .inc_by(u64::cast_from(expiry_metrics.writers_expired));
491
492 if let Some(gc) = garbage_collection.as_ref() {
493 debug!("Assigned gc request: {:?}", gc);
494 }
495
496 let maintenance = RoutineMaintenance {
497 garbage_collection,
498 write_rollup,
499 };
500
501 ApplyCmdResult::Committed((diff, state, work_ret, maintenance))
502 }
503 Ok((CaSResult::ExpectationMismatch, _diff)) => {
504 ApplyCmdResult::ExpectationMismatch(state.seqno())
505 }
506 Err(err) => ApplyCmdResult::Indeterminate(err),
507 }
508 }
509
510 fn compute_next_state_locked<
511 R,
512 E,
513 WorkFn: FnMut(SeqNo, &PersistConfig, &mut StateCollections<T>) -> ControlFlow<E, R>,
514 >(
515 state: &TypedState<K, V, T, D>,
516 work_fn: &mut WorkFn,
517 metrics: &Metrics,
518 cmd: &CmdMetrics,
519 cfg: &PersistConfig,
520 ) -> Result<NextState<K, V, T, D, R>, (SeqNo, E)> {
521 let is_write = cmd.name == metrics.cmds.compare_and_append.name;
522 let is_rollup = cmd.name == metrics.cmds.add_rollup.name;
523 let is_become_tombstone = cmd.name == metrics.cmds.become_tombstone.name;
524
525 let gc_config = GcConfig {
526 use_active_gc: GC_USE_ACTIVE_GC.get(cfg),
527 fallback_threshold_ms: u64::cast_from(GC_FALLBACK_THRESHOLD_MS.get(cfg)),
528 min_versions: GC_MIN_VERSIONS.get(cfg),
529 max_versions: GC_MAX_VERSIONS.get(cfg),
530 };
531
532 let use_active_rollup = ROLLUP_USE_ACTIVE_ROLLUP.get(cfg);
533 let rollup_threshold = ROLLUP_THRESHOLD.get(cfg);
534 let rollup_fallback_threshold_ms = u64::cast_from(ROLLUP_FALLBACK_THRESHOLD_MS.get(cfg));
535
536 let expected = state.seqno;
537 let was_tombstone_before = state.collections.is_tombstone();
538
539 let (work_ret, mut new_state) = match state.clone_apply(cfg, work_fn) {
540 Continue(x) => x,
541 Break(err) => {
542 return Err((expected, err));
543 }
544 };
545 let expiry_metrics = new_state.expire_at((cfg.now)());
546 new_state.state.collections.trace.roundtrip_structure = true;
547
548 if was_tombstone_before && !(is_rollup || is_become_tombstone) {
558 panic!(
559 "cmd {} unexpectedly tried to commit a new state on a tombstone: {:?}",
560 cmd.name, state
561 );
562 }
563
564 let now = (cfg.now)();
565 let write_rollup = new_state.need_rollup(
566 rollup_threshold,
567 use_active_rollup,
568 rollup_fallback_threshold_ms,
569 now,
570 );
571
572 if let Some(write_rollup_seqno) = write_rollup {
573 if use_active_rollup {
574 new_state.collections.active_rollup = Some(ActiveRollup {
575 seqno: write_rollup_seqno,
576 start_ms: now,
577 });
578 }
579 }
580
581 let garbage_collection = new_state.maybe_gc(is_write, now, gc_config);
587
588 if let Some(gc) = garbage_collection.as_ref() {
589 if gc_config.use_active_gc {
590 new_state.collections.active_gc = Some(ActiveGc {
591 seqno: gc.new_seqno_since,
592 start_ms: now,
593 });
594 }
595 }
596
597 let diff = StateDiff::from_diff(&state.state, &new_state);
600 #[cfg(any(test, debug_assertions))]
603 {
604 if let Err(err) = StateDiff::validate_roundtrip(metrics, state, &diff, &new_state) {
605 panic!("validate_roundtrips failed: {}", err);
606 }
607 }
608
609 assert_eq!(
610 expected.next(),
611 new_state.seqno(),
612 "successive states should have successive seqnos"
613 );
614 Ok(NextState {
615 diff,
616 state: new_state,
617 expiry_metrics,
618 garbage_collection,
619 write_rollup,
620 work_ret,
621 })
622 }
623
624 pub fn update_state(&self, new_state: TypedState<K, V, T, D>) {
625 let (seqno_before, seqno_after) =
626 self.state
627 .write_lock(&self.metrics.locks.applier_write, |state| {
628 let seqno_before = state.seqno;
629 if seqno_before < new_state.seqno {
630 *state = new_state;
631 }
632 let seqno_after = state.seqno;
633 (seqno_before, seqno_after)
634 });
635
636 assert!(
637 seqno_before <= seqno_after,
638 "state seqno regressed: {} vs {}",
639 seqno_before,
640 seqno_after
641 );
642 }
643
644 pub async fn fetch_and_update_state(&self, seqno_hint: Option<SeqNo>) {
648 let current_seqno = self.seqno();
649 let seqno_before = match seqno_hint {
650 None => current_seqno,
651 Some(hint) => {
652 if hint < current_seqno {
655 self.metrics.state.update_state_noop_path.inc();
656 return;
657 }
658 current_seqno
659 }
660 };
661
662 let diffs_to_current = self
663 .state_versions
664 .fetch_all_live_diffs_gt_seqno::<K, V, T, D>(&self.shard_id, seqno_before)
665 .await;
666
667 if diffs_to_current.is_empty() {
669 self.metrics.state.update_state_empty_path.inc();
670 return;
671 }
672
673 let new_seqno = self
674 .state
675 .write_lock(&self.metrics.locks.applier_write, |state| {
676 state.apply_encoded_diffs(&self.cfg, &self.metrics, &diffs_to_current);
677 state.seqno
678 });
679
680 assert!(
681 seqno_before <= new_seqno,
682 "state seqno regressed: {} vs {}",
683 seqno_before,
684 new_seqno
685 );
686
687 if seqno_before < new_seqno {
690 self.metrics.state.update_state_fast_path.inc();
691 return;
692 }
693
694 let new_state = self
698 .state_versions
699 .fetch_current_state(&self.shard_id, diffs_to_current)
700 .await
701 .check_codecs::<K, V, D>(&self.shard_id)
702 .expect("shard codecs should not change");
703
704 let new_seqno = self
705 .state
706 .write_lock(&self.metrics.locks.applier_write, |state| {
707 if state.seqno < new_state.seqno {
708 *state = new_state;
709 }
710 state.seqno
711 });
712
713 self.metrics.state.update_state_slow_path.inc();
714 assert!(
715 seqno_before <= new_seqno,
716 "state seqno regressed: {} vs {}",
717 seqno_before,
718 new_seqno
719 );
720 }
721}
722
723enum ApplyCmdResult<K, V, T, D, R, E> {
724 Committed((VersionedData, TypedState<K, V, T, D>, R, RoutineMaintenance)),
725 SkippedStateTransition((SeqNo, E, RoutineMaintenance)),
726 Indeterminate(Indeterminate),
727 ExpectationMismatch(SeqNo),
728}
729
730struct NextState<K, V, T, D, R> {
731 diff: StateDiff<T>,
732 state: TypedState<K, V, T, D>,
733 expiry_metrics: ExpiryMetrics,
734 write_rollup: Option<SeqNo>,
735 garbage_collection: Option<GcReq>,
736 work_ret: R,
737}