1use std::collections::{BTreeMap, BTreeSet};
13use std::future::Future;
14use std::sync::Arc;
15
16use bytes::Bytes;
17use mz_compute_types::dataflows::DataflowDescription;
18use mz_durable_cache::{DurableCache, DurableCacheCodec};
19use mz_dyncfg::ConfigSet;
20use mz_expr::OptimizedMirRelationExpr;
21use mz_ore::channel::trigger;
22use mz_ore::soft_panic_or_log;
23use mz_ore::task::spawn;
24use mz_persist_client::PersistClient;
25use mz_persist_client::cli::admin::{
26 EXPRESSION_CACHE_FORCE_COMPACTION_FUEL, EXPRESSION_CACHE_FORCE_COMPACTION_WAIT,
27};
28use mz_persist_types::codec_impls::VecU8Schema;
29use mz_persist_types::{Codec, ShardId};
30use mz_repr::optimize::OptimizerFeatures;
31use mz_repr::{GlobalId, RelationVersion};
32use mz_transform::dataflow::DataflowMetainfo;
33use mz_transform::notice::OptimizerNotice;
34use semver::Version;
35use serde::{Deserialize, Serialize};
36use tokio::sync::mpsc;
37use tracing::{debug, warn};
38
39#[derive(
40 Debug,
41 Clone,
42 PartialEq,
43 Eq,
44 PartialOrd,
45 Ord,
46 Hash,
47 Serialize,
48 Deserialize
49)]
50enum ExpressionType {
51 Local,
52 Global,
53}
54
55#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
57pub struct LocalExpressions {
58 pub local_mir: OptimizedMirRelationExpr,
59 pub optimizer_features: OptimizerFeatures,
60 pub item_version: RelationVersion,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
66pub struct GlobalExpressions {
67 pub global_mir: DataflowDescription<OptimizedMirRelationExpr>,
68 pub physical_plan: DataflowDescription<mz_compute_types::plan::LirRelationExpr>,
69 pub dataflow_metainfos: DataflowMetainfo<Arc<OptimizerNotice>>,
70 pub optimizer_features: OptimizerFeatures,
71 pub item_version: RelationVersion,
73}
74
75impl GlobalExpressions {
76 fn index_imports(&self) -> impl Iterator<Item = &GlobalId> {
77 self.global_mir
78 .index_imports
79 .keys()
80 .chain(self.physical_plan.index_imports.keys())
81 }
82}
83
84pub fn latest_item_version(versions: &BTreeMap<RelationVersion, GlobalId>) -> RelationVersion {
86 versions
87 .last_key_value()
88 .map_or_else(RelationVersion::root, |(version, _)| *version)
89}
90
91#[derive(
92 Debug,
93 Clone,
94 PartialEq,
95 Eq,
96 PartialOrd,
97 Ord,
98 Hash,
99 Serialize,
100 Deserialize
101)]
102struct CacheKey {
103 build_version: String,
104 id: GlobalId,
105 expr_type: ExpressionType,
106}
107
108#[derive(Debug, PartialEq, Eq)]
109struct ExpressionCodec;
110
111impl DurableCacheCodec for ExpressionCodec {
112 type Key = CacheKey;
113 type Val = Bytes;
116 type KeyCodec = Bytes;
117 type ValCodec = Bytes;
118
119 fn schemas() -> (
120 <Self::KeyCodec as Codec>::Schema,
121 <Self::ValCodec as Codec>::Schema,
122 ) {
123 (VecU8Schema::default(), VecU8Schema::default())
124 }
125
126 fn encode(key: &Self::Key, val: &Self::Val) -> (Self::KeyCodec, Self::ValCodec) {
127 let key = bincode::serialize(key).expect("must serialize");
128 (Bytes::from(key), val.clone())
129 }
130
131 fn decode(key: &Self::KeyCodec, val: &Self::ValCodec) -> (Self::Key, Self::Val) {
132 let key = bincode::deserialize(key).expect("must deserialize");
133 (key, val.clone())
134 }
135}
136
137#[derive(Debug, Clone)]
139pub struct ExpressionCacheConfig {
140 pub build_version: Version,
141 pub persist: PersistClient,
142 pub shard_id: ShardId,
143 pub current_items: BTreeMap<GlobalId, RelationVersion>,
147 pub remove_prior_versions: bool,
149 pub compact_shard: bool,
150 pub dyncfgs: ConfigSet,
151}
152
153pub struct ExpressionCache {
155 build_version: Version,
156 durable_cache: DurableCache<ExpressionCodec>,
157}
158
159impl ExpressionCache {
160 pub async fn open(
185 ExpressionCacheConfig {
186 build_version,
187 persist,
188 shard_id,
189 current_items,
190 remove_prior_versions,
191 compact_shard,
192 dyncfgs,
193 }: ExpressionCacheConfig,
194 ) -> (
195 Self,
196 BTreeMap<GlobalId, LocalExpressions>,
197 BTreeMap<GlobalId, GlobalExpressions>,
198 ) {
199 let durable_cache = DurableCache::new(&persist, shard_id, "expressions").await;
200 let mut cache = Self {
201 build_version,
202 durable_cache,
203 };
204
205 const RETRIES: usize = 100;
206 for _ in 0..RETRIES {
207 match cache
208 .try_open(
209 ¤t_items,
210 remove_prior_versions,
211 compact_shard,
212 &dyncfgs,
213 )
214 .await
215 {
216 Ok((local_expressions, global_expressions)) => {
217 return (cache, local_expressions, global_expressions);
218 }
219 Err(err) => debug!("failed to open cache: {err} ... retrying"),
220 }
221 }
222
223 panic!("Unable to open expression cache after {RETRIES} retries");
224 }
225
226 async fn try_open(
227 &mut self,
228 current_items: &BTreeMap<GlobalId, RelationVersion>,
229 remove_prior_versions: bool,
230 compact_shard: bool,
231 dyncfgs: &ConfigSet,
232 ) -> Result<
233 (
234 BTreeMap<GlobalId, LocalExpressions>,
235 BTreeMap<GlobalId, GlobalExpressions>,
236 ),
237 mz_durable_cache::Error,
238 > {
239 let mut keys_to_remove = Vec::new();
240 let mut local_expressions = BTreeMap::new();
241 let mut global_expressions = BTreeMap::new();
242
243 for (key, expressions) in self.durable_cache.entries_local() {
244 let build_version = match key.build_version.parse::<Version>() {
245 Ok(build_version) => build_version,
246 Err(err) => {
247 warn!("unable to parse build version: {key:?}: {err:?}");
248 keys_to_remove.push((key.clone(), None));
249 continue;
250 }
251 };
252 if build_version == self.build_version {
253 match key.expr_type {
255 ExpressionType::Local => {
256 let expressions: LocalExpressions = match bincode::deserialize(expressions)
257 {
258 Ok(expressions) => expressions,
259 Err(err) => {
260 soft_panic_or_log!(
261 "unable to deserialize local expressions: ({key:?}, {expressions:?}): {err:?}"
262 );
263 continue;
264 }
265 };
266 if current_items.get(&key.id) != Some(&expressions.item_version) {
269 keys_to_remove.push((key.clone(), None));
270 } else {
271 local_expressions.insert(key.id, expressions);
272 }
273 }
274 ExpressionType::Global => {
275 let expressions: GlobalExpressions = match bincode::deserialize(expressions)
276 {
277 Ok(expressions) => expressions,
278 Err(err) => {
279 soft_panic_or_log!(
280 "unable to deserialize global expressions: ({key:?}, {expressions:?}): {err:?}"
281 );
282 continue;
283 }
284 };
285 if current_items.get(&key.id) != Some(&expressions.item_version)
289 || !expressions
290 .index_imports()
291 .all(|id| current_items.contains_key(id))
292 {
293 keys_to_remove.push((key.clone(), None));
294 } else {
295 global_expressions.insert(key.id, expressions);
296 }
297 }
298 }
299 } else if remove_prior_versions {
300 keys_to_remove.push((key.clone(), None));
302 }
303 }
304 let keys_to_remove: Vec<_> = keys_to_remove
305 .iter()
306 .map(|(key, expressions)| (key, expressions.as_ref()))
307 .collect();
308 self.durable_cache.try_set_many(&keys_to_remove).await?;
309
310 if remove_prior_versions {
311 self.durable_cache.upgrade_version().await;
314 }
315
316 if compact_shard {
317 let fuel = EXPRESSION_CACHE_FORCE_COMPACTION_FUEL.handle(dyncfgs);
318 let wait = EXPRESSION_CACHE_FORCE_COMPACTION_WAIT.handle(dyncfgs);
319 self.durable_cache
320 .dangerous_compact_shard(move || fuel.get(), move || wait.get())
321 .await;
322 }
323
324 Ok((local_expressions, global_expressions))
325 }
326
327 async fn update(
333 &mut self,
334 new_local_expressions: Vec<(GlobalId, LocalExpressions)>,
335 new_global_expressions: Vec<(GlobalId, GlobalExpressions)>,
336 invalidate_ids: BTreeSet<GlobalId>,
337 ) {
338 let mut entries = BTreeMap::new();
339 let build_version = self.build_version.to_string();
340 for id in invalidate_ids {
343 entries.insert(
344 CacheKey {
345 id,
346 build_version: build_version.clone(),
347 expr_type: ExpressionType::Local,
348 },
349 None,
350 );
351 entries.insert(
352 CacheKey {
353 id,
354 build_version: build_version.clone(),
355 expr_type: ExpressionType::Global,
356 },
357 None,
358 );
359 }
360 for (id, expressions) in new_local_expressions {
361 let expressions = match bincode::serialize(&expressions) {
362 Ok(expressions) => Bytes::from(expressions),
363 Err(err) => {
364 soft_panic_or_log!(
365 "unable to serialize local expressions: {expressions:?}: {err:?}"
366 );
367 continue;
368 }
369 };
370 entries.insert(
371 CacheKey {
372 id,
373 build_version: build_version.clone(),
374 expr_type: ExpressionType::Local,
375 },
376 Some(expressions),
377 );
378 }
379 for (id, expressions) in new_global_expressions {
380 let expressions = match bincode::serialize(&expressions) {
381 Ok(expressions) => Bytes::from(expressions),
382 Err(err) => {
383 soft_panic_or_log!(
384 "unable to serialize global expressions: {expressions:?}: {err:?}"
385 );
386 continue;
387 }
388 };
389 entries.insert(
390 CacheKey {
391 id,
392 build_version: build_version.clone(),
393 expr_type: ExpressionType::Global,
394 },
395 Some(expressions),
396 );
397 }
398 let entries: Vec<_> = entries
399 .iter()
400 .map(|(key, expressions)| (key, expressions.as_ref()))
401 .collect();
402 self.durable_cache.set_many(&entries).await
403 }
404}
405
406enum CacheOperation {
408 Update {
410 new_local_expressions: Vec<(GlobalId, LocalExpressions)>,
411 new_global_expressions: Vec<(GlobalId, GlobalExpressions)>,
412 invalidate_ids: BTreeSet<GlobalId>,
413 trigger: trigger::Trigger,
414 },
415}
416
417#[derive(Debug, Clone)]
418pub struct ExpressionCacheHandle {
419 tx: mpsc::UnboundedSender<CacheOperation>,
420}
421
422impl ExpressionCacheHandle {
423 pub async fn spawn_expression_cache(
427 config: ExpressionCacheConfig,
428 ) -> (
429 Self,
430 BTreeMap<GlobalId, LocalExpressions>,
431 BTreeMap<GlobalId, GlobalExpressions>,
432 ) {
433 let (mut cache, local_expressions, global_expressions) =
434 ExpressionCache::open(config).await;
435 let (tx, mut rx) = mpsc::unbounded_channel();
436 spawn(|| "expression-cache-task", async move {
437 while let Some(op) = rx.recv().await {
438 match op {
439 CacheOperation::Update {
440 new_local_expressions,
441 new_global_expressions,
442 invalidate_ids,
443 trigger: _trigger,
444 } => {
445 cache
446 .update(
447 new_local_expressions,
448 new_global_expressions,
449 invalidate_ids,
450 )
451 .await
452 }
453 }
454 }
455 });
456
457 (Self { tx }, local_expressions, global_expressions)
458 }
459
460 pub fn update(
461 &self,
462 new_local_expressions: Vec<(GlobalId, LocalExpressions)>,
463 new_global_expressions: Vec<(GlobalId, GlobalExpressions)>,
464 invalidate_ids: BTreeSet<GlobalId>,
465 ) -> impl Future<Output = ()> + use<> {
466 let (trigger, trigger_rx) = trigger::channel();
467 let op = CacheOperation::Update {
468 new_local_expressions,
469 new_global_expressions,
470 invalidate_ids,
471 trigger,
472 };
473 let _ = self.tx.send(op);
475 trigger_rx
476 }
477}
478
479#[cfg(test)]
480mod tests {
481 use std::collections::{BTreeMap, BTreeSet};
482
483 use bytes::Bytes;
484 use mz_compute_types::dataflows::{DataflowDescription, IndexDesc, IndexImport};
485 use mz_durable_cache::DurableCacheCodec;
486 use mz_dyncfg::ConfigSet;
487 use mz_expr::{MirRelationExpr, OptimizedMirRelationExpr};
488 use mz_persist_client::PersistClient;
489 use mz_persist_types::ShardId;
490 use mz_repr::{Datum, GlobalId, ReprRelationType, ReprScalarType};
491 use semver::Version;
492
493 use super::*;
494
495 #[mz_ore::test(tokio::test)]
496 #[cfg_attr(miri, ignore)] async fn expression_cache() {
498 let first_build_version = Version::new(0, 1, 0);
499 let second_build_version = Version::new(0, 2, 0);
500 let persist = PersistClient::new_for_tests().await;
501 let shard_id = ShardId::new();
502
503 let mut current_items = BTreeMap::new();
504 let mut remove_prior_versions = false;
505 let compact_shard = false;
507 let dyncfgs = &mz_persist_client::cfg::all_dyncfgs(ConfigSet::default());
508 let spawn = |build_version: &Version,
509 current_items: &BTreeMap<GlobalId, RelationVersion>,
510 remove_prior_versions: bool| {
511 ExpressionCacheHandle::spawn_expression_cache(ExpressionCacheConfig {
512 build_version: build_version.clone(),
513 persist: persist.clone(),
514 shard_id,
515 current_items: current_items.clone(),
516 remove_prior_versions,
517 compact_shard,
518 dyncfgs: dyncfgs.clone(),
519 })
520 };
521
522 let mut next_id = 0;
523
524 let (mut local_exps, mut global_exps) = {
525 let (cache, local_exprs, global_exprs) =
527 spawn(&first_build_version, ¤t_items, remove_prior_versions).await;
528 assert_eq!(local_exprs, BTreeMap::new(), "new cache should be empty");
529 assert_eq!(global_exprs, BTreeMap::new(), "new cache should be empty");
530
531 let mut local_exps = BTreeMap::new();
533 let mut global_exps = BTreeMap::new();
534 for _ in 0..4 {
535 let id = GlobalId::User(next_id);
536 let local_exp = gen_local_expressions();
537 let global_exp = gen_global_expressions();
538
539 cache
540 .update(
541 vec![(id, local_exp.clone())],
542 vec![(id, global_exp.clone())],
543 BTreeSet::new(),
544 )
545 .await;
546
547 current_items.insert(id, RelationVersion::root());
548 current_items.extend(
549 global_exp
550 .index_imports()
551 .map(|id| (*id, RelationVersion::root())),
552 );
553 local_exps.insert(id, local_exp);
554 global_exps.insert(id, global_exp);
555
556 next_id += 1;
557 }
558 (local_exps, global_exps)
559 };
560
561 {
562 let (_cache, local_entries, global_entries) =
564 spawn(&first_build_version, ¤t_items, remove_prior_versions).await;
565 assert_eq!(
566 local_entries, local_exps,
567 "local expression with non-matching optimizer features should be removed during reconciliation"
568 );
569 assert_eq!(
570 global_entries, global_exps,
571 "global expression with non-matching optimizer features should be removed during reconciliation"
572 );
573 }
574
575 {
576 let id_to_remove = local_exps.keys().next().expect("not empty").clone();
578 current_items.remove(&id_to_remove);
579 let _removed_local_exp = local_exps.remove(&id_to_remove);
580 let _removed_global_exp = global_exps.remove(&id_to_remove);
581
582 let (_cache, local_entries, global_entries) =
584 spawn(&first_build_version, ¤t_items, remove_prior_versions).await;
585 assert_eq!(
586 local_entries, local_exps,
587 "dropped local objects should be removed during reconciliation"
588 );
589 assert_eq!(
590 global_entries, global_exps,
591 "dropped global objects should be removed during reconciliation"
592 );
593 }
594
595 {
596 let id_to_bump = global_exps.keys().next_back().expect("not empty").clone();
598 current_items.insert(id_to_bump, RelationVersion::root().bump());
599 let _removed_local_exp = local_exps.remove(&id_to_bump);
600 let _removed_global_exp = global_exps.remove(&id_to_bump);
601
602 let (_cache, local_entries, global_entries) =
604 spawn(&first_build_version, ¤t_items, remove_prior_versions).await;
605 assert_eq!(
606 local_entries, local_exps,
607 "local expressions of an earlier item version should be removed during reconciliation"
608 );
609 assert_eq!(
610 global_entries, global_exps,
611 "global expressions of an earlier item version should be removed during reconciliation"
612 );
613
614 current_items.insert(id_to_bump, RelationVersion::root());
616 let (_cache, local_entries, global_entries) =
617 spawn(&first_build_version, ¤t_items, remove_prior_versions).await;
618 assert_eq!(
619 local_entries, local_exps,
620 "local expressions of an earlier item version should be durably removed"
621 );
622 assert_eq!(
623 global_entries, global_exps,
624 "global expressions of an earlier item version should be durably removed"
625 );
626 }
627
628 {
629 let global_exp_to_remove = global_exps.keys().next().expect("not empty").clone();
631 let removed_global_exp = global_exps
632 .remove(&global_exp_to_remove)
633 .expect("known to exist");
634 let dependency_to_remove = removed_global_exp
635 .index_imports()
636 .next()
637 .expect("generator always makes non-empty index imports");
638 current_items.remove(dependency_to_remove);
639
640 let _removed_local_exp = local_exps.remove(dependency_to_remove);
642 let _removed_global_exp = global_exps.remove(dependency_to_remove);
643 global_exps.retain(|_, exp| {
645 let index_imports: BTreeSet<_> = exp.index_imports().collect();
646 !index_imports.contains(&dependency_to_remove)
647 });
648
649 let (_cache, local_entries, global_entries) =
651 spawn(&first_build_version, ¤t_items, remove_prior_versions).await;
652 assert_eq!(
653 local_entries, local_exps,
654 "dropped object dependencies should NOT remove local expressions"
655 );
656 assert_eq!(
657 global_entries, global_exps,
658 "dropped object dependencies should remove global expressions"
659 );
660 }
661
662 let (new_gen_local_exps, new_gen_global_exps) = {
663 let (cache, local_entries, global_entries) =
665 spawn(&second_build_version, ¤t_items, remove_prior_versions).await;
666 assert_eq!(
667 local_entries,
668 BTreeMap::new(),
669 "new build version should be empty"
670 );
671 assert_eq!(
672 global_entries,
673 BTreeMap::new(),
674 "new build version should be empty"
675 );
676
677 let mut local_exps = BTreeMap::new();
679 let mut global_exps = BTreeMap::new();
680 for _ in 0..2 {
681 let id = GlobalId::User(next_id);
682 let local_exp = gen_local_expressions();
683 let global_exp = gen_global_expressions();
684
685 cache
686 .update(
687 vec![(id, local_exp.clone())],
688 vec![(id, global_exp.clone())],
689 BTreeSet::new(),
690 )
691 .await;
692
693 current_items.insert(id, RelationVersion::root());
694 current_items.extend(
695 global_exp
696 .index_imports()
697 .map(|id| (*id, RelationVersion::root())),
698 );
699 local_exps.insert(id, local_exp);
700 global_exps.insert(id, global_exp);
701
702 next_id += 1;
703 }
704 (local_exps, global_exps)
705 };
706
707 {
708 let (_cache, local_entries, global_entries) =
710 spawn(&first_build_version, ¤t_items, remove_prior_versions).await;
711 assert_eq!(
712 local_entries, local_exps,
713 "Previous build version local expressions should still exist"
714 );
715 assert_eq!(
716 global_entries, global_exps,
717 "Previous build version global expressions should still exist"
718 );
719 }
720
721 {
722 remove_prior_versions = true;
724 let (_cache, local_entries, global_entries) =
725 spawn(&second_build_version, ¤t_items, remove_prior_versions).await;
726 assert_eq!(
727 local_entries, new_gen_local_exps,
728 "new build version local expressions should be persisted"
729 );
730 assert_eq!(
731 global_entries, new_gen_global_exps,
732 "new build version global expressions should be persisted"
733 );
734 }
735
736 {
737 let (_cache, local_entries, global_entries) =
739 spawn(&first_build_version, ¤t_items, remove_prior_versions).await;
740 assert_eq!(
741 local_entries,
742 BTreeMap::new(),
743 "Previous build version local expressions should be cleared"
744 );
745 assert_eq!(
746 global_entries,
747 BTreeMap::new(),
748 "Previous build version global expressions should be cleared"
749 );
750 }
751 }
752
753 #[mz_ore::test]
754 fn local_expr_cache_roundtrip() {
755 let key = CacheKey {
756 id: GlobalId::User(1),
757 build_version: "1.2.3".into(),
758 expr_type: ExpressionType::Local,
759 };
760 let val = gen_local_expressions();
761
762 let bincode_val = Bytes::from(bincode::serialize(&val).expect("must serialize"));
763 let (encoded_key, encoded_val) = ExpressionCodec::encode(&key, &bincode_val);
764 let (decoded_key, decoded_val) = ExpressionCodec::decode(&encoded_key, &encoded_val);
765 let decoded_val: LocalExpressions =
766 bincode::deserialize(&decoded_val).expect("local expressions should roundtrip");
767
768 assert_eq!(key, decoded_key);
769 assert_eq!(val, decoded_val);
770 }
771
772 #[mz_ore::test]
773 fn global_expr_cache_roundtrip() {
774 let key = CacheKey {
775 id: GlobalId::User(1),
776 build_version: "1.2.3".into(),
777 expr_type: ExpressionType::Global,
778 };
779 let val = gen_global_expressions();
780
781 let bincode_val = Bytes::from(bincode::serialize(&val).expect("must serialize"));
782 let (encoded_key, encoded_val) = ExpressionCodec::encode(&key, &bincode_val);
783 let (decoded_key, decoded_val) = ExpressionCodec::decode(&encoded_key, &encoded_val);
784 let decoded_val: GlobalExpressions =
785 bincode::deserialize(&decoded_val).expect("global expressions should roundtrip");
786
787 assert_eq!(key, decoded_key);
788 assert_eq!(val, decoded_val);
789 }
790
791 fn gen_local_expressions() -> LocalExpressions {
797 let datum = Datum::UInt64(rand::random());
798
799 LocalExpressions {
800 local_mir: OptimizedMirRelationExpr(MirRelationExpr::constant(
801 vec![vec![datum]],
802 ReprRelationType::new(vec![ReprScalarType::UInt64.nullable(false)]),
803 )),
804 optimizer_features: Default::default(),
805 item_version: RelationVersion::root(),
806 }
807 }
808
809 fn gen_global_expressions() -> GlobalExpressions {
815 let name = format!("test-{}", rand::random::<u64>());
816
817 let mut global_mir = DataflowDescription::new(name.clone());
818 let mut physical_plan = DataflowDescription::new(name);
819
820 let index_imports = BTreeMap::from_iter([(
822 GlobalId::User(2),
823 IndexImport {
824 desc: IndexDesc {
825 on_id: GlobalId::User(1),
826 key: Default::default(),
827 },
828 typ: ReprRelationType::empty(),
829 monotonic: false,
830 with_snapshot: true,
831 },
832 )]);
833 global_mir.index_imports = index_imports.clone();
834 physical_plan.index_imports = index_imports;
835
836 GlobalExpressions {
837 global_mir,
838 physical_plan,
839 dataflow_metainfos: Default::default(),
840 optimizer_features: Default::default(),
841 item_version: RelationVersion::root(),
842 }
843 }
844}