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 if baseline.snapshot_count() == 0 {
623 return Err(BaselinePoolContractError::EmptyCanisterSet);
624 }
625 Ok(Self {
628 canister_ids: baseline.snapshot_canister_ids().collect(),
629 cycle_policy,
630 })
631 }
632
633 #[must_use]
635 pub fn canister_ids(&self) -> &[Principal] {
636 &self.canister_ids
637 }
638
639 #[must_use]
641 pub const fn cycle_policy(&self) -> CycleResetPolicy {
642 self.cycle_policy
643 }
644}
645
646impl ReadinessReceipt {
647 pub fn try_new(identity: impl Into<String>) -> Result<Self, BaselinePoolContractError> {
649 Ok(Self {
650 identity: nonempty_receipt_identity("readiness", identity.into())?,
651 })
652 }
653
654 #[must_use]
656 pub fn identity(&self) -> &str {
657 &self.identity
658 }
659}
660
661impl ValidationReceipt {
662 pub fn try_new(
664 recipe_id: FixtureRecipeId,
665 invariant_identity: impl Into<String>,
666 ) -> Result<Self, BaselinePoolContractError> {
667 Ok(Self {
668 recipe_id,
669 invariant_identity: nonempty_receipt_identity("validation", invariant_identity.into())?,
670 })
671 }
672
673 #[must_use]
675 pub const fn recipe_id(&self) -> &FixtureRecipeId {
676 &self.recipe_id
677 }
678
679 #[must_use]
681 pub fn invariant_identity(&self) -> &str {
682 &self.invariant_identity
683 }
684}
685
686impl BaselinePreparationStage {
687 #[must_use]
692 pub fn default_rebuild_reason(self) -> RebuildReason {
693 match self {
694 Self::RestoreCanisters => RebuildReason::SnapshotRestoreFailure,
695 Self::ResetNonSnapshotState => RebuildReason::ResetFailure,
696 Self::DriveToReadiness => RebuildReason::ReadinessFailure,
697 Self::ValidateRestored | Self::ValidateBuilt => {
698 RebuildReason::InvariantValidationFailure
699 }
700 Self::Build => RebuildReason::RecipeClassified {
701 code: "build".to_owned(),
702 },
703 }
704 }
705}
706
707impl BaselinePoolTimings {
708 #[must_use]
710 pub const fn wait(self) -> Duration {
711 self.wait
712 }
713
714 #[must_use]
716 pub const fn build(self) -> Option<Duration> {
717 self.build
718 }
719
720 #[must_use]
722 pub const fn restore(self) -> Option<Duration> {
723 self.restore
724 }
725
726 #[must_use]
728 pub const fn reset(self) -> Option<Duration> {
729 self.reset
730 }
731
732 #[must_use]
734 pub const fn readiness(self) -> Option<Duration> {
735 self.readiness
736 }
737
738 #[must_use]
740 pub const fn validation(self) -> Option<Duration> {
741 self.validation
742 }
743
744 #[must_use]
746 pub const fn stale_teardown(self) -> Option<Duration> {
747 self.stale_teardown
748 }
749
750 #[must_use]
752 pub const fn total(self) -> Duration {
753 self.total
754 }
755}
756
757impl BaselinePoolOutcome {
758 #[must_use]
760 pub const fn slot(&self) -> usize {
761 match self {
762 Self::Built { slot, .. } | Self::Restored { slot, .. } | Self::Rebuilt { slot, .. } => {
763 *slot
764 }
765 }
766 }
767
768 #[must_use]
770 pub const fn timings(&self) -> BaselinePoolTimings {
771 match self {
772 Self::Built { timings, .. }
773 | Self::Restored { timings, .. }
774 | Self::Rebuilt { timings, .. } => *timings,
775 }
776 }
777}
778
779impl std::fmt::Display for BaselinePoolTimings {
780 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
781 write!(
782 formatter,
783 "total={:?} wait={:?} build={:?} restore={:?} reset={:?} readiness={:?} validation={:?} stale_teardown={:?}",
784 self.total,
785 self.wait,
786 self.build,
787 self.restore,
788 self.reset,
789 self.readiness,
790 self.validation,
791 self.stale_teardown,
792 )
793 }
794}
795
796impl std::fmt::Display for BaselinePoolOutcome {
797 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
798 match self {
799 Self::Built { slot, timings } => write!(formatter, "built slot={slot} {timings}"),
800 Self::Restored { slot, timings } => {
801 write!(formatter, "restored slot={slot} {timings}")
802 }
803 Self::Rebuilt {
804 slot,
805 reason,
806 timings,
807 } => write!(formatter, "rebuilt slot={slot} reason={reason:?} {timings}"),
808 }
809 }
810}
811
812impl<E> BaselinePoolError<E> {
813 #[must_use]
815 pub const fn timings(&self) -> BaselinePoolTimings {
816 match self {
817 Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
818 }
819 }
820}
821
822impl<R> CachedPocketIcBaselinePool<R>
823where
824 R: PocketIcBaselineRecipe,
825{
826 #[must_use]
828 pub fn new(capacity: NonZeroUsize, recipe: R) -> Self {
829 Self {
830 recipe,
831 slots: BoundedSlotPool::new(capacity),
832 }
833 }
834
835 #[must_use]
837 pub fn recipe_id(&self) -> &FixtureRecipeId {
838 self.recipe.id()
839 }
840
841 #[must_use]
843 pub fn capacity(&self) -> NonZeroUsize {
844 self.slots.capacity()
845 }
846
847 pub fn acquire(
859 &self,
860 ) -> Result<
861 (CachedPocketIcBaselinePoolGuard<'_, R>, BaselinePoolOutcome),
862 BaselinePoolError<R::Error>,
863 > {
864 let total_started = Instant::now();
865 let mut slot = self.slots.acquire();
866 let mut timings = BaselinePoolTimings {
867 wait: slot.wait(),
868 ..BaselinePoolTimings::default()
869 };
870 let slot_index = slot.slot_index();
871
872 if slot.is_reusable() {
873 match self.prepare_reused(&slot, &mut timings) {
874 Ok(()) => {
875 timings.total = total_started.elapsed();
876 return Ok((
877 CachedPocketIcBaselinePoolGuard { slot },
878 BaselinePoolOutcome::Restored {
879 slot: slot_index,
880 timings,
881 },
882 ));
883 }
884 Err(original) => {
885 let disposition = self.failure_disposition(&original);
886 match disposition {
887 FailureDisposition::Fatal => {
888 Self::discard_stale_slot(&mut slot, &mut timings);
892 timings.total = total_started.elapsed();
893 return Err(BaselinePoolError::Preparation {
894 error: original,
895 timings: Box::new(timings),
896 });
897 }
898 FailureDisposition::Rebuild(reason) => {
899 Self::discard_stale_slot(&mut slot, &mut timings);
900 if let Err(rebuild) = self.build_slot(&mut slot, &mut timings) {
901 timings.total = total_started.elapsed();
902 return Err(BaselinePoolError::RecoveryFailed {
903 original: Box::new(original),
904 rebuild: Box::new(rebuild),
905 timings: Box::new(timings),
906 });
907 }
908 timings.total = total_started.elapsed();
909 return Ok((
910 CachedPocketIcBaselinePoolGuard { slot },
911 BaselinePoolOutcome::Rebuilt {
912 slot: slot_index,
913 reason,
914 timings,
915 },
916 ));
917 }
918 }
919 }
920 }
921 }
922
923 let rebuild_reason = if slot.invalidated_by_unwind() {
924 Some(RebuildReason::UnwindWhileLeased)
925 } else {
926 slot.get()
927 .and_then(|slot| slot.invalidation_reason.clone())
928 .or_else(|| {
929 slot.is_populated()
930 .then_some(RebuildReason::ExplicitLeaseInvalidation)
931 })
932 };
933 if slot.is_populated() {
934 Self::discard_stale_slot(&mut slot, &mut timings);
935 }
936 if let Err(error) = self.build_slot(&mut slot, &mut timings) {
937 timings.total = total_started.elapsed();
938 return Err(BaselinePoolError::Preparation {
939 error,
940 timings: Box::new(timings),
941 });
942 }
943 timings.total = total_started.elapsed();
944
945 let outcome = rebuild_reason.map_or_else(
946 || BaselinePoolOutcome::Built {
947 slot: slot_index,
948 timings,
949 },
950 |reason| BaselinePoolOutcome::Rebuilt {
951 slot: slot_index,
952 reason,
953 timings,
954 },
955 );
956 Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
957 }
958
959 fn prepare_reused(
960 &self,
961 slot: &BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
962 timings: &mut BaselinePoolTimings,
963 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
964 let baseline = &slot
965 .get()
966 .expect("reusable baseline slot must be populated")
967 .baseline;
968
969 let started = Instant::now();
970 let restore = self.recipe.restore_canisters(baseline);
971 add_timing(&mut timings.restore, started.elapsed());
972 let canisters = restore.map_err(|source| BaselinePoolPreparationError::Recipe {
973 stage: BaselinePreparationStage::RestoreCanisters,
974 source,
975 })?;
976 if !baseline
977 .snapshot_canister_ids()
978 .eq(canisters.canister_ids().iter().copied())
979 {
980 return Err(BaselinePoolPreparationError::Contract(
981 BaselinePoolContractError::RestoreCanisterSetMismatch {
982 expected: baseline.snapshot_canister_ids().collect(),
983 actual: canisters.canister_ids().to_vec(),
984 },
985 ));
986 }
987
988 let started = Instant::now();
989 let reset_result = self.recipe.reset_non_snapshot_state(baseline);
990 add_timing(&mut timings.reset, started.elapsed());
991 let reset = reset_result.map_err(|source| BaselinePoolPreparationError::Recipe {
992 stage: BaselinePreparationStage::ResetNonSnapshotState,
993 source,
994 })?;
995
996 let started = Instant::now();
997 let readiness_result = self.recipe.drive_to_readiness(baseline);
998 add_timing(&mut timings.readiness, started.elapsed());
999 let readiness =
1000 readiness_result.map_err(|source| BaselinePoolPreparationError::Recipe {
1001 stage: BaselinePreparationStage::DriveToReadiness,
1002 source,
1003 })?;
1004 self.recipe
1005 .reset_requirements()
1006 .verify(&canisters, &reset)
1007 .map_err(BaselinePoolPreparationError::Contract)?;
1008
1009 let preparation = PreparedBaseline::Restored {
1010 canisters,
1011 reset,
1012 readiness,
1013 };
1014 self.validate_baseline(
1015 baseline,
1016 &preparation,
1017 BaselinePreparationStage::ValidateRestored,
1018 timings,
1019 )?;
1020 Ok(())
1021 }
1022
1023 fn build_slot(
1024 &self,
1025 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1026 timings: &mut BaselinePoolTimings,
1027 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1028 let started = Instant::now();
1029 let build = self.recipe.build();
1030 add_timing(&mut timings.build, started.elapsed());
1031 let baseline = build.map_err(|source| BaselinePoolPreparationError::Recipe {
1032 stage: BaselinePreparationStage::Build,
1033 source,
1034 })?;
1035
1036 if let Err(error) = self.validate_baseline(
1037 &baseline,
1038 &PreparedBaseline::Built,
1039 BaselinePreparationStage::ValidateBuilt,
1040 timings,
1041 ) {
1042 drop_baseline_safely(baseline);
1043 return Err(error);
1044 }
1045 let replaced = slot.replace(BaselineSlot {
1046 baseline,
1047 invalidation_reason: None,
1048 });
1049 debug_assert!(replaced.is_none());
1050 Ok(())
1051 }
1052
1053 fn validate_baseline(
1054 &self,
1055 baseline: &CachedPocketIcBaseline<R::Metadata>,
1056 preparation: &PreparedBaseline,
1057 stage: BaselinePreparationStage,
1058 timings: &mut BaselinePoolTimings,
1059 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1060 let started = Instant::now();
1061 let validation = self.recipe.validate(baseline, preparation);
1062 add_timing(&mut timings.validation, started.elapsed());
1063 let receipt =
1064 validation.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })?;
1065 if receipt.recipe_id() != self.recipe.id() {
1066 return Err(BaselinePoolPreparationError::Contract(
1067 BaselinePoolContractError::RecipeIdentityMismatch {
1068 expected: self.recipe.id().clone(),
1069 actual: receipt.recipe_id().clone(),
1070 },
1071 ));
1072 }
1073 Ok(())
1074 }
1075
1076 fn failure_disposition(
1077 &self,
1078 error: &BaselinePoolPreparationError<R::Error>,
1079 ) -> FailureDisposition {
1080 match error {
1081 BaselinePoolPreparationError::Recipe { stage, source } => {
1082 self.recipe.classify_failure(*stage, source)
1083 }
1084 BaselinePoolPreparationError::Contract(
1085 BaselinePoolContractError::RecipeIdentityMismatch { .. },
1086 ) => FailureDisposition::Fatal,
1087 BaselinePoolPreparationError::Contract(_) => {
1088 FailureDisposition::Rebuild(rebuild_reason_for_error(error))
1089 }
1090 }
1091 }
1092
1093 fn discard_stale_slot(
1094 slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1095 timings: &mut BaselinePoolTimings,
1096 ) {
1097 let started = Instant::now();
1098 if let Some(stale) = slot.take() {
1099 drop_baseline_safely(stale.baseline);
1100 }
1101 timings.stale_teardown = Some(started.elapsed());
1102 }
1103}
1104
1105impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
1106where
1107 R: PocketIcBaselineRecipe,
1108{
1109 #[must_use]
1111 pub const fn slot(&self) -> usize {
1112 self.slot.slot_index()
1113 }
1114
1115 pub fn invalidate(&mut self, reason: RebuildReason) {
1117 if let Some(slot) = self.slot.get_mut() {
1118 slot.invalidation_reason = Some(reason);
1119 }
1120 self.slot.invalidate();
1121 }
1122}
1123
1124impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
1125where
1126 R: PocketIcBaselineRecipe,
1127{
1128 type Target = CachedPocketIcBaseline<R::Metadata>;
1129
1130 fn deref(&self) -> &Self::Target {
1131 &self
1132 .slot
1133 .get()
1134 .expect("leased baseline pool slot must be populated")
1135 .baseline
1136 }
1137}
1138
1139fn nonempty_receipt_identity(
1140 receipt: &'static str,
1141 identity: String,
1142) -> Result<String, BaselinePoolContractError> {
1143 if identity.trim().is_empty() {
1144 return Err(BaselinePoolContractError::EmptyReceiptIdentity { receipt });
1145 }
1146 Ok(identity)
1147}
1148
1149const fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
1150 *total = saturating_add_optional_duration(*total, Some(elapsed));
1151}
1152
1153fn rebuild_reason_for_error<E>(error: &BaselinePoolPreparationError<E>) -> RebuildReason {
1154 match error {
1155 BaselinePoolPreparationError::Recipe { stage, .. } => stage.default_rebuild_reason(),
1156 BaselinePoolPreparationError::Contract(
1157 BaselinePoolContractError::MissingResetDomain { .. }
1158 | BaselinePoolContractError::ResetPolicyMismatch { .. }
1159 | BaselinePoolContractError::CyclePolicyMismatch { .. }
1160 | BaselinePoolContractError::DuplicateResetDomain { .. }
1161 | BaselinePoolContractError::RestoreCanisterSetMismatch { .. },
1162 ) => RebuildReason::ResetCoverageMismatch,
1163 BaselinePoolPreparationError::Contract(_) => RebuildReason::InvariantValidationFailure,
1164 }
1165}
1166
1167fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
1168 let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
1169}
1170
1171impl std::fmt::Display for FixtureRecipeId {
1172 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1173 formatter.write_str(&self.0)
1174 }
1175}
1176
1177impl std::fmt::Display for BaselinePreparationStage {
1178 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1179 formatter.write_str(match self {
1180 Self::Build => "build",
1181 Self::RestoreCanisters => "canister restore",
1182 Self::ResetNonSnapshotState => "non-snapshot reset",
1183 Self::DriveToReadiness => "readiness",
1184 Self::ValidateBuilt => "built-baseline validation",
1185 Self::ValidateRestored => "restored-baseline validation",
1186 })
1187 }
1188}
1189
1190impl std::fmt::Display for BaselinePoolContractError {
1191 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1192 match self {
1193 Self::EmptyRecipeIdentity => formatter.write_str("fixture recipe identity is empty"),
1194 Self::EmptyReceiptIdentity { receipt } => {
1195 write!(formatter, "{receipt} receipt identity is empty")
1196 }
1197 Self::DuplicateResetDomain { domain } => {
1198 write!(
1199 formatter,
1200 "reset domain {domain:?} was reported more than once"
1201 )
1202 }
1203 Self::DuplicateCanisterId { canister_id } => {
1204 write!(formatter, "restore receipt repeats canister {canister_id}")
1205 }
1206 Self::EmptyCanisterSet => formatter.write_str("restore receipt contains no canisters"),
1207 Self::CyclePolicyMismatch { required, achieved } => write!(
1208 formatter,
1209 "restore cycle policy {achieved:?} does not satisfy {required:?}",
1210 ),
1211 Self::RestoreCanisterSetMismatch { expected, actual } => write!(
1212 formatter,
1213 "restore receipt identified canisters {actual:?}, expected captured set {expected:?}",
1214 ),
1215 Self::MissingResetDomain { domain } => {
1216 write!(
1217 formatter,
1218 "required reset domain {domain:?} was not achieved"
1219 )
1220 }
1221 Self::ResetPolicyMismatch {
1222 requirement,
1223 achievement,
1224 } => write!(
1225 formatter,
1226 "reset achievement {achievement:?} does not satisfy {requirement:?}",
1227 ),
1228 Self::RecipeIdentityMismatch { expected, actual } => write!(
1229 formatter,
1230 "validation receipt used recipe `{actual}` instead of `{expected}`",
1231 ),
1232 }
1233 }
1234}
1235
1236impl std::error::Error for BaselinePoolContractError {}
1237
1238impl<E> std::fmt::Display for BaselinePoolPreparationError<E>
1239where
1240 E: std::fmt::Display,
1241{
1242 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1243 match self {
1244 Self::Recipe { stage, source } => {
1245 write!(formatter, "baseline {stage} failed: {source}")
1246 }
1247 Self::Contract(error) => write!(formatter, "baseline pool contract failed: {error}"),
1248 }
1249 }
1250}
1251
1252impl<E> std::error::Error for BaselinePoolPreparationError<E>
1253where
1254 E: std::error::Error + 'static,
1255{
1256 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1257 match self {
1258 Self::Recipe { source, .. } => Some(source),
1259 Self::Contract(error) => Some(error),
1260 }
1261 }
1262}
1263
1264impl<E> std::fmt::Display for BaselinePoolError<E>
1265where
1266 E: std::fmt::Display,
1267{
1268 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1269 match self {
1270 Self::Preparation { error, .. } => error.fmt(formatter),
1271 Self::RecoveryFailed {
1272 original, rebuild, ..
1273 } => write!(
1274 formatter,
1275 "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
1276 ),
1277 }
1278 }
1279}
1280
1281impl<E> std::error::Error for BaselinePoolError<E>
1282where
1283 E: std::error::Error + 'static,
1284{
1285 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1286 match self {
1287 Self::Preparation { error, .. } => Some(error),
1288 Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
1289 }
1290 }
1291}
1292
1293#[cfg(test)]
1294mod tests {
1295 use super::{
1296 BaselinePoolContractError, CanisterRestoreReceipt, CycleResetPolicy, FixtureRecipeId,
1297 ResetAchievement, ResetDomainKind, ResetReceipt, ResetRequirement, ResetRequirements,
1298 TimeResetPolicy,
1299 };
1300 use candid::Principal;
1301
1302 #[test]
1303 fn recipe_identity_must_be_nonempty() {
1304 assert!(matches!(
1305 FixtureRecipeId::try_new(" "),
1306 Err(BaselinePoolContractError::EmptyRecipeIdentity)
1307 ));
1308 }
1309
1310 #[test]
1311 fn caller_restore_receipts_validate_and_order_canister_ids() {
1312 let first = Principal::from_slice(&[1]);
1313 let second = Principal::from_slice(&[2]);
1314 let policy = CycleResetPolicy::TopUpTo(123);
1315 let receipt = CanisterRestoreReceipt::try_new([second, first], policy).unwrap();
1316 assert_eq!(receipt.canister_ids(), [first, second]);
1317 assert_eq!(receipt.cycle_policy(), policy);
1318 assert_eq!(
1319 CanisterRestoreReceipt::try_new([], policy),
1320 Err(BaselinePoolContractError::EmptyCanisterSet),
1321 );
1322 assert_eq!(
1323 CanisterRestoreReceipt::try_new([second, first, second], policy),
1324 Err(BaselinePoolContractError::DuplicateCanisterId {
1325 canister_id: second,
1326 }),
1327 );
1328 }
1329
1330 #[test]
1331 fn reset_requirements_reject_duplicate_domains() {
1332 let result = ResetRequirements::try_new(
1333 CycleResetPolicy::PreserveCurrent,
1334 [
1335 ResetRequirement::PocketIcTime(TimeResetPolicy::PreserveCurrent),
1336 ResetRequirement::PocketIcTime(TimeResetPolicy::RebuildOnMutation),
1337 ],
1338 );
1339 assert!(matches!(
1340 result,
1341 Err(BaselinePoolContractError::DuplicateResetDomain {
1342 domain: ResetDomainKind::PocketIcTime,
1343 })
1344 ));
1345 }
1346
1347 #[test]
1348 fn required_policy_must_match_achieved_policy() {
1349 let requirements = ResetRequirements::try_new(
1350 CycleResetPolicy::PreserveCurrent,
1351 [ResetRequirement::PocketIcTime(
1352 TimeResetPolicy::PreserveCurrent,
1353 )],
1354 )
1355 .unwrap();
1356 let restore = CanisterRestoreReceipt::try_new(
1357 [Principal::anonymous()],
1358 CycleResetPolicy::PreserveCurrent,
1359 )
1360 .unwrap();
1361 let receipt = ResetReceipt::try_new([ResetAchievement::PocketIcTime(
1362 TimeResetPolicy::RebuildOnMutation,
1363 )])
1364 .unwrap();
1365 assert!(matches!(
1366 requirements.verify(&restore, &receipt),
1367 Err(BaselinePoolContractError::ResetPolicyMismatch { .. })
1368 ));
1369 }
1370
1371 #[test]
1372 fn restore_cycle_policy_must_match_required_policy() {
1373 let requirements =
1374 ResetRequirements::try_new(CycleResetPolicy::RestoreExactBaseline, []).unwrap();
1375 let restore = CanisterRestoreReceipt::try_new(
1376 [Principal::anonymous()],
1377 CycleResetPolicy::PreserveCurrent,
1378 )
1379 .unwrap();
1380 assert!(matches!(
1381 requirements.verify(&restore, &ResetReceipt::empty()),
1382 Err(BaselinePoolContractError::CyclePolicyMismatch {
1383 required: CycleResetPolicy::RestoreExactBaseline,
1384 achieved: CycleResetPolicy::PreserveCurrent,
1385 })
1386 ));
1387 }
1388}