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