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,
12 bounded_pool::{BoundedSlotLease, BoundedSlotPool},
13};
14
15#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
17pub struct FixtureRecipeId(String);
18
19#[non_exhaustive]
21#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
22pub enum ResetDomainKind {
23 CanisterSnapshots,
25 CanisterCycles,
27 PocketIcTime,
29 ExtraCanisters,
31 PendingMessages,
33 SubnetState,
35 ExternalResources,
37}
38
39#[non_exhaustive]
41#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42pub enum CycleResetPolicy {
43 PreserveCurrent,
47 TopUpTo(u128),
50 RestoreExactBaseline,
52 RebuildOnMutation,
54}
55
56#[non_exhaustive]
58#[derive(Clone, Copy, Debug, Eq, PartialEq)]
59pub enum TimeResetPolicy {
60 PreserveCurrent,
62 RestoreBaseline,
64 RebuildOnMutation,
66}
67
68#[non_exhaustive]
70#[derive(Clone, Copy, Debug, Eq, PartialEq)]
71pub enum ExtraCanisterPolicy {
72 RequireBaselineSet,
74 RemoveTracked,
76 RebuildOnChange,
78}
79
80#[non_exhaustive]
82#[derive(Clone, Copy, Debug, Eq, PartialEq)]
83pub enum StateResetPolicy {
84 ResetByRecipe,
86 ValidateUnchanged,
88 IrrelevantByRecipeContract,
90 RebuildOnChange,
92}
93
94#[non_exhaustive]
96#[derive(Clone, Debug, Eq, PartialEq)]
97pub enum ResetRequirement {
98 CanisterSnapshots,
100 CanisterCycles(CycleResetPolicy),
102 PocketIcTime(TimeResetPolicy),
104 ExtraCanisters(ExtraCanisterPolicy),
106 PendingMessages(StateResetPolicy),
108 SubnetState(StateResetPolicy),
110 ExternalResources(StateResetPolicy),
112}
113
114#[non_exhaustive]
116#[derive(Clone, Debug, Eq, PartialEq)]
117pub enum ResetAchievement {
118 CanisterSnapshots,
120 CanisterCycles(CycleResetPolicy),
122 PocketIcTime(TimeResetPolicy),
124 ExtraCanisters(ExtraCanisterPolicy),
126 PendingMessages(StateResetPolicy),
128 SubnetState(StateResetPolicy),
130 ExternalResources(StateResetPolicy),
132}
133
134#[derive(Clone, Debug, Eq, PartialEq)]
136pub struct ResetRequirements(BTreeMap<ResetDomainKind, ResetRequirement>);
137
138#[derive(Clone, Debug, Default, Eq, PartialEq)]
140pub struct ResetReceipt(BTreeMap<ResetDomainKind, ResetAchievement>);
141
142#[derive(Clone, Debug, Eq, PartialEq)]
144pub struct CanisterRestoreReceipt {
145 canister_ids: Vec<Principal>,
146 cycle_policy: CycleResetPolicy,
147}
148
149#[derive(Clone, Debug, Eq, PartialEq)]
151pub struct ReadinessReceipt {
152 identity: String,
153}
154
155#[derive(Clone, Debug, Eq, PartialEq)]
157pub struct ValidationReceipt {
158 recipe_id: FixtureRecipeId,
159 invariant_identity: String,
160}
161
162#[non_exhaustive]
164#[derive(Clone, Debug, Eq, PartialEq)]
165pub enum BaselinePoolContractError {
166 EmptyRecipeIdentity,
168 EmptyReceiptIdentity { receipt: &'static str },
170 DuplicateResetDomain { domain: ResetDomainKind },
172 DuplicateCanisterId { canister_id: Principal },
174 EmptyCanisterSet,
176 UndeclaredRequiredResetDomain { domain: ResetDomainKind },
178 RestoreCanisterSetMismatch {
180 expected: Vec<Principal>,
181 actual: Vec<Principal>,
182 },
183 MissingResetDomain { domain: ResetDomainKind },
185 ResetPolicyMismatch {
187 requirement: ResetRequirement,
188 achievement: ResetAchievement,
189 },
190 RecipeIdentityMismatch {
192 expected: FixtureRecipeId,
193 actual: FixtureRecipeId,
194 },
195}
196
197#[non_exhaustive]
199#[derive(Clone, Debug, Eq, PartialEq)]
200pub enum PreparedBaseline {
201 Built,
203 Restored {
205 canisters: CanisterRestoreReceipt,
207 reset: ResetReceipt,
209 readiness: ReadinessReceipt,
211 },
212}
213
214#[non_exhaustive]
216#[derive(Clone, Copy, Debug, Eq, PartialEq)]
217pub enum BaselinePreparationStage {
218 Build,
220 RestoreCanisters,
222 ResetNonSnapshotState,
224 DriveToReadiness,
226 ValidateBuilt,
228 ValidateRestored,
230}
231
232#[non_exhaustive]
234#[derive(Clone, Debug, Eq, PartialEq)]
235pub enum RebuildReason {
236 DeadPocketIcTransport,
238 SnapshotRestoreFailure,
240 ResetFailure,
242 ReadinessFailure,
244 ResetCoverageMismatch,
246 InvariantValidationFailure,
248 ExplicitLeaseInvalidation,
250 UnwindWhileLeased,
252 RecipeClassified { code: String },
254}
255
256#[non_exhaustive]
258#[derive(Clone, Debug, Eq, PartialEq)]
259pub enum FailureDisposition {
260 Fatal,
262 Rebuild(RebuildReason),
264}
265
266#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
268pub struct BaselinePoolTimings {
269 wait: Duration,
270 build: Option<Duration>,
271 restore: Option<Duration>,
272 reset: Option<Duration>,
273 readiness: Option<Duration>,
274 validation: Option<Duration>,
275 stale_teardown: Option<Duration>,
276 total: Duration,
277}
278
279#[non_exhaustive]
281#[derive(Clone, Debug, Eq, PartialEq)]
282pub enum BaselinePoolOutcome {
283 Built {
285 slot: usize,
287 timings: BaselinePoolTimings,
289 },
290 Restored {
292 slot: usize,
294 timings: BaselinePoolTimings,
296 },
297 Rebuilt {
299 slot: usize,
301 reason: RebuildReason,
303 timings: BaselinePoolTimings,
305 },
306}
307
308#[non_exhaustive]
310#[derive(Debug)]
311pub enum BaselinePoolPreparationError<E> {
312 Recipe {
314 stage: BaselinePreparationStage,
316 source: E,
318 },
319 Contract(BaselinePoolContractError),
321}
322
323#[non_exhaustive]
325#[derive(Debug)]
326pub enum BaselinePoolError<E> {
327 Preparation {
329 error: BaselinePoolPreparationError<E>,
331 timings: Box<BaselinePoolTimings>,
333 },
334 RecoveryFailed {
336 original: Box<BaselinePoolPreparationError<E>>,
338 rebuild: Box<BaselinePoolPreparationError<E>>,
340 timings: Box<BaselinePoolTimings>,
342 },
343}
344
345pub trait PocketIcBaselineRecipe: Send + Sync + 'static {
347 type Metadata: Send + 'static;
349 type Error: std::error::Error + Send + Sync + 'static;
351
352 fn id(&self) -> &FixtureRecipeId;
354
355 fn reset_requirements(&self) -> &ResetRequirements;
357
358 fn build(&self) -> Result<CachedPocketIcBaseline<Self::Metadata>, Self::Error>;
360
361 fn restore_canisters(
363 &self,
364 baseline: &CachedPocketIcBaseline<Self::Metadata>,
365 ) -> Result<CanisterRestoreReceipt, Self::Error>;
366
367 fn reset_non_snapshot_state(
369 &self,
370 baseline: &CachedPocketIcBaseline<Self::Metadata>,
371 ) -> Result<ResetReceipt, Self::Error>;
372
373 fn drive_to_readiness(
375 &self,
376 baseline: &CachedPocketIcBaseline<Self::Metadata>,
377 ) -> Result<ReadinessReceipt, Self::Error>;
378
379 fn validate(
381 &self,
382 baseline: &CachedPocketIcBaseline<Self::Metadata>,
383 preparation: &PreparedBaseline,
384 ) -> Result<ValidationReceipt, Self::Error>;
385
386 fn classify_failure(
388 &self,
389 stage: BaselinePreparationStage,
390 _error: &Self::Error,
391 ) -> FailureDisposition {
392 FailureDisposition::Rebuild(stage.default_rebuild_reason())
393 }
394}
395
396pub struct CachedPocketIcBaselinePool<R>
408where
409 R: PocketIcBaselineRecipe,
410{
411 recipe: R,
412 slots: BoundedSlotPool<BaselineSlot<R::Metadata>>,
413}
414
415struct BaselineSlot<M> {
416 baseline: CachedPocketIcBaseline<M>,
417 invalidation_reason: Option<RebuildReason>,
418}
419
420pub struct CachedPocketIcBaselinePoolGuard<'a, R>
422where
423 R: PocketIcBaselineRecipe,
424{
425 slot: BoundedSlotLease<'a, BaselineSlot<R::Metadata>>,
426}
427
428impl FixtureRecipeId {
429 pub fn try_new(identity: impl Into<String>) -> Result<Self, BaselinePoolContractError> {
431 let identity = identity.into();
432 if identity.trim().is_empty() {
433 return Err(BaselinePoolContractError::EmptyRecipeIdentity);
434 }
435 Ok(Self(identity))
436 }
437
438 #[must_use]
440 pub fn as_str(&self) -> &str {
441 &self.0
442 }
443}
444
445impl ResetRequirement {
446 #[must_use]
448 pub const fn domain(&self) -> ResetDomainKind {
449 match self {
450 Self::CanisterSnapshots => ResetDomainKind::CanisterSnapshots,
451 Self::CanisterCycles(_) => ResetDomainKind::CanisterCycles,
452 Self::PocketIcTime(_) => ResetDomainKind::PocketIcTime,
453 Self::ExtraCanisters(_) => ResetDomainKind::ExtraCanisters,
454 Self::PendingMessages(_) => ResetDomainKind::PendingMessages,
455 Self::SubnetState(_) => ResetDomainKind::SubnetState,
456 Self::ExternalResources(_) => ResetDomainKind::ExternalResources,
457 }
458 }
459}
460
461impl ResetAchievement {
462 #[must_use]
464 pub const fn domain(&self) -> ResetDomainKind {
465 match self {
466 Self::CanisterSnapshots => ResetDomainKind::CanisterSnapshots,
467 Self::CanisterCycles(_) => ResetDomainKind::CanisterCycles,
468 Self::PocketIcTime(_) => ResetDomainKind::PocketIcTime,
469 Self::ExtraCanisters(_) => ResetDomainKind::ExtraCanisters,
470 Self::PendingMessages(_) => ResetDomainKind::PendingMessages,
471 Self::SubnetState(_) => ResetDomainKind::SubnetState,
472 Self::ExternalResources(_) => ResetDomainKind::ExternalResources,
473 }
474 }
475
476 fn satisfies(&self, requirement: &ResetRequirement) -> bool {
477 match (requirement, self) {
478 (ResetRequirement::CanisterSnapshots, Self::CanisterSnapshots) => true,
479 (ResetRequirement::CanisterCycles(left), Self::CanisterCycles(right)) => left == right,
480 (ResetRequirement::PocketIcTime(left), Self::PocketIcTime(right)) => left == right,
481 (ResetRequirement::ExtraCanisters(left), Self::ExtraCanisters(right)) => left == right,
482 (ResetRequirement::PendingMessages(left), Self::PendingMessages(right))
483 | (ResetRequirement::SubnetState(left), Self::SubnetState(right))
484 | (ResetRequirement::ExternalResources(left), Self::ExternalResources(right)) => {
485 left == right
486 }
487 _ => false,
488 }
489 }
490}
491
492impl ResetRequirements {
493 pub fn try_new<I>(requirements: I) -> Result<Self, BaselinePoolContractError>
498 where
499 I: IntoIterator<Item = ResetRequirement>,
500 {
501 let mut domains = BTreeMap::new();
502 for requirement in requirements {
503 let domain = requirement.domain();
504 if domains.insert(domain, requirement).is_some() {
505 return Err(BaselinePoolContractError::DuplicateResetDomain { domain });
506 }
507 }
508 for domain in [
509 ResetDomainKind::CanisterSnapshots,
510 ResetDomainKind::CanisterCycles,
511 ] {
512 if !domains.contains_key(&domain) {
513 return Err(BaselinePoolContractError::UndeclaredRequiredResetDomain { domain });
514 }
515 }
516 Ok(Self(domains))
517 }
518
519 #[must_use]
521 pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetRequirement> {
522 self.0.get(&domain)
523 }
524
525 pub fn iter(&self) -> impl Iterator<Item = &ResetRequirement> {
527 self.0.values()
528 }
529
530 fn verify(&self, receipt: &ResetReceipt) -> Result<(), BaselinePoolContractError> {
531 for (domain, requirement) in &self.0 {
532 let Some(achievement) = receipt.0.get(domain) else {
533 return Err(BaselinePoolContractError::MissingResetDomain { domain: *domain });
534 };
535 if !achievement.satisfies(requirement) {
536 return Err(BaselinePoolContractError::ResetPolicyMismatch {
537 requirement: requirement.clone(),
538 achievement: achievement.clone(),
539 });
540 }
541 }
542 Ok(())
543 }
544}
545
546impl ResetReceipt {
547 pub fn try_new<I>(achievements: I) -> Result<Self, BaselinePoolContractError>
553 where
554 I: IntoIterator<Item = ResetAchievement>,
555 {
556 let mut domains = BTreeMap::new();
557 for achievement in achievements {
558 let domain = achievement.domain();
559 if domains.insert(domain, achievement).is_some() {
560 return Err(BaselinePoolContractError::DuplicateResetDomain { domain });
561 }
562 }
563 Ok(Self(domains))
564 }
565
566 #[must_use]
568 pub const fn empty() -> Self {
569 Self(BTreeMap::new())
570 }
571
572 #[must_use]
574 pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetAchievement> {
575 self.0.get(&domain)
576 }
577
578 pub fn iter(&self) -> impl Iterator<Item = &ResetAchievement> {
580 self.0.values()
581 }
582
583 fn include_restore(
584 &mut self,
585 restore: &CanisterRestoreReceipt,
586 ) -> Result<(), BaselinePoolContractError> {
587 self.insert(ResetAchievement::CanisterSnapshots)?;
588 self.insert(ResetAchievement::CanisterCycles(restore.cycle_policy))
589 }
590
591 fn insert(&mut self, achievement: ResetAchievement) -> Result<(), BaselinePoolContractError> {
592 let domain = achievement.domain();
593 if self.0.insert(domain, achievement).is_some() {
594 return Err(BaselinePoolContractError::DuplicateResetDomain { domain });
595 }
596 Ok(())
597 }
598}
599
600impl CanisterRestoreReceipt {
601 pub fn try_new<I>(
603 canister_ids: I,
604 cycle_policy: CycleResetPolicy,
605 ) -> Result<Self, BaselinePoolContractError>
606 where
607 I: IntoIterator<Item = Principal>,
608 {
609 let mut unique = BTreeSet::new();
610 for canister_id in canister_ids {
611 if !unique.insert(canister_id) {
612 return Err(BaselinePoolContractError::DuplicateCanisterId { canister_id });
613 }
614 }
615 if unique.is_empty() {
616 return Err(BaselinePoolContractError::EmptyCanisterSet);
617 }
618 Ok(Self {
619 canister_ids: unique.into_iter().collect(),
620 cycle_policy,
621 })
622 }
623
624 pub fn try_from_baseline<M>(
632 baseline: &CachedPocketIcBaseline<M>,
633 cycle_policy: CycleResetPolicy,
634 ) -> Result<Self, BaselinePoolContractError> {
635 Self::try_new(baseline.snapshot_canister_ids(), cycle_policy)
636 }
637
638 #[must_use]
640 pub fn canister_ids(&self) -> &[Principal] {
641 &self.canister_ids
642 }
643
644 #[must_use]
646 pub const fn cycle_policy(&self) -> CycleResetPolicy {
647 self.cycle_policy
648 }
649}
650
651impl ReadinessReceipt {
652 pub fn try_new(identity: impl Into<String>) -> Result<Self, BaselinePoolContractError> {
654 Ok(Self {
655 identity: nonempty_receipt_identity("readiness", identity.into())?,
656 })
657 }
658
659 #[must_use]
661 pub fn identity(&self) -> &str {
662 &self.identity
663 }
664}
665
666impl ValidationReceipt {
667 pub fn try_new(
669 recipe_id: FixtureRecipeId,
670 invariant_identity: impl Into<String>,
671 ) -> Result<Self, BaselinePoolContractError> {
672 Ok(Self {
673 recipe_id,
674 invariant_identity: nonempty_receipt_identity("validation", invariant_identity.into())?,
675 })
676 }
677
678 #[must_use]
680 pub const fn recipe_id(&self) -> &FixtureRecipeId {
681 &self.recipe_id
682 }
683
684 #[must_use]
686 pub fn invariant_identity(&self) -> &str {
687 &self.invariant_identity
688 }
689}
690
691impl BaselinePreparationStage {
692 #[must_use]
697 pub fn default_rebuild_reason(self) -> RebuildReason {
698 match self {
699 Self::RestoreCanisters => RebuildReason::SnapshotRestoreFailure,
700 Self::ResetNonSnapshotState => RebuildReason::ResetFailure,
701 Self::DriveToReadiness => RebuildReason::ReadinessFailure,
702 Self::ValidateRestored | Self::ValidateBuilt => {
703 RebuildReason::InvariantValidationFailure
704 }
705 Self::Build => RebuildReason::RecipeClassified {
706 code: "build".to_owned(),
707 },
708 }
709 }
710}
711
712impl BaselinePoolTimings {
713 #[must_use]
715 pub const fn wait(self) -> Duration {
716 self.wait
717 }
718
719 #[must_use]
721 pub const fn build(self) -> Option<Duration> {
722 self.build
723 }
724
725 #[must_use]
727 pub const fn restore(self) -> Option<Duration> {
728 self.restore
729 }
730
731 #[must_use]
733 pub const fn reset(self) -> Option<Duration> {
734 self.reset
735 }
736
737 #[must_use]
739 pub const fn readiness(self) -> Option<Duration> {
740 self.readiness
741 }
742
743 #[must_use]
745 pub const fn validation(self) -> Option<Duration> {
746 self.validation
747 }
748
749 #[must_use]
751 pub const fn stale_teardown(self) -> Option<Duration> {
752 self.stale_teardown
753 }
754
755 #[must_use]
757 pub const fn total(self) -> Duration {
758 self.total
759 }
760}
761
762impl BaselinePoolOutcome {
763 #[must_use]
765 pub const fn slot(&self) -> usize {
766 match self {
767 Self::Built { slot, .. } | Self::Restored { slot, .. } | Self::Rebuilt { slot, .. } => {
768 *slot
769 }
770 }
771 }
772
773 #[must_use]
775 pub const fn timings(&self) -> BaselinePoolTimings {
776 match self {
777 Self::Built { timings, .. }
778 | Self::Restored { timings, .. }
779 | Self::Rebuilt { timings, .. } => *timings,
780 }
781 }
782}
783
784impl<E> BaselinePoolError<E> {
785 #[must_use]
787 pub const fn timings(&self) -> BaselinePoolTimings {
788 match self {
789 Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
790 }
791 }
792}
793
794impl<R> CachedPocketIcBaselinePool<R>
795where
796 R: PocketIcBaselineRecipe,
797{
798 #[must_use]
800 pub fn new(capacity: NonZeroUsize, recipe: R) -> Self {
801 Self {
802 recipe,
803 slots: BoundedSlotPool::new(capacity),
804 }
805 }
806
807 #[must_use]
809 pub fn recipe_id(&self) -> &FixtureRecipeId {
810 self.recipe.id()
811 }
812
813 #[must_use]
815 pub fn capacity(&self) -> NonZeroUsize {
816 self.slots.capacity()
817 }
818
819 pub fn acquire(
831 &self,
832 ) -> Result<
833 (CachedPocketIcBaselinePoolGuard<'_, R>, BaselinePoolOutcome),
834 BaselinePoolError<R::Error>,
835 > {
836 let total_started = Instant::now();
837 let mut slot = self.slots.acquire();
838 let mut timings = BaselinePoolTimings {
839 wait: slot.wait(),
840 ..BaselinePoolTimings::default()
841 };
842 let slot_index = slot.slot_index();
843
844 if slot.is_reusable() {
845 match self.prepare_reused(&mut slot, &mut timings) {
846 Ok(()) => {
847 timings.total = total_started.elapsed();
848 return Ok((
849 CachedPocketIcBaselinePoolGuard { slot },
850 BaselinePoolOutcome::Restored {
851 slot: slot_index,
852 timings,
853 },
854 ));
855 }
856 Err(original) => {
857 let disposition = self.failure_disposition(&original);
858 match disposition {
859 FailureDisposition::Fatal => {
860 Self::discard_stale_slot(&mut slot, &mut timings);
864 timings.total = total_started.elapsed();
865 return Err(BaselinePoolError::Preparation {
866 error: original,
867 timings: Box::new(timings),
868 });
869 }
870 FailureDisposition::Rebuild(reason) => {
871 Self::discard_stale_slot(&mut slot, &mut timings);
872 if let Err(rebuild) = self.build_slot(&mut slot, &mut timings) {
873 timings.total = total_started.elapsed();
874 return Err(BaselinePoolError::RecoveryFailed {
875 original: Box::new(original),
876 rebuild: Box::new(rebuild),
877 timings: Box::new(timings),
878 });
879 }
880 timings.total = total_started.elapsed();
881 return Ok((
882 CachedPocketIcBaselinePoolGuard { slot },
883 BaselinePoolOutcome::Rebuilt {
884 slot: slot_index,
885 reason,
886 timings,
887 },
888 ));
889 }
890 }
891 }
892 }
893 }
894
895 let rebuild_reason = if slot.invalidated_by_unwind() {
896 Some(RebuildReason::UnwindWhileLeased)
897 } else {
898 slot.get()
899 .and_then(|slot| slot.invalidation_reason.clone())
900 .or_else(|| {
901 slot.is_populated()
902 .then_some(RebuildReason::ExplicitLeaseInvalidation)
903 })
904 };
905 if slot.is_populated() {
906 Self::discard_stale_slot(&mut slot, &mut timings);
907 }
908 if let Err(error) = self.build_slot(&mut slot, &mut timings) {
909 timings.total = total_started.elapsed();
910 return Err(BaselinePoolError::Preparation {
911 error,
912 timings: Box::new(timings),
913 });
914 }
915 timings.total = total_started.elapsed();
916
917 let outcome = rebuild_reason.map_or_else(
918 || BaselinePoolOutcome::Built {
919 slot: slot_index,
920 timings,
921 },
922 |reason| BaselinePoolOutcome::Rebuilt {
923 slot: slot_index,
924 reason,
925 timings,
926 },
927 );
928 Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
929 }
930
931 fn prepare_reused(
932 &self,
933 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
934 timings: &mut BaselinePoolTimings,
935 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
936 let baseline = &slot
937 .get()
938 .expect("reusable baseline slot must be populated")
939 .baseline;
940
941 let started = Instant::now();
942 let restore = self.recipe.restore_canisters(baseline);
943 add_timing(&mut timings.restore, started.elapsed());
944 let canisters = restore.map_err(|source| BaselinePoolPreparationError::Recipe {
945 stage: BaselinePreparationStage::RestoreCanisters,
946 source,
947 })?;
948 let expected_canisters = baseline.snapshot_canister_ids().collect::<Vec<_>>();
949 if canisters.canister_ids() != expected_canisters {
950 return Err(BaselinePoolPreparationError::Contract(
951 BaselinePoolContractError::RestoreCanisterSetMismatch {
952 expected: expected_canisters,
953 actual: canisters.canister_ids().to_vec(),
954 },
955 ));
956 }
957
958 let started = Instant::now();
959 let reset_result = self.recipe.reset_non_snapshot_state(baseline);
960 add_timing(&mut timings.reset, started.elapsed());
961 let mut reset = reset_result.map_err(|source| BaselinePoolPreparationError::Recipe {
962 stage: BaselinePreparationStage::ResetNonSnapshotState,
963 source,
964 })?;
965 reset
966 .include_restore(&canisters)
967 .map_err(BaselinePoolPreparationError::Contract)?;
968
969 let started = Instant::now();
970 let readiness_result = self.recipe.drive_to_readiness(baseline);
971 add_timing(&mut timings.readiness, started.elapsed());
972 let readiness =
973 readiness_result.map_err(|source| BaselinePoolPreparationError::Recipe {
974 stage: BaselinePreparationStage::DriveToReadiness,
975 source,
976 })?;
977 self.recipe
978 .reset_requirements()
979 .verify(&reset)
980 .map_err(BaselinePoolPreparationError::Contract)?;
981
982 let preparation = PreparedBaseline::Restored {
983 canisters,
984 reset,
985 readiness,
986 };
987 self.validate_baseline(
988 baseline,
989 &preparation,
990 BaselinePreparationStage::ValidateRestored,
991 timings,
992 )?;
993 slot.get_mut()
994 .expect("validated baseline slot must remain populated")
995 .invalidation_reason = None;
996 Ok(())
997 }
998
999 fn build_slot(
1000 &self,
1001 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1002 timings: &mut BaselinePoolTimings,
1003 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1004 let started = Instant::now();
1005 let build = self.recipe.build();
1006 add_timing(&mut timings.build, started.elapsed());
1007 let baseline = build.map_err(|source| BaselinePoolPreparationError::Recipe {
1008 stage: BaselinePreparationStage::Build,
1009 source,
1010 })?;
1011
1012 if let Err(error) = self.validate_baseline(
1013 &baseline,
1014 &PreparedBaseline::Built,
1015 BaselinePreparationStage::ValidateBuilt,
1016 timings,
1017 ) {
1018 drop_baseline_safely(baseline);
1019 return Err(error);
1020 }
1021 let replaced = slot.replace(BaselineSlot {
1022 baseline,
1023 invalidation_reason: None,
1024 });
1025 debug_assert!(replaced.is_none());
1026 Ok(())
1027 }
1028
1029 fn validate_baseline(
1030 &self,
1031 baseline: &CachedPocketIcBaseline<R::Metadata>,
1032 preparation: &PreparedBaseline,
1033 stage: BaselinePreparationStage,
1034 timings: &mut BaselinePoolTimings,
1035 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1036 let started = Instant::now();
1037 let validation = self.recipe.validate(baseline, preparation);
1038 add_timing(&mut timings.validation, started.elapsed());
1039 let receipt =
1040 validation.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })?;
1041 if receipt.recipe_id() != self.recipe.id() {
1042 return Err(BaselinePoolPreparationError::Contract(
1043 BaselinePoolContractError::RecipeIdentityMismatch {
1044 expected: self.recipe.id().clone(),
1045 actual: receipt.recipe_id().clone(),
1046 },
1047 ));
1048 }
1049 Ok(())
1050 }
1051
1052 fn failure_disposition(
1053 &self,
1054 error: &BaselinePoolPreparationError<R::Error>,
1055 ) -> FailureDisposition {
1056 match error {
1057 BaselinePoolPreparationError::Recipe { stage, source } => {
1058 self.recipe.classify_failure(*stage, source)
1059 }
1060 BaselinePoolPreparationError::Contract(
1061 BaselinePoolContractError::RecipeIdentityMismatch { .. },
1062 ) => FailureDisposition::Fatal,
1063 BaselinePoolPreparationError::Contract(_) => {
1064 FailureDisposition::Rebuild(rebuild_reason_for_error(error))
1065 }
1066 }
1067 }
1068
1069 fn discard_stale_slot(
1070 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1071 timings: &mut BaselinePoolTimings,
1072 ) {
1073 let started = Instant::now();
1074 if let Some(stale) = slot.take() {
1075 drop_baseline_safely(stale.baseline);
1076 }
1077 timings.stale_teardown = Some(started.elapsed());
1078 }
1079}
1080
1081impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
1082where
1083 R: PocketIcBaselineRecipe,
1084{
1085 #[must_use]
1087 pub const fn slot(&self) -> usize {
1088 self.slot.slot_index()
1089 }
1090
1091 pub fn invalidate(&mut self, reason: RebuildReason) {
1093 if let Some(slot) = self.slot.get_mut() {
1094 slot.invalidation_reason = Some(reason);
1095 }
1096 self.slot.invalidate();
1097 }
1098}
1099
1100impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
1101where
1102 R: PocketIcBaselineRecipe,
1103{
1104 type Target = CachedPocketIcBaseline<R::Metadata>;
1105
1106 fn deref(&self) -> &Self::Target {
1107 &self
1108 .slot
1109 .get()
1110 .expect("leased baseline pool slot must be populated")
1111 .baseline
1112 }
1113}
1114
1115fn nonempty_receipt_identity(
1116 receipt: &'static str,
1117 identity: String,
1118) -> Result<String, BaselinePoolContractError> {
1119 if identity.trim().is_empty() {
1120 return Err(BaselinePoolContractError::EmptyReceiptIdentity { receipt });
1121 }
1122 Ok(identity)
1123}
1124
1125fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
1126 *total = Some(total.unwrap_or_default().saturating_add(elapsed));
1127}
1128
1129fn rebuild_reason_for_error<E>(error: &BaselinePoolPreparationError<E>) -> RebuildReason {
1130 match error {
1131 BaselinePoolPreparationError::Recipe { stage, .. } => stage.default_rebuild_reason(),
1132 BaselinePoolPreparationError::Contract(
1133 BaselinePoolContractError::MissingResetDomain { .. }
1134 | BaselinePoolContractError::ResetPolicyMismatch { .. }
1135 | BaselinePoolContractError::DuplicateResetDomain { .. }
1136 | BaselinePoolContractError::RestoreCanisterSetMismatch { .. },
1137 ) => RebuildReason::ResetCoverageMismatch,
1138 BaselinePoolPreparationError::Contract(_) => RebuildReason::InvariantValidationFailure,
1139 }
1140}
1141
1142fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
1143 let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
1144}
1145
1146impl std::fmt::Display for FixtureRecipeId {
1147 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1148 formatter.write_str(&self.0)
1149 }
1150}
1151
1152impl std::fmt::Display for BaselinePreparationStage {
1153 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1154 formatter.write_str(match self {
1155 Self::Build => "build",
1156 Self::RestoreCanisters => "canister restore",
1157 Self::ResetNonSnapshotState => "non-snapshot reset",
1158 Self::DriveToReadiness => "readiness",
1159 Self::ValidateBuilt => "built-baseline validation",
1160 Self::ValidateRestored => "restored-baseline validation",
1161 })
1162 }
1163}
1164
1165impl std::fmt::Display for BaselinePoolContractError {
1166 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1167 match self {
1168 Self::EmptyRecipeIdentity => formatter.write_str("fixture recipe identity is empty"),
1169 Self::EmptyReceiptIdentity { receipt } => {
1170 write!(formatter, "{receipt} receipt identity is empty")
1171 }
1172 Self::DuplicateResetDomain { domain } => {
1173 write!(
1174 formatter,
1175 "reset domain {domain:?} was reported more than once"
1176 )
1177 }
1178 Self::DuplicateCanisterId { canister_id } => {
1179 write!(formatter, "restore receipt repeats canister {canister_id}")
1180 }
1181 Self::EmptyCanisterSet => formatter.write_str("restore receipt contains no canisters"),
1182 Self::UndeclaredRequiredResetDomain { domain } => write!(
1183 formatter,
1184 "baseline recipe does not declare required reset domain {domain:?}",
1185 ),
1186 Self::RestoreCanisterSetMismatch { expected, actual } => write!(
1187 formatter,
1188 "restore receipt identified canisters {actual:?}, expected captured set {expected:?}",
1189 ),
1190 Self::MissingResetDomain { domain } => {
1191 write!(
1192 formatter,
1193 "required reset domain {domain:?} was not achieved"
1194 )
1195 }
1196 Self::ResetPolicyMismatch {
1197 requirement,
1198 achievement,
1199 } => write!(
1200 formatter,
1201 "reset achievement {achievement:?} does not satisfy {requirement:?}",
1202 ),
1203 Self::RecipeIdentityMismatch { expected, actual } => write!(
1204 formatter,
1205 "validation receipt used recipe `{actual}` instead of `{expected}`",
1206 ),
1207 }
1208 }
1209}
1210
1211impl std::error::Error for BaselinePoolContractError {}
1212
1213impl<E> std::fmt::Display for BaselinePoolPreparationError<E>
1214where
1215 E: std::fmt::Display,
1216{
1217 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1218 match self {
1219 Self::Recipe { stage, source } => {
1220 write!(formatter, "baseline {stage} failed: {source}")
1221 }
1222 Self::Contract(error) => write!(formatter, "baseline pool contract failed: {error}"),
1223 }
1224 }
1225}
1226
1227impl<E> std::error::Error for BaselinePoolPreparationError<E>
1228where
1229 E: std::error::Error + 'static,
1230{
1231 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1232 match self {
1233 Self::Recipe { source, .. } => Some(source),
1234 Self::Contract(error) => Some(error),
1235 }
1236 }
1237}
1238
1239impl<E> std::fmt::Display for BaselinePoolError<E>
1240where
1241 E: std::fmt::Display,
1242{
1243 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1244 match self {
1245 Self::Preparation { error, .. } => error.fmt(formatter),
1246 Self::RecoveryFailed {
1247 original, rebuild, ..
1248 } => write!(
1249 formatter,
1250 "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
1251 ),
1252 }
1253 }
1254}
1255
1256impl<E> std::error::Error for BaselinePoolError<E>
1257where
1258 E: std::error::Error + 'static,
1259{
1260 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1261 match self {
1262 Self::Preparation { error, .. } => Some(error),
1263 Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
1264 }
1265 }
1266}
1267
1268#[cfg(test)]
1269mod tests {
1270 use super::{
1271 BaselinePoolContractError, CycleResetPolicy, FixtureRecipeId, ResetAchievement,
1272 ResetDomainKind, ResetReceipt, ResetRequirement, ResetRequirements,
1273 };
1274
1275 #[test]
1276 fn recipe_identity_must_be_nonempty() {
1277 assert!(matches!(
1278 FixtureRecipeId::try_new(" "),
1279 Err(BaselinePoolContractError::EmptyRecipeIdentity)
1280 ));
1281 }
1282
1283 #[test]
1284 fn reset_requirements_reject_duplicate_domains() {
1285 let result = ResetRequirements::try_new([
1286 ResetRequirement::CanisterCycles(CycleResetPolicy::PreserveCurrent),
1287 ResetRequirement::CanisterCycles(CycleResetPolicy::RestoreExactBaseline),
1288 ]);
1289 assert!(matches!(
1290 result,
1291 Err(BaselinePoolContractError::DuplicateResetDomain {
1292 domain: ResetDomainKind::CanisterCycles,
1293 })
1294 ));
1295 }
1296
1297 #[test]
1298 fn required_policy_must_match_achieved_policy() {
1299 let requirements = ResetRequirements::try_new([
1300 ResetRequirement::CanisterSnapshots,
1301 ResetRequirement::CanisterCycles(CycleResetPolicy::RestoreExactBaseline),
1302 ])
1303 .unwrap();
1304 let receipt = ResetReceipt::try_new([
1305 ResetAchievement::CanisterSnapshots,
1306 ResetAchievement::CanisterCycles(CycleResetPolicy::PreserveCurrent),
1307 ])
1308 .unwrap();
1309
1310 assert!(matches!(
1311 requirements.verify(&receipt),
1312 Err(BaselinePoolContractError::ResetPolicyMismatch { .. })
1313 ));
1314 }
1315}