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 std::fmt::Display for BaselinePoolTimings {
785 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
786 write!(
787 formatter,
788 "total={:?} wait={:?} build={:?} restore={:?} reset={:?} readiness={:?} validation={:?} stale_teardown={:?}",
789 self.total,
790 self.wait,
791 self.build,
792 self.restore,
793 self.reset,
794 self.readiness,
795 self.validation,
796 self.stale_teardown,
797 )
798 }
799}
800
801impl std::fmt::Display for BaselinePoolOutcome {
802 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
803 match self {
804 Self::Built { slot, timings } => write!(formatter, "built slot={slot} {timings}"),
805 Self::Restored { slot, timings } => {
806 write!(formatter, "restored slot={slot} {timings}")
807 }
808 Self::Rebuilt {
809 slot,
810 reason,
811 timings,
812 } => write!(formatter, "rebuilt slot={slot} reason={reason:?} {timings}"),
813 }
814 }
815}
816
817impl<E> BaselinePoolError<E> {
818 #[must_use]
820 pub const fn timings(&self) -> BaselinePoolTimings {
821 match self {
822 Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
823 }
824 }
825}
826
827impl<R> CachedPocketIcBaselinePool<R>
828where
829 R: PocketIcBaselineRecipe,
830{
831 #[must_use]
833 pub fn new(capacity: NonZeroUsize, recipe: R) -> Self {
834 Self {
835 recipe,
836 slots: BoundedSlotPool::new(capacity),
837 }
838 }
839
840 #[must_use]
842 pub fn recipe_id(&self) -> &FixtureRecipeId {
843 self.recipe.id()
844 }
845
846 #[must_use]
848 pub fn capacity(&self) -> NonZeroUsize {
849 self.slots.capacity()
850 }
851
852 pub fn acquire(
864 &self,
865 ) -> Result<
866 (CachedPocketIcBaselinePoolGuard<'_, R>, BaselinePoolOutcome),
867 BaselinePoolError<R::Error>,
868 > {
869 let total_started = Instant::now();
870 let mut slot = self.slots.acquire();
871 let mut timings = BaselinePoolTimings {
872 wait: slot.wait(),
873 ..BaselinePoolTimings::default()
874 };
875 let slot_index = slot.slot_index();
876
877 if slot.is_reusable() {
878 match self.prepare_reused(&mut slot, &mut timings) {
879 Ok(()) => {
880 timings.total = total_started.elapsed();
881 return Ok((
882 CachedPocketIcBaselinePoolGuard { slot },
883 BaselinePoolOutcome::Restored {
884 slot: slot_index,
885 timings,
886 },
887 ));
888 }
889 Err(original) => {
890 let disposition = self.failure_disposition(&original);
891 match disposition {
892 FailureDisposition::Fatal => {
893 Self::discard_stale_slot(&mut slot, &mut timings);
897 timings.total = total_started.elapsed();
898 return Err(BaselinePoolError::Preparation {
899 error: original,
900 timings: Box::new(timings),
901 });
902 }
903 FailureDisposition::Rebuild(reason) => {
904 Self::discard_stale_slot(&mut slot, &mut timings);
905 if let Err(rebuild) = self.build_slot(&mut slot, &mut timings) {
906 timings.total = total_started.elapsed();
907 return Err(BaselinePoolError::RecoveryFailed {
908 original: Box::new(original),
909 rebuild: Box::new(rebuild),
910 timings: Box::new(timings),
911 });
912 }
913 timings.total = total_started.elapsed();
914 return Ok((
915 CachedPocketIcBaselinePoolGuard { slot },
916 BaselinePoolOutcome::Rebuilt {
917 slot: slot_index,
918 reason,
919 timings,
920 },
921 ));
922 }
923 }
924 }
925 }
926 }
927
928 let rebuild_reason = if slot.invalidated_by_unwind() {
929 Some(RebuildReason::UnwindWhileLeased)
930 } else {
931 slot.get()
932 .and_then(|slot| slot.invalidation_reason.clone())
933 .or_else(|| {
934 slot.is_populated()
935 .then_some(RebuildReason::ExplicitLeaseInvalidation)
936 })
937 };
938 if slot.is_populated() {
939 Self::discard_stale_slot(&mut slot, &mut timings);
940 }
941 if let Err(error) = self.build_slot(&mut slot, &mut timings) {
942 timings.total = total_started.elapsed();
943 return Err(BaselinePoolError::Preparation {
944 error,
945 timings: Box::new(timings),
946 });
947 }
948 timings.total = total_started.elapsed();
949
950 let outcome = rebuild_reason.map_or_else(
951 || BaselinePoolOutcome::Built {
952 slot: slot_index,
953 timings,
954 },
955 |reason| BaselinePoolOutcome::Rebuilt {
956 slot: slot_index,
957 reason,
958 timings,
959 },
960 );
961 Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
962 }
963
964 fn prepare_reused(
965 &self,
966 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
967 timings: &mut BaselinePoolTimings,
968 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
969 let baseline = &slot
970 .get()
971 .expect("reusable baseline slot must be populated")
972 .baseline;
973
974 let started = Instant::now();
975 let restore = self.recipe.restore_canisters(baseline);
976 add_timing(&mut timings.restore, started.elapsed());
977 let canisters = restore.map_err(|source| BaselinePoolPreparationError::Recipe {
978 stage: BaselinePreparationStage::RestoreCanisters,
979 source,
980 })?;
981 let expected_canisters = baseline.snapshot_canister_ids().collect::<Vec<_>>();
982 if canisters.canister_ids() != expected_canisters {
983 return Err(BaselinePoolPreparationError::Contract(
984 BaselinePoolContractError::RestoreCanisterSetMismatch {
985 expected: expected_canisters,
986 actual: canisters.canister_ids().to_vec(),
987 },
988 ));
989 }
990
991 let started = Instant::now();
992 let reset_result = self.recipe.reset_non_snapshot_state(baseline);
993 add_timing(&mut timings.reset, started.elapsed());
994 let mut reset = reset_result.map_err(|source| BaselinePoolPreparationError::Recipe {
995 stage: BaselinePreparationStage::ResetNonSnapshotState,
996 source,
997 })?;
998 reset
999 .include_restore(&canisters)
1000 .map_err(BaselinePoolPreparationError::Contract)?;
1001
1002 let started = Instant::now();
1003 let readiness_result = self.recipe.drive_to_readiness(baseline);
1004 add_timing(&mut timings.readiness, started.elapsed());
1005 let readiness =
1006 readiness_result.map_err(|source| BaselinePoolPreparationError::Recipe {
1007 stage: BaselinePreparationStage::DriveToReadiness,
1008 source,
1009 })?;
1010 self.recipe
1011 .reset_requirements()
1012 .verify(&reset)
1013 .map_err(BaselinePoolPreparationError::Contract)?;
1014
1015 let preparation = PreparedBaseline::Restored {
1016 canisters,
1017 reset,
1018 readiness,
1019 };
1020 self.validate_baseline(
1021 baseline,
1022 &preparation,
1023 BaselinePreparationStage::ValidateRestored,
1024 timings,
1025 )?;
1026 slot.get_mut()
1027 .expect("validated baseline slot must remain populated")
1028 .invalidation_reason = None;
1029 Ok(())
1030 }
1031
1032 fn build_slot(
1033 &self,
1034 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1035 timings: &mut BaselinePoolTimings,
1036 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1037 let started = Instant::now();
1038 let build = self.recipe.build();
1039 add_timing(&mut timings.build, started.elapsed());
1040 let baseline = build.map_err(|source| BaselinePoolPreparationError::Recipe {
1041 stage: BaselinePreparationStage::Build,
1042 source,
1043 })?;
1044
1045 if let Err(error) = self.validate_baseline(
1046 &baseline,
1047 &PreparedBaseline::Built,
1048 BaselinePreparationStage::ValidateBuilt,
1049 timings,
1050 ) {
1051 drop_baseline_safely(baseline);
1052 return Err(error);
1053 }
1054 let replaced = slot.replace(BaselineSlot {
1055 baseline,
1056 invalidation_reason: None,
1057 });
1058 debug_assert!(replaced.is_none());
1059 Ok(())
1060 }
1061
1062 fn validate_baseline(
1063 &self,
1064 baseline: &CachedPocketIcBaseline<R::Metadata>,
1065 preparation: &PreparedBaseline,
1066 stage: BaselinePreparationStage,
1067 timings: &mut BaselinePoolTimings,
1068 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1069 let started = Instant::now();
1070 let validation = self.recipe.validate(baseline, preparation);
1071 add_timing(&mut timings.validation, started.elapsed());
1072 let receipt =
1073 validation.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })?;
1074 if receipt.recipe_id() != self.recipe.id() {
1075 return Err(BaselinePoolPreparationError::Contract(
1076 BaselinePoolContractError::RecipeIdentityMismatch {
1077 expected: self.recipe.id().clone(),
1078 actual: receipt.recipe_id().clone(),
1079 },
1080 ));
1081 }
1082 Ok(())
1083 }
1084
1085 fn failure_disposition(
1086 &self,
1087 error: &BaselinePoolPreparationError<R::Error>,
1088 ) -> FailureDisposition {
1089 match error {
1090 BaselinePoolPreparationError::Recipe { stage, source } => {
1091 self.recipe.classify_failure(*stage, source)
1092 }
1093 BaselinePoolPreparationError::Contract(
1094 BaselinePoolContractError::RecipeIdentityMismatch { .. },
1095 ) => FailureDisposition::Fatal,
1096 BaselinePoolPreparationError::Contract(_) => {
1097 FailureDisposition::Rebuild(rebuild_reason_for_error(error))
1098 }
1099 }
1100 }
1101
1102 fn discard_stale_slot(
1103 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1104 timings: &mut BaselinePoolTimings,
1105 ) {
1106 let started = Instant::now();
1107 if let Some(stale) = slot.take() {
1108 drop_baseline_safely(stale.baseline);
1109 }
1110 timings.stale_teardown = Some(started.elapsed());
1111 }
1112}
1113
1114impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
1115where
1116 R: PocketIcBaselineRecipe,
1117{
1118 #[must_use]
1120 pub const fn slot(&self) -> usize {
1121 self.slot.slot_index()
1122 }
1123
1124 pub fn invalidate(&mut self, reason: RebuildReason) {
1126 if let Some(slot) = self.slot.get_mut() {
1127 slot.invalidation_reason = Some(reason);
1128 }
1129 self.slot.invalidate();
1130 }
1131}
1132
1133impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
1134where
1135 R: PocketIcBaselineRecipe,
1136{
1137 type Target = CachedPocketIcBaseline<R::Metadata>;
1138
1139 fn deref(&self) -> &Self::Target {
1140 &self
1141 .slot
1142 .get()
1143 .expect("leased baseline pool slot must be populated")
1144 .baseline
1145 }
1146}
1147
1148fn nonempty_receipt_identity(
1149 receipt: &'static str,
1150 identity: String,
1151) -> Result<String, BaselinePoolContractError> {
1152 if identity.trim().is_empty() {
1153 return Err(BaselinePoolContractError::EmptyReceiptIdentity { receipt });
1154 }
1155 Ok(identity)
1156}
1157
1158fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
1159 *total = Some(total.unwrap_or_default().saturating_add(elapsed));
1160}
1161
1162fn rebuild_reason_for_error<E>(error: &BaselinePoolPreparationError<E>) -> RebuildReason {
1163 match error {
1164 BaselinePoolPreparationError::Recipe { stage, .. } => stage.default_rebuild_reason(),
1165 BaselinePoolPreparationError::Contract(
1166 BaselinePoolContractError::MissingResetDomain { .. }
1167 | BaselinePoolContractError::ResetPolicyMismatch { .. }
1168 | BaselinePoolContractError::DuplicateResetDomain { .. }
1169 | BaselinePoolContractError::RestoreCanisterSetMismatch { .. },
1170 ) => RebuildReason::ResetCoverageMismatch,
1171 BaselinePoolPreparationError::Contract(_) => RebuildReason::InvariantValidationFailure,
1172 }
1173}
1174
1175fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
1176 let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
1177}
1178
1179impl std::fmt::Display for FixtureRecipeId {
1180 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1181 formatter.write_str(&self.0)
1182 }
1183}
1184
1185impl std::fmt::Display for BaselinePreparationStage {
1186 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1187 formatter.write_str(match self {
1188 Self::Build => "build",
1189 Self::RestoreCanisters => "canister restore",
1190 Self::ResetNonSnapshotState => "non-snapshot reset",
1191 Self::DriveToReadiness => "readiness",
1192 Self::ValidateBuilt => "built-baseline validation",
1193 Self::ValidateRestored => "restored-baseline validation",
1194 })
1195 }
1196}
1197
1198impl std::fmt::Display for BaselinePoolContractError {
1199 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1200 match self {
1201 Self::EmptyRecipeIdentity => formatter.write_str("fixture recipe identity is empty"),
1202 Self::EmptyReceiptIdentity { receipt } => {
1203 write!(formatter, "{receipt} receipt identity is empty")
1204 }
1205 Self::DuplicateResetDomain { domain } => {
1206 write!(
1207 formatter,
1208 "reset domain {domain:?} was reported more than once"
1209 )
1210 }
1211 Self::DuplicateCanisterId { canister_id } => {
1212 write!(formatter, "restore receipt repeats canister {canister_id}")
1213 }
1214 Self::EmptyCanisterSet => formatter.write_str("restore receipt contains no canisters"),
1215 Self::UndeclaredRequiredResetDomain { domain } => write!(
1216 formatter,
1217 "baseline recipe does not declare required reset domain {domain:?}",
1218 ),
1219 Self::RestoreCanisterSetMismatch { expected, actual } => write!(
1220 formatter,
1221 "restore receipt identified canisters {actual:?}, expected captured set {expected:?}",
1222 ),
1223 Self::MissingResetDomain { domain } => {
1224 write!(
1225 formatter,
1226 "required reset domain {domain:?} was not achieved"
1227 )
1228 }
1229 Self::ResetPolicyMismatch {
1230 requirement,
1231 achievement,
1232 } => write!(
1233 formatter,
1234 "reset achievement {achievement:?} does not satisfy {requirement:?}",
1235 ),
1236 Self::RecipeIdentityMismatch { expected, actual } => write!(
1237 formatter,
1238 "validation receipt used recipe `{actual}` instead of `{expected}`",
1239 ),
1240 }
1241 }
1242}
1243
1244impl std::error::Error for BaselinePoolContractError {}
1245
1246impl<E> std::fmt::Display for BaselinePoolPreparationError<E>
1247where
1248 E: std::fmt::Display,
1249{
1250 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1251 match self {
1252 Self::Recipe { stage, source } => {
1253 write!(formatter, "baseline {stage} failed: {source}")
1254 }
1255 Self::Contract(error) => write!(formatter, "baseline pool contract failed: {error}"),
1256 }
1257 }
1258}
1259
1260impl<E> std::error::Error for BaselinePoolPreparationError<E>
1261where
1262 E: std::error::Error + 'static,
1263{
1264 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1265 match self {
1266 Self::Recipe { source, .. } => Some(source),
1267 Self::Contract(error) => Some(error),
1268 }
1269 }
1270}
1271
1272impl<E> std::fmt::Display for BaselinePoolError<E>
1273where
1274 E: std::fmt::Display,
1275{
1276 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1277 match self {
1278 Self::Preparation { error, .. } => error.fmt(formatter),
1279 Self::RecoveryFailed {
1280 original, rebuild, ..
1281 } => write!(
1282 formatter,
1283 "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
1284 ),
1285 }
1286 }
1287}
1288
1289impl<E> std::error::Error for BaselinePoolError<E>
1290where
1291 E: std::error::Error + 'static,
1292{
1293 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1294 match self {
1295 Self::Preparation { error, .. } => Some(error),
1296 Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
1297 }
1298 }
1299}
1300
1301#[cfg(test)]
1302mod tests {
1303 use super::{
1304 BaselinePoolContractError, CycleResetPolicy, FixtureRecipeId, ResetAchievement,
1305 ResetDomainKind, ResetReceipt, ResetRequirement, ResetRequirements,
1306 };
1307
1308 #[test]
1309 fn recipe_identity_must_be_nonempty() {
1310 assert!(matches!(
1311 FixtureRecipeId::try_new(" "),
1312 Err(BaselinePoolContractError::EmptyRecipeIdentity)
1313 ));
1314 }
1315
1316 #[test]
1317 fn reset_requirements_reject_duplicate_domains() {
1318 let result = ResetRequirements::try_new([
1319 ResetRequirement::CanisterCycles(CycleResetPolicy::PreserveCurrent),
1320 ResetRequirement::CanisterCycles(CycleResetPolicy::RestoreExactBaseline),
1321 ]);
1322 assert!(matches!(
1323 result,
1324 Err(BaselinePoolContractError::DuplicateResetDomain {
1325 domain: ResetDomainKind::CanisterCycles,
1326 })
1327 ));
1328 }
1329
1330 #[test]
1331 fn required_policy_must_match_achieved_policy() {
1332 let requirements = ResetRequirements::try_new([
1333 ResetRequirement::CanisterSnapshots,
1334 ResetRequirement::CanisterCycles(CycleResetPolicy::RestoreExactBaseline),
1335 ])
1336 .unwrap();
1337 let receipt = ResetReceipt::try_new([
1338 ResetAchievement::CanisterSnapshots,
1339 ResetAchievement::CanisterCycles(CycleResetPolicy::PreserveCurrent),
1340 ])
1341 .unwrap();
1342
1343 assert!(matches!(
1344 requirements.verify(&receipt),
1345 Err(BaselinePoolContractError::ResetPolicyMismatch { .. })
1346 ));
1347 }
1348}