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