1use candid::Principal;
2use std::{
3 collections::{BTreeMap, BTreeSet},
4 num::NonZeroUsize,
5 ops::Deref,
6 panic::{AssertUnwindSafe, catch_unwind},
7 time::{Duration, Instant},
8};
9
10use super::{
11 CachedPocketIcBaseline, ControllerSnapshots, SnapshotRestoreFunding,
12 bounded_pool::{BoundedSlotLease, BoundedSlotPool},
13};
14
15#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
17pub struct FixtureRecipeId(String);
18
19impl From<SnapshotRestoreFunding> for CycleResetPolicy {
20 fn from(funding: SnapshotRestoreFunding) -> Self {
21 match funding {
22 SnapshotRestoreFunding::Preserve => Self::PreserveCurrent,
23 SnapshotRestoreFunding::TopUpTo { minimum_cycles } => Self::TopUpTo(minimum_cycles),
24 }
25 }
26}
27
28#[non_exhaustive]
30#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
31pub enum ResetDomainKind {
32 PocketIcTime,
34 ExtraCanisters,
36 PendingMessages,
38 SubnetState,
40 ExternalResources,
42}
43
44#[non_exhaustive]
46#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum CycleResetPolicy {
48 PreserveCurrent,
52 TopUpTo(u128),
55 RestoreExactBaseline,
57 RebuildOnMutation,
59}
60
61#[non_exhaustive]
63#[derive(Clone, Copy, Debug, Eq, PartialEq)]
64pub enum TimeResetPolicy {
65 PreserveCurrent,
67 RestoreBaseline,
69 RebuildOnMutation,
71}
72
73#[non_exhaustive]
75#[derive(Clone, Copy, Debug, Eq, PartialEq)]
76pub enum ExtraCanisterPolicy {
77 RequireBaselineSet,
79 RemoveTracked,
81 RebuildOnChange,
83}
84
85#[non_exhaustive]
87#[derive(Clone, Copy, Debug, Eq, PartialEq)]
88pub enum StateResetPolicy {
89 ResetByRecipe,
91 ValidateUnchanged,
93 IrrelevantByRecipeContract,
95 RebuildOnChange,
97}
98
99#[non_exhaustive]
126#[derive(Clone, Debug, Eq, PartialEq)]
127pub enum ResetDomainPolicy {
128 PocketIcTime(TimeResetPolicy),
130 ExtraCanisters(ExtraCanisterPolicy),
132 PendingMessages(StateResetPolicy),
134 SubnetState(StateResetPolicy),
136 ExternalResources(StateResetPolicy),
138}
139
140#[derive(Clone, Debug, Eq, PartialEq)]
142pub struct ResetRequirements {
143 cycle_policy: CycleResetPolicy,
144 domains: BTreeMap<ResetDomainKind, ResetDomainPolicy>,
145}
146
147#[derive(Clone, Debug, Default, Eq, PartialEq)]
149pub struct ResetReceipt(BTreeMap<ResetDomainKind, ResetDomainPolicy>);
150
151#[derive(Clone, Debug, Eq, PartialEq)]
153pub struct CanisterRestoreReceipt {
154 canister_ids: Vec<Principal>,
155 cycle_policy: CycleResetPolicy,
156}
157
158#[derive(Clone, Debug, Eq, PartialEq)]
160pub struct ReadinessReceipt {
161 identity: String,
162}
163
164#[derive(Clone, Debug, Eq, PartialEq)]
166pub struct ValidationReceipt {
167 recipe_id: FixtureRecipeId,
168 invariant_identity: String,
169}
170
171#[non_exhaustive]
173#[derive(Clone, Debug, Eq, PartialEq)]
174pub enum BaselinePoolContractError {
175 EmptyRecipeIdentity,
177 EmptyReceiptIdentity { receipt: &'static str },
179 DuplicateResetDomain { domain: ResetDomainKind },
181 DuplicateCanisterId { canister_id: Principal },
183 EmptyCanisterSet,
185 CyclePolicyMismatch {
187 required: CycleResetPolicy,
188 achieved: CycleResetPolicy,
189 },
190 RestoreCanisterSetMismatch {
192 expected: Vec<Principal>,
193 actual: Vec<Principal>,
194 },
195 MissingResetDomain { domain: ResetDomainKind },
197 ResetPolicyMismatch {
199 requirement: ResetDomainPolicy,
200 achievement: ResetDomainPolicy,
201 },
202 RecipeIdentityMismatch {
204 expected: FixtureRecipeId,
205 actual: FixtureRecipeId,
206 },
207}
208
209#[non_exhaustive]
211#[derive(Clone, Debug, Eq, PartialEq)]
212pub enum PreparedBaseline {
213 Built,
215 Restored {
217 canisters: CanisterRestoreReceipt,
219 reset: ResetReceipt,
221 readiness: ReadinessReceipt,
223 },
224}
225
226#[non_exhaustive]
228#[derive(Clone, Copy, Debug, Eq, PartialEq)]
229pub enum BaselinePreparationStage {
230 Build,
232 RestoreCanisters,
234 ResetNonSnapshotState,
236 DriveToReadiness,
238 ValidateBuilt,
240 ValidateRestored,
242}
243
244#[non_exhaustive]
246#[derive(Clone, Debug, Eq, PartialEq)]
247pub enum RebuildReason {
248 DeadPocketIcTransport,
250 SnapshotRestoreFailure,
252 ResetFailure,
254 ReadinessFailure,
256 ResetCoverageMismatch,
258 InvariantValidationFailure,
260 ExplicitLeaseInvalidation,
262 UnwindWhileLeased,
264 RecipeClassified { code: String },
266}
267
268#[non_exhaustive]
270#[derive(Clone, Debug, Eq, PartialEq)]
271pub enum FailureDisposition {
272 Fatal,
274 Rebuild(RebuildReason),
276}
277
278#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
280pub struct BaselinePoolTimings {
281 wait: Duration,
282 build: Option<Duration>,
283 restore: Option<Duration>,
284 reset: Option<Duration>,
285 readiness: Option<Duration>,
286 validation: Option<Duration>,
287 stale_teardown: Option<Duration>,
288 total: Duration,
289}
290
291#[non_exhaustive]
293#[derive(Clone, Debug, Eq, PartialEq)]
294pub enum BaselinePoolOutcome {
295 Built {
297 slot: usize,
299 timings: BaselinePoolTimings,
301 },
302 Restored {
304 slot: usize,
306 timings: BaselinePoolTimings,
308 },
309 Rebuilt {
311 slot: usize,
313 reason: RebuildReason,
315 timings: BaselinePoolTimings,
317 },
318}
319
320#[non_exhaustive]
322#[derive(Debug)]
323pub enum BaselinePoolPreparationError<E> {
324 Recipe {
326 stage: BaselinePreparationStage,
328 source: E,
330 },
331 Contract(BaselinePoolContractError),
333}
334
335#[non_exhaustive]
337#[derive(Debug)]
338pub enum BaselinePoolError<E> {
339 Preparation {
341 error: BaselinePoolPreparationError<E>,
343 timings: Box<BaselinePoolTimings>,
345 },
346 RecoveryFailed {
348 original: Box<BaselinePoolPreparationError<E>>,
350 rebuild: Box<BaselinePoolPreparationError<E>>,
352 timings: Box<BaselinePoolTimings>,
354 },
355}
356
357pub trait PocketIcBaselineRecipe: Send + Sync + 'static {
359 type Metadata: Send + 'static;
361 type Error: std::error::Error + Send + Sync + 'static;
363
364 fn id(&self) -> &FixtureRecipeId;
366
367 fn reset_requirements(&self) -> &ResetRequirements;
369
370 fn build(&self) -> Result<CachedPocketIcBaseline<Self::Metadata>, Self::Error>;
372
373 fn restore_canisters(
375 &self,
376 baseline: &CachedPocketIcBaseline<Self::Metadata>,
377 ) -> Result<CanisterRestoreReceipt, Self::Error>;
378
379 fn reset_non_snapshot_state(
381 &self,
382 baseline: &CachedPocketIcBaseline<Self::Metadata>,
383 ) -> Result<ResetReceipt, Self::Error>;
384
385 fn drive_to_readiness(
387 &self,
388 baseline: &CachedPocketIcBaseline<Self::Metadata>,
389 ) -> Result<ReadinessReceipt, Self::Error>;
390
391 fn validate(
393 &self,
394 baseline: &CachedPocketIcBaseline<Self::Metadata>,
395 preparation: &PreparedBaseline,
396 ) -> Result<ValidationReceipt, Self::Error>;
397
398 fn classify_failure(
400 &self,
401 stage: BaselinePreparationStage,
402 _error: &Self::Error,
403 ) -> FailureDisposition {
404 FailureDisposition::Rebuild(stage.default_rebuild_reason())
405 }
406}
407
408pub struct CachedPocketIcBaselinePool<R>
420where
421 R: PocketIcBaselineRecipe,
422{
423 recipe: R,
424 slots: BoundedSlotPool<BaselineSlot<R::Metadata>>,
425}
426
427struct BaselineSlot<M> {
428 baseline: CachedPocketIcBaseline<M>,
429 invalidation_reason: Option<RebuildReason>,
430}
431
432pub struct CachedPocketIcBaselinePoolGuard<'a, R>
434where
435 R: PocketIcBaselineRecipe,
436{
437 slot: BoundedSlotLease<'a, BaselineSlot<R::Metadata>>,
438}
439
440impl FixtureRecipeId {
441 pub fn try_new(identity: impl Into<String>) -> Result<Self, BaselinePoolContractError> {
443 let identity = identity.into();
444 if identity.trim().is_empty() {
445 return Err(BaselinePoolContractError::EmptyRecipeIdentity);
446 }
447 Ok(Self(identity))
448 }
449
450 #[must_use]
452 pub fn as_str(&self) -> &str {
453 &self.0
454 }
455}
456
457impl ResetDomainPolicy {
458 #[must_use]
460 pub const fn domain(&self) -> ResetDomainKind {
461 match self {
462 Self::PocketIcTime(_) => ResetDomainKind::PocketIcTime,
463 Self::ExtraCanisters(_) => ResetDomainKind::ExtraCanisters,
464 Self::PendingMessages(_) => ResetDomainKind::PendingMessages,
465 Self::SubnetState(_) => ResetDomainKind::SubnetState,
466 Self::ExternalResources(_) => ResetDomainKind::ExternalResources,
467 }
468 }
469}
470
471fn collect_reset_domains(
472 policies: impl IntoIterator<Item = ResetDomainPolicy>,
473) -> Result<BTreeMap<ResetDomainKind, ResetDomainPolicy>, BaselinePoolContractError> {
474 let mut domains = BTreeMap::new();
475 for policy in policies {
476 let domain = policy.domain();
477 if domains.insert(domain, policy).is_some() {
478 return Err(BaselinePoolContractError::DuplicateResetDomain { domain });
479 }
480 }
481 Ok(domains)
482}
483
484impl ResetRequirements {
485 pub fn try_new<I>(
490 cycle_policy: CycleResetPolicy,
491 requirements: I,
492 ) -> Result<Self, BaselinePoolContractError>
493 where
494 I: IntoIterator<Item = ResetDomainPolicy>,
495 {
496 Ok(Self {
497 cycle_policy,
498 domains: collect_reset_domains(requirements)?,
499 })
500 }
501
502 #[must_use]
504 pub const fn cycle_policy(&self) -> CycleResetPolicy {
505 self.cycle_policy
506 }
507
508 #[must_use]
510 pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetDomainPolicy> {
511 self.domains.get(&domain)
512 }
513
514 pub fn iter(&self) -> impl Iterator<Item = &ResetDomainPolicy> {
516 self.domains.values()
517 }
518
519 fn verify(
520 &self,
521 restore: &CanisterRestoreReceipt,
522 receipt: &ResetReceipt,
523 ) -> Result<(), BaselinePoolContractError> {
524 if restore.cycle_policy != self.cycle_policy {
525 return Err(BaselinePoolContractError::CyclePolicyMismatch {
526 required: self.cycle_policy,
527 achieved: restore.cycle_policy,
528 });
529 }
530 for (domain, requirement) in &self.domains {
531 let Some(achievement) = receipt.0.get(domain) else {
532 return Err(BaselinePoolContractError::MissingResetDomain { domain: *domain });
533 };
534 if achievement != requirement {
535 return Err(BaselinePoolContractError::ResetPolicyMismatch {
536 requirement: requirement.clone(),
537 achievement: achievement.clone(),
538 });
539 }
540 }
541 Ok(())
542 }
543}
544
545impl ResetReceipt {
546 pub fn try_new<I>(achievements: I) -> Result<Self, BaselinePoolContractError>
551 where
552 I: IntoIterator<Item = ResetDomainPolicy>,
553 {
554 Ok(Self(collect_reset_domains(achievements)?))
555 }
556
557 #[must_use]
559 pub const fn empty() -> Self {
560 Self(BTreeMap::new())
561 }
562
563 #[must_use]
565 pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetDomainPolicy> {
566 self.0.get(&domain)
567 }
568
569 pub fn iter(&self) -> impl Iterator<Item = &ResetDomainPolicy> {
571 self.0.values()
572 }
573}
574
575impl CanisterRestoreReceipt {
576 pub fn try_new<I>(
578 canister_ids: I,
579 cycle_policy: CycleResetPolicy,
580 ) -> Result<Self, BaselinePoolContractError>
581 where
582 I: IntoIterator<Item = Principal>,
583 {
584 let mut unique = BTreeSet::new();
585 for canister_id in canister_ids {
586 if !unique.insert(canister_id) {
587 return Err(BaselinePoolContractError::DuplicateCanisterId { canister_id });
588 }
589 }
590 if unique.is_empty() {
591 return Err(BaselinePoolContractError::EmptyCanisterSet);
592 }
593 Ok(Self {
594 canister_ids: unique.into_iter().collect(),
595 cycle_policy,
596 })
597 }
598
599 pub(super) fn from_restore(
602 snapshots: &ControllerSnapshots,
603 funding: SnapshotRestoreFunding,
604 ) -> Self {
605 Self {
606 canister_ids: snapshots.canister_ids().collect(),
607 cycle_policy: funding.into(),
608 }
609 }
610
611 #[must_use]
613 pub fn canister_ids(&self) -> &[Principal] {
614 &self.canister_ids
615 }
616
617 #[must_use]
619 pub const fn cycle_policy(&self) -> CycleResetPolicy {
620 self.cycle_policy
621 }
622}
623
624impl ReadinessReceipt {
625 pub fn try_new(identity: impl Into<String>) -> Result<Self, BaselinePoolContractError> {
627 Ok(Self {
628 identity: nonempty_receipt_identity("readiness", identity.into())?,
629 })
630 }
631
632 #[must_use]
634 pub fn identity(&self) -> &str {
635 &self.identity
636 }
637}
638
639impl ValidationReceipt {
640 pub fn try_new(
642 recipe_id: FixtureRecipeId,
643 invariant_identity: impl Into<String>,
644 ) -> Result<Self, BaselinePoolContractError> {
645 Ok(Self {
646 recipe_id,
647 invariant_identity: nonempty_receipt_identity("validation", invariant_identity.into())?,
648 })
649 }
650
651 #[must_use]
653 pub const fn recipe_id(&self) -> &FixtureRecipeId {
654 &self.recipe_id
655 }
656
657 #[must_use]
659 pub fn invariant_identity(&self) -> &str {
660 &self.invariant_identity
661 }
662}
663
664impl BaselinePreparationStage {
665 #[must_use]
670 pub fn default_rebuild_reason(self) -> RebuildReason {
671 match self {
672 Self::RestoreCanisters => RebuildReason::SnapshotRestoreFailure,
673 Self::ResetNonSnapshotState => RebuildReason::ResetFailure,
674 Self::DriveToReadiness => RebuildReason::ReadinessFailure,
675 Self::ValidateRestored | Self::ValidateBuilt => {
676 RebuildReason::InvariantValidationFailure
677 }
678 Self::Build => RebuildReason::RecipeClassified {
679 code: "build".to_owned(),
680 },
681 }
682 }
683}
684
685impl BaselinePoolTimings {
686 #[must_use]
688 pub const fn wait(self) -> Duration {
689 self.wait
690 }
691
692 #[must_use]
694 pub const fn build(self) -> Option<Duration> {
695 self.build
696 }
697
698 #[must_use]
700 pub const fn restore(self) -> Option<Duration> {
701 self.restore
702 }
703
704 #[must_use]
706 pub const fn reset(self) -> Option<Duration> {
707 self.reset
708 }
709
710 #[must_use]
712 pub const fn readiness(self) -> Option<Duration> {
713 self.readiness
714 }
715
716 #[must_use]
718 pub const fn validation(self) -> Option<Duration> {
719 self.validation
720 }
721
722 #[must_use]
724 pub const fn stale_teardown(self) -> Option<Duration> {
725 self.stale_teardown
726 }
727
728 #[must_use]
730 pub const fn total(self) -> Duration {
731 self.total
732 }
733}
734
735impl BaselinePoolOutcome {
736 #[must_use]
738 pub const fn slot(&self) -> usize {
739 match self {
740 Self::Built { slot, .. } | Self::Restored { slot, .. } | Self::Rebuilt { slot, .. } => {
741 *slot
742 }
743 }
744 }
745
746 #[must_use]
748 pub const fn timings(&self) -> BaselinePoolTimings {
749 match self {
750 Self::Built { timings, .. }
751 | Self::Restored { timings, .. }
752 | Self::Rebuilt { timings, .. } => *timings,
753 }
754 }
755}
756
757impl std::fmt::Display for BaselinePoolTimings {
758 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
759 write!(
760 formatter,
761 "total={:?} wait={:?} build={:?} restore={:?} reset={:?} readiness={:?} validation={:?} stale_teardown={:?}",
762 self.total,
763 self.wait,
764 self.build,
765 self.restore,
766 self.reset,
767 self.readiness,
768 self.validation,
769 self.stale_teardown,
770 )
771 }
772}
773
774impl std::fmt::Display for BaselinePoolOutcome {
775 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
776 match self {
777 Self::Built { slot, timings } => write!(formatter, "built slot={slot} {timings}"),
778 Self::Restored { slot, timings } => {
779 write!(formatter, "restored slot={slot} {timings}")
780 }
781 Self::Rebuilt {
782 slot,
783 reason,
784 timings,
785 } => write!(formatter, "rebuilt slot={slot} reason={reason:?} {timings}"),
786 }
787 }
788}
789
790impl<E> BaselinePoolError<E> {
791 #[must_use]
793 pub const fn timings(&self) -> BaselinePoolTimings {
794 match self {
795 Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
796 }
797 }
798}
799
800impl<R> CachedPocketIcBaselinePool<R>
801where
802 R: PocketIcBaselineRecipe,
803{
804 #[must_use]
809 pub const fn new(capacity: NonZeroUsize, recipe: R) -> Self {
810 Self {
811 recipe,
812 slots: BoundedSlotPool::new(capacity),
813 }
814 }
815
816 #[must_use]
818 pub fn recipe_id(&self) -> &FixtureRecipeId {
819 self.recipe.id()
820 }
821
822 #[must_use]
824 pub const fn capacity(&self) -> NonZeroUsize {
825 self.slots.capacity()
826 }
827
828 pub fn acquire(
840 &self,
841 ) -> Result<
842 (CachedPocketIcBaselinePoolGuard<'_, R>, BaselinePoolOutcome),
843 BaselinePoolError<R::Error>,
844 > {
845 let total_started = Instant::now();
846 let mut slot = self.slots.acquire();
847 let mut timings = BaselinePoolTimings {
848 wait: slot.wait(),
849 ..BaselinePoolTimings::default()
850 };
851 let slot_index = slot.slot_index();
852
853 let (rebuild_reason, original) = if slot.is_reusable() {
854 match self.prepare_reused(&slot, &mut timings) {
855 Ok(()) => {
856 timings.total = total_started.elapsed();
857 return Ok((
858 CachedPocketIcBaselinePoolGuard { slot },
859 BaselinePoolOutcome::Restored {
860 slot: slot_index,
861 timings,
862 },
863 ));
864 }
865 Err(original) => match self.failure_disposition(&original) {
866 FailureDisposition::Fatal => {
867 Self::discard_stale_slot(&mut slot, &mut timings);
870 timings.total = total_started.elapsed();
871 return Err(BaselinePoolError::Preparation {
872 error: original,
873 timings: Box::new(timings),
874 });
875 }
876 FailureDisposition::Rebuild(reason) => (Some(reason), Some(original)),
877 },
878 }
879 } else {
880 let reason = if slot.invalidated_by_unwind() {
881 Some(RebuildReason::UnwindWhileLeased)
882 } else {
883 slot.get()
884 .and_then(|slot| slot.invalidation_reason.clone())
885 .or_else(|| {
886 slot.is_populated()
887 .then_some(RebuildReason::ExplicitLeaseInvalidation)
888 })
889 };
890 (reason, None)
891 };
892
893 if slot.is_populated() {
894 Self::discard_stale_slot(&mut slot, &mut timings);
895 }
896 let build = self.build_slot(&mut slot, &mut timings);
897 timings.total = total_started.elapsed();
898 if let Err(error) = build {
899 return Err(match original {
900 Some(original) => BaselinePoolError::RecoveryFailed {
901 original: Box::new(original),
902 rebuild: Box::new(error),
903 timings: Box::new(timings),
904 },
905 None => BaselinePoolError::Preparation {
906 error,
907 timings: Box::new(timings),
908 },
909 });
910 }
911
912 let outcome = rebuild_reason.map_or_else(
913 || BaselinePoolOutcome::Built {
914 slot: slot_index,
915 timings,
916 },
917 |reason| BaselinePoolOutcome::Rebuilt {
918 slot: slot_index,
919 reason,
920 timings,
921 },
922 );
923 Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
924 }
925
926 fn prepare_reused(
927 &self,
928 slot: &BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
929 timings: &mut BaselinePoolTimings,
930 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
931 let baseline = &slot
932 .get()
933 .expect("reusable baseline slot must be populated")
934 .baseline;
935
936 let started = Instant::now();
937 let restore = self.recipe.restore_canisters(baseline);
938 add_timing(&mut timings.restore, started.elapsed());
939 let canisters = restore.map_err(|source| BaselinePoolPreparationError::Recipe {
940 stage: BaselinePreparationStage::RestoreCanisters,
941 source,
942 })?;
943 if !baseline
944 .snapshot_canister_ids()
945 .eq(canisters.canister_ids().iter().copied())
946 {
947 return Err(BaselinePoolPreparationError::Contract(
948 BaselinePoolContractError::RestoreCanisterSetMismatch {
949 expected: baseline.snapshot_canister_ids().collect(),
950 actual: canisters.canister_ids().to_vec(),
951 },
952 ));
953 }
954
955 let started = Instant::now();
956 let reset_result = self.recipe.reset_non_snapshot_state(baseline);
957 add_timing(&mut timings.reset, started.elapsed());
958 let reset = reset_result.map_err(|source| BaselinePoolPreparationError::Recipe {
959 stage: BaselinePreparationStage::ResetNonSnapshotState,
960 source,
961 })?;
962
963 let started = Instant::now();
964 let readiness_result = self.recipe.drive_to_readiness(baseline);
965 add_timing(&mut timings.readiness, started.elapsed());
966 let readiness =
967 readiness_result.map_err(|source| BaselinePoolPreparationError::Recipe {
968 stage: BaselinePreparationStage::DriveToReadiness,
969 source,
970 })?;
971 self.recipe
972 .reset_requirements()
973 .verify(&canisters, &reset)
974 .map_err(BaselinePoolPreparationError::Contract)?;
975
976 let preparation = PreparedBaseline::Restored {
977 canisters,
978 reset,
979 readiness,
980 };
981 self.validate_baseline(
982 baseline,
983 &preparation,
984 BaselinePreparationStage::ValidateRestored,
985 timings,
986 )?;
987 Ok(())
988 }
989
990 fn build_slot(
991 &self,
992 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
993 timings: &mut BaselinePoolTimings,
994 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
995 let started = Instant::now();
996 let build = self.recipe.build();
997 add_timing(&mut timings.build, started.elapsed());
998 let baseline = build.map_err(|source| BaselinePoolPreparationError::Recipe {
999 stage: BaselinePreparationStage::Build,
1000 source,
1001 })?;
1002
1003 if let Err(error) = self.validate_baseline(
1004 &baseline,
1005 &PreparedBaseline::Built,
1006 BaselinePreparationStage::ValidateBuilt,
1007 timings,
1008 ) {
1009 drop_baseline_safely(baseline);
1010 return Err(error);
1011 }
1012 let replaced = slot.replace(BaselineSlot {
1013 baseline,
1014 invalidation_reason: None,
1015 });
1016 debug_assert!(replaced.is_none());
1017 Ok(())
1018 }
1019
1020 fn validate_baseline(
1021 &self,
1022 baseline: &CachedPocketIcBaseline<R::Metadata>,
1023 preparation: &PreparedBaseline,
1024 stage: BaselinePreparationStage,
1025 timings: &mut BaselinePoolTimings,
1026 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1027 let started = Instant::now();
1028 let validation = self.recipe.validate(baseline, preparation);
1029 add_timing(&mut timings.validation, started.elapsed());
1030 let receipt =
1031 validation.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })?;
1032 if receipt.recipe_id() != self.recipe.id() {
1033 return Err(BaselinePoolPreparationError::Contract(
1034 BaselinePoolContractError::RecipeIdentityMismatch {
1035 expected: self.recipe.id().clone(),
1036 actual: receipt.recipe_id().clone(),
1037 },
1038 ));
1039 }
1040 Ok(())
1041 }
1042
1043 fn failure_disposition(
1044 &self,
1045 error: &BaselinePoolPreparationError<R::Error>,
1046 ) -> FailureDisposition {
1047 match error {
1048 BaselinePoolPreparationError::Recipe { stage, source } => {
1049 self.recipe.classify_failure(*stage, source)
1050 }
1051 BaselinePoolPreparationError::Contract(
1052 BaselinePoolContractError::RecipeIdentityMismatch { .. },
1053 ) => FailureDisposition::Fatal,
1054 BaselinePoolPreparationError::Contract(_) => {
1055 FailureDisposition::Rebuild(rebuild_reason_for_error(error))
1056 }
1057 }
1058 }
1059
1060 fn discard_stale_slot(
1061 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1062 timings: &mut BaselinePoolTimings,
1063 ) {
1064 let started = Instant::now();
1065 if let Some(stale) = slot.take() {
1066 drop_baseline_safely(stale.baseline);
1067 }
1068 timings.stale_teardown = Some(started.elapsed());
1069 }
1070}
1071
1072impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
1073where
1074 R: PocketIcBaselineRecipe,
1075{
1076 #[must_use]
1078 pub const fn slot(&self) -> usize {
1079 self.slot.slot_index()
1080 }
1081
1082 pub fn invalidate(&mut self, reason: RebuildReason) {
1084 if let Some(slot) = self.slot.get_mut() {
1085 slot.invalidation_reason = Some(reason);
1086 }
1087 self.slot.invalidate();
1088 }
1089}
1090
1091impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
1092where
1093 R: PocketIcBaselineRecipe,
1094{
1095 type Target = CachedPocketIcBaseline<R::Metadata>;
1096
1097 fn deref(&self) -> &Self::Target {
1098 &self
1099 .slot
1100 .get()
1101 .expect("leased baseline pool slot must be populated")
1102 .baseline
1103 }
1104}
1105
1106fn nonempty_receipt_identity(
1107 receipt: &'static str,
1108 identity: String,
1109) -> Result<String, BaselinePoolContractError> {
1110 if identity.trim().is_empty() {
1111 return Err(BaselinePoolContractError::EmptyReceiptIdentity { receipt });
1112 }
1113 Ok(identity)
1114}
1115
1116const fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
1117 *total = Some(match *total {
1118 None => elapsed,
1119 Some(previous) => previous.saturating_add(elapsed),
1120 });
1121}
1122
1123fn rebuild_reason_for_error<E>(error: &BaselinePoolPreparationError<E>) -> RebuildReason {
1124 match error {
1125 BaselinePoolPreparationError::Recipe { stage, .. } => stage.default_rebuild_reason(),
1126 BaselinePoolPreparationError::Contract(
1127 BaselinePoolContractError::MissingResetDomain { .. }
1128 | BaselinePoolContractError::ResetPolicyMismatch { .. }
1129 | BaselinePoolContractError::CyclePolicyMismatch { .. }
1130 | BaselinePoolContractError::DuplicateResetDomain { .. }
1131 | BaselinePoolContractError::RestoreCanisterSetMismatch { .. },
1132 ) => RebuildReason::ResetCoverageMismatch,
1133 BaselinePoolPreparationError::Contract(_) => RebuildReason::InvariantValidationFailure,
1134 }
1135}
1136
1137fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
1138 let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
1139}
1140
1141impl std::fmt::Display for FixtureRecipeId {
1142 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1143 formatter.write_str(&self.0)
1144 }
1145}
1146
1147impl std::fmt::Display for BaselinePreparationStage {
1148 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1149 formatter.write_str(match self {
1150 Self::Build => "build",
1151 Self::RestoreCanisters => "canister restore",
1152 Self::ResetNonSnapshotState => "non-snapshot reset",
1153 Self::DriveToReadiness => "readiness",
1154 Self::ValidateBuilt => "built-baseline validation",
1155 Self::ValidateRestored => "restored-baseline validation",
1156 })
1157 }
1158}
1159
1160impl std::fmt::Display for BaselinePoolContractError {
1161 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1162 match self {
1163 Self::EmptyRecipeIdentity => formatter.write_str("fixture recipe identity is empty"),
1164 Self::EmptyReceiptIdentity { receipt } => {
1165 write!(formatter, "{receipt} receipt identity is empty")
1166 }
1167 Self::DuplicateResetDomain { domain } => {
1168 write!(
1169 formatter,
1170 "reset domain {domain:?} was reported more than once"
1171 )
1172 }
1173 Self::DuplicateCanisterId { canister_id } => {
1174 write!(formatter, "restore receipt repeats canister {canister_id}")
1175 }
1176 Self::EmptyCanisterSet => formatter.write_str("restore receipt contains no canisters"),
1177 Self::CyclePolicyMismatch { required, achieved } => write!(
1178 formatter,
1179 "restore cycle policy {achieved:?} does not satisfy {required:?}",
1180 ),
1181 Self::RestoreCanisterSetMismatch { expected, actual } => write!(
1182 formatter,
1183 "restore receipt identified canisters {actual:?}, expected captured set {expected:?}",
1184 ),
1185 Self::MissingResetDomain { domain } => {
1186 write!(
1187 formatter,
1188 "required reset domain {domain:?} was not achieved"
1189 )
1190 }
1191 Self::ResetPolicyMismatch {
1192 requirement,
1193 achievement,
1194 } => write!(
1195 formatter,
1196 "reset achievement {achievement:?} does not satisfy {requirement:?}",
1197 ),
1198 Self::RecipeIdentityMismatch { expected, actual } => write!(
1199 formatter,
1200 "validation receipt used recipe `{actual}` instead of `{expected}`",
1201 ),
1202 }
1203 }
1204}
1205
1206impl std::error::Error for BaselinePoolContractError {}
1207
1208impl<E> std::fmt::Display for BaselinePoolPreparationError<E>
1209where
1210 E: std::fmt::Display,
1211{
1212 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1213 match self {
1214 Self::Recipe { stage, source } => {
1215 write!(formatter, "baseline {stage} failed: {source}")
1216 }
1217 Self::Contract(error) => write!(formatter, "baseline pool contract failed: {error}"),
1218 }
1219 }
1220}
1221
1222impl<E> std::error::Error for BaselinePoolPreparationError<E>
1223where
1224 E: std::error::Error + 'static,
1225{
1226 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1227 match self {
1228 Self::Recipe { source, .. } => Some(source),
1229 Self::Contract(error) => Some(error),
1230 }
1231 }
1232}
1233
1234impl<E> std::fmt::Display for BaselinePoolError<E>
1235where
1236 E: std::fmt::Display,
1237{
1238 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1239 match self {
1240 Self::Preparation { error, .. } => error.fmt(formatter),
1241 Self::RecoveryFailed {
1242 original, rebuild, ..
1243 } => write!(
1244 formatter,
1245 "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
1246 ),
1247 }
1248 }
1249}
1250
1251impl<E> std::error::Error for BaselinePoolError<E>
1252where
1253 E: std::error::Error + 'static,
1254{
1255 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1256 match self {
1257 Self::Preparation { error, .. } => Some(error),
1258 Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
1259 }
1260 }
1261}
1262
1263#[cfg(test)]
1264mod tests {
1265 use super::{
1266 BaselinePoolContractError, BaselinePoolTimings, CanisterRestoreReceipt, CycleResetPolicy,
1267 ExtraCanisterPolicy, FixtureRecipeId, ResetDomainKind, ResetDomainPolicy, ResetReceipt,
1268 ResetRequirements, StateResetPolicy, TimeResetPolicy, add_timing,
1269 };
1270 use candid::Principal;
1271 use std::time::Duration;
1272
1273 #[test]
1274 fn phase_timing_preserves_absence_zero_and_saturation() {
1275 let mut timings = BaselinePoolTimings::default();
1276 assert_eq!(timings.restore(), None);
1277 add_timing(&mut timings.restore, Duration::ZERO);
1278 assert_eq!(timings.restore(), Some(Duration::ZERO));
1279 add_timing(&mut timings.restore, Duration::from_secs(2));
1280 assert_eq!(timings.restore(), Some(Duration::from_secs(2)));
1281 add_timing(&mut timings.restore, Duration::MAX);
1282 assert_eq!(timings.restore(), Some(Duration::MAX));
1283 assert_eq!(timings.build(), None);
1284 }
1285
1286 #[test]
1287 fn recipe_identity_must_be_nonempty() {
1288 assert!(matches!(
1289 FixtureRecipeId::try_new(" "),
1290 Err(BaselinePoolContractError::EmptyRecipeIdentity)
1291 ));
1292 }
1293
1294 #[test]
1295 fn caller_restore_receipts_validate_and_order_canister_ids() {
1296 let first = Principal::from_slice(&[1]);
1297 let second = Principal::from_slice(&[2]);
1298 let policy = CycleResetPolicy::TopUpTo(123);
1299 let receipt = CanisterRestoreReceipt::try_new([second, first], policy).unwrap();
1300 assert_eq!(receipt.canister_ids(), [first, second]);
1301 assert_eq!(receipt.cycle_policy(), policy);
1302 assert_eq!(
1303 CanisterRestoreReceipt::try_new([], policy),
1304 Err(BaselinePoolContractError::EmptyCanisterSet),
1305 );
1306 assert_eq!(
1307 CanisterRestoreReceipt::try_new([second, first, second], policy),
1308 Err(BaselinePoolContractError::DuplicateCanisterId {
1309 canister_id: second,
1310 }),
1311 );
1312 }
1313
1314 #[test]
1315 fn reset_declarations_and_receipts_reject_duplicate_domains() {
1316 let policies = [
1317 ResetDomainPolicy::PocketIcTime(TimeResetPolicy::PreserveCurrent),
1318 ResetDomainPolicy::PocketIcTime(TimeResetPolicy::RebuildOnMutation),
1319 ];
1320 let error = BaselinePoolContractError::DuplicateResetDomain {
1321 domain: ResetDomainKind::PocketIcTime,
1322 };
1323 assert_eq!(
1324 ResetRequirements::try_new(CycleResetPolicy::PreserveCurrent, policies.clone()),
1325 Err(error.clone()),
1326 );
1327 assert_eq!(ResetReceipt::try_new(policies), Err(error));
1328 }
1329
1330 #[test]
1331 fn required_policy_must_match_achieved_policy() {
1332 let restore = CanisterRestoreReceipt::try_new(
1333 [Principal::anonymous()],
1334 CycleResetPolicy::PreserveCurrent,
1335 )
1336 .unwrap();
1337 for (requirement, achievement) in [
1338 (
1339 ResetDomainPolicy::PocketIcTime(TimeResetPolicy::PreserveCurrent),
1340 ResetDomainPolicy::PocketIcTime(TimeResetPolicy::RebuildOnMutation),
1341 ),
1342 (
1343 ResetDomainPolicy::ExtraCanisters(ExtraCanisterPolicy::RequireBaselineSet),
1344 ResetDomainPolicy::ExtraCanisters(ExtraCanisterPolicy::RemoveTracked),
1345 ),
1346 (
1347 ResetDomainPolicy::PendingMessages(StateResetPolicy::ValidateUnchanged),
1348 ResetDomainPolicy::PendingMessages(StateResetPolicy::IrrelevantByRecipeContract),
1349 ),
1350 (
1351 ResetDomainPolicy::SubnetState(StateResetPolicy::ResetByRecipe),
1352 ResetDomainPolicy::SubnetState(StateResetPolicy::RebuildOnChange),
1353 ),
1354 (
1355 ResetDomainPolicy::ExternalResources(StateResetPolicy::ValidateUnchanged),
1356 ResetDomainPolicy::ExternalResources(StateResetPolicy::ResetByRecipe),
1357 ),
1358 ] {
1359 let requirements = ResetRequirements::try_new(
1360 CycleResetPolicy::PreserveCurrent,
1361 [requirement.clone()],
1362 )
1363 .unwrap();
1364 let receipt = ResetReceipt::try_new([requirement.clone()]).unwrap();
1365 assert_eq!(requirements.verify(&restore, &receipt), Ok(()));
1366 assert_eq!(
1367 requirements.verify(&restore, &ResetReceipt::empty()),
1368 Err(BaselinePoolContractError::MissingResetDomain {
1369 domain: requirement.domain(),
1370 }),
1371 );
1372 let receipt = ResetReceipt::try_new([achievement.clone()]).unwrap();
1373 assert_eq!(
1374 requirements.verify(&restore, &receipt),
1375 Err(BaselinePoolContractError::ResetPolicyMismatch {
1376 requirement,
1377 achievement,
1378 }),
1379 );
1380 }
1381 }
1382
1383 #[test]
1384 fn restore_cycle_policy_must_match_required_policy() {
1385 let requirements =
1386 ResetRequirements::try_new(CycleResetPolicy::RestoreExactBaseline, []).unwrap();
1387 let restore = CanisterRestoreReceipt::try_new(
1388 [Principal::anonymous()],
1389 CycleResetPolicy::PreserveCurrent,
1390 )
1391 .unwrap();
1392 assert!(matches!(
1393 requirements.verify(&restore, &ResetReceipt::empty()),
1394 Err(BaselinePoolContractError::CyclePolicyMismatch {
1395 required: CycleResetPolicy::RestoreExactBaseline,
1396 achieved: CycleResetPolicy::PreserveCurrent,
1397 })
1398 ));
1399 }
1400}