1use candid::Principal;
2use std::{
3 collections::{BTreeMap, BTreeSet},
4 num::NonZeroUsize,
5 ops::Deref,
6 panic::{AssertUnwindSafe, catch_unwind},
7 time::{Duration, Instant},
8};
9
10use super::{
11 CachedPocketIcBaseline, ControllerSnapshots, SnapshotRestoreFunding,
12 bounded_pool::{BoundedSlotLease, BoundedSlotPool},
13};
14
15#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
17pub struct FixtureRecipeId(String);
18
19impl From<SnapshotRestoreFunding> for CycleResetPolicy {
20 fn from(funding: SnapshotRestoreFunding) -> Self {
21 match funding {
22 SnapshotRestoreFunding::Preserve => Self::PreserveCurrent,
23 SnapshotRestoreFunding::TopUpTo { minimum_cycles } => Self::TopUpTo(minimum_cycles),
24 }
25 }
26}
27
28#[non_exhaustive]
30#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
31pub enum ResetDomainKind {
32 PocketIcTime,
34 ExtraCanisters,
36 PendingMessages,
38 SubnetState,
40 ExternalResources,
42}
43
44#[non_exhaustive]
46#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum CycleResetPolicy {
48 PreserveCurrent,
52 TopUpTo(u128),
55 RestoreExactBaseline,
57 RebuildOnMutation,
59}
60
61#[non_exhaustive]
63#[derive(Clone, Copy, Debug, Eq, PartialEq)]
64pub enum TimeResetPolicy {
65 PreserveCurrent,
67 RestoreBaseline,
69 RebuildOnMutation,
71}
72
73#[non_exhaustive]
75#[derive(Clone, Copy, Debug, Eq, PartialEq)]
76pub enum ExtraCanisterPolicy {
77 RequireBaselineSet,
79 RemoveTracked,
81 RebuildOnChange,
83}
84
85#[non_exhaustive]
87#[derive(Clone, Copy, Debug, Eq, PartialEq)]
88pub enum StateResetPolicy {
89 ResetByRecipe,
91 ValidateUnchanged,
93 IrrelevantByRecipeContract,
95 RebuildOnChange,
97}
98
99#[non_exhaustive]
126#[derive(Clone, Debug, Eq, PartialEq)]
127pub enum ResetDomainPolicy {
128 PocketIcTime(TimeResetPolicy),
130 ExtraCanisters(ExtraCanisterPolicy),
132 PendingMessages(StateResetPolicy),
134 SubnetState(StateResetPolicy),
136 ExternalResources(StateResetPolicy),
138}
139
140#[derive(Clone, Debug, Eq, PartialEq)]
142pub struct ResetRequirements {
143 cycle_policy: CycleResetPolicy,
144 domains: BTreeMap<ResetDomainKind, ResetDomainPolicy>,
145}
146
147#[derive(Clone, Debug, Default, Eq, PartialEq)]
149pub struct ResetReceipt(BTreeMap<ResetDomainKind, ResetDomainPolicy>);
150
151#[derive(Clone, Debug, Eq, PartialEq)]
153pub struct CanisterRestoreReceipt {
154 canister_ids: Vec<Principal>,
155 cycle_policy: CycleResetPolicy,
156}
157
158#[non_exhaustive]
160#[derive(Clone, Debug, Eq, PartialEq)]
161pub enum BaselinePoolContractError {
162 EmptyRecipeIdentity,
164 DuplicateResetDomain { domain: ResetDomainKind },
166 DuplicateCanisterId { canister_id: Principal },
168 EmptyCanisterSet,
170 CyclePolicyMismatch {
172 required: CycleResetPolicy,
173 achieved: CycleResetPolicy,
174 },
175 RestoreCanisterSetMismatch {
177 expected: Vec<Principal>,
178 actual: Vec<Principal>,
179 },
180 MissingResetDomain { domain: ResetDomainKind },
182 ResetPolicyMismatch {
184 requirement: ResetDomainPolicy,
185 achievement: ResetDomainPolicy,
186 },
187}
188
189#[non_exhaustive]
191#[derive(Clone, Debug, Eq, PartialEq)]
192pub enum PreparedBaseline {
193 Built,
195 Restored {
197 canisters: CanisterRestoreReceipt,
199 reset: ResetReceipt,
201 },
202}
203
204#[non_exhaustive]
206#[derive(Clone, Copy, Debug, Eq, PartialEq)]
207pub enum BaselinePreparationStage {
208 Build,
210 RestoreCanisters,
212 ResetNonSnapshotState,
214 DriveToReadiness,
216 ValidateBuilt,
218 ValidateRestored,
220}
221
222#[non_exhaustive]
224#[derive(Clone, Debug, Eq, PartialEq)]
225pub enum RebuildReason {
226 DeadPocketIcTransport,
228 SnapshotRestoreFailure,
230 ResetFailure,
232 ReadinessFailure,
234 ResetCoverageMismatch,
236 InvariantValidationFailure,
238 ExplicitLeaseInvalidation,
240 UnwindWhileLeased,
242 RecipeClassified { code: String },
244}
245
246#[non_exhaustive]
248#[derive(Clone, Debug, Eq, PartialEq)]
249pub enum FailureDisposition {
250 Fatal,
252 Rebuild(RebuildReason),
254}
255
256#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
258pub struct BaselinePoolTimings {
259 wait: Duration,
260 build: Option<Duration>,
261 restore: Option<Duration>,
262 reset: Option<Duration>,
263 readiness: Option<Duration>,
264 validation: Option<Duration>,
265 stale_teardown: Option<Duration>,
266 total: Duration,
267}
268
269#[derive(Clone, Debug, Eq, PartialEq)]
271pub struct BaselinePoolOutcome {
272 pub slot: usize,
274 pub timings: BaselinePoolTimings,
276 pub disposition: BaselinePoolDisposition,
278}
279
280#[non_exhaustive]
282#[derive(Clone, Debug, Eq, PartialEq)]
283pub enum BaselinePoolDisposition {
284 Built,
286 Restored,
288 Rebuilt(RebuildReason),
290}
291
292#[non_exhaustive]
294#[derive(Debug)]
295pub enum BaselinePoolPreparationError<E> {
296 Recipe {
298 stage: BaselinePreparationStage,
300 source: E,
302 },
303 Contract(BaselinePoolContractError),
305}
306
307#[non_exhaustive]
309#[derive(Debug)]
310pub enum BaselinePoolError<E> {
311 Preparation {
313 error: BaselinePoolPreparationError<E>,
315 timings: Box<BaselinePoolTimings>,
317 },
318 RecoveryFailed {
320 original: Box<BaselinePoolPreparationError<E>>,
322 rebuild: Box<BaselinePoolPreparationError<E>>,
324 timings: Box<BaselinePoolTimings>,
326 },
327}
328
329pub trait PocketIcBaselineRecipe: Send + Sync + 'static {
331 type Metadata: Send + 'static;
333 type Error: std::error::Error + Send + Sync + 'static;
335
336 fn id(&self) -> &FixtureRecipeId;
338
339 fn reset_requirements(&self) -> &ResetRequirements;
341
342 fn build(&self) -> Result<CachedPocketIcBaseline<Self::Metadata>, Self::Error>;
344
345 fn restore_canisters(
347 &self,
348 baseline: &CachedPocketIcBaseline<Self::Metadata>,
349 ) -> Result<CanisterRestoreReceipt, Self::Error>;
350
351 fn reset_non_snapshot_state(
353 &self,
354 baseline: &CachedPocketIcBaseline<Self::Metadata>,
355 ) -> Result<ResetReceipt, Self::Error>;
356
357 fn drive_to_readiness(
360 &self,
361 baseline: &CachedPocketIcBaseline<Self::Metadata>,
362 ) -> Result<(), Self::Error>;
363
364 fn validate(
367 &self,
368 baseline: &CachedPocketIcBaseline<Self::Metadata>,
369 preparation: &PreparedBaseline,
370 ) -> Result<(), Self::Error>;
371
372 fn classify_failure(
374 &self,
375 stage: BaselinePreparationStage,
376 _error: &Self::Error,
377 ) -> FailureDisposition {
378 FailureDisposition::Rebuild(stage.default_rebuild_reason())
379 }
380}
381
382pub struct CachedPocketIcBaselinePool<R>
394where
395 R: PocketIcBaselineRecipe,
396{
397 recipe: R,
398 slots: BoundedSlotPool<CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
399}
400
401pub struct CachedPocketIcBaselinePoolGuard<'a, R>
403where
404 R: PocketIcBaselineRecipe,
405{
406 slot: BoundedSlotLease<'a, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
407}
408
409impl FixtureRecipeId {
410 pub fn try_new(identity: impl Into<String>) -> Result<Self, BaselinePoolContractError> {
412 let identity = identity.into();
413 if identity.trim().is_empty() {
414 return Err(BaselinePoolContractError::EmptyRecipeIdentity);
415 }
416 Ok(Self(identity))
417 }
418
419 #[must_use]
421 pub fn as_str(&self) -> &str {
422 &self.0
423 }
424}
425
426impl ResetDomainPolicy {
427 #[must_use]
429 pub const fn domain(&self) -> ResetDomainKind {
430 match self {
431 Self::PocketIcTime(_) => ResetDomainKind::PocketIcTime,
432 Self::ExtraCanisters(_) => ResetDomainKind::ExtraCanisters,
433 Self::PendingMessages(_) => ResetDomainKind::PendingMessages,
434 Self::SubnetState(_) => ResetDomainKind::SubnetState,
435 Self::ExternalResources(_) => ResetDomainKind::ExternalResources,
436 }
437 }
438}
439
440fn collect_reset_domains(
441 policies: impl IntoIterator<Item = ResetDomainPolicy>,
442) -> Result<BTreeMap<ResetDomainKind, ResetDomainPolicy>, BaselinePoolContractError> {
443 let mut domains = BTreeMap::new();
444 for policy in policies {
445 let domain = policy.domain();
446 if domains.insert(domain, policy).is_some() {
447 return Err(BaselinePoolContractError::DuplicateResetDomain { domain });
448 }
449 }
450 Ok(domains)
451}
452
453impl ResetRequirements {
454 pub fn try_new<I>(
459 cycle_policy: CycleResetPolicy,
460 requirements: I,
461 ) -> Result<Self, BaselinePoolContractError>
462 where
463 I: IntoIterator<Item = ResetDomainPolicy>,
464 {
465 Ok(Self {
466 cycle_policy,
467 domains: collect_reset_domains(requirements)?,
468 })
469 }
470
471 #[must_use]
473 pub const fn cycle_policy(&self) -> CycleResetPolicy {
474 self.cycle_policy
475 }
476
477 #[must_use]
479 pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetDomainPolicy> {
480 self.domains.get(&domain)
481 }
482
483 pub fn iter(&self) -> impl Iterator<Item = &ResetDomainPolicy> {
485 self.domains.values()
486 }
487
488 fn verify(
489 &self,
490 restore: &CanisterRestoreReceipt,
491 receipt: &ResetReceipt,
492 ) -> Result<(), BaselinePoolContractError> {
493 if restore.cycle_policy != self.cycle_policy {
494 return Err(BaselinePoolContractError::CyclePolicyMismatch {
495 required: self.cycle_policy,
496 achieved: restore.cycle_policy,
497 });
498 }
499 for (domain, requirement) in &self.domains {
500 let Some(achievement) = receipt.0.get(domain) else {
501 return Err(BaselinePoolContractError::MissingResetDomain { domain: *domain });
502 };
503 if achievement != requirement {
504 return Err(BaselinePoolContractError::ResetPolicyMismatch {
505 requirement: requirement.clone(),
506 achievement: achievement.clone(),
507 });
508 }
509 }
510 Ok(())
511 }
512}
513
514impl ResetReceipt {
515 pub fn try_new<I>(achievements: I) -> Result<Self, BaselinePoolContractError>
520 where
521 I: IntoIterator<Item = ResetDomainPolicy>,
522 {
523 Ok(Self(collect_reset_domains(achievements)?))
524 }
525
526 #[must_use]
528 pub const fn empty() -> Self {
529 Self(BTreeMap::new())
530 }
531
532 #[must_use]
534 pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetDomainPolicy> {
535 self.0.get(&domain)
536 }
537
538 pub fn iter(&self) -> impl Iterator<Item = &ResetDomainPolicy> {
540 self.0.values()
541 }
542}
543
544impl CanisterRestoreReceipt {
545 pub fn try_new<I>(
547 canister_ids: I,
548 cycle_policy: CycleResetPolicy,
549 ) -> Result<Self, BaselinePoolContractError>
550 where
551 I: IntoIterator<Item = Principal>,
552 {
553 let mut unique = BTreeSet::new();
554 for canister_id in canister_ids {
555 if !unique.insert(canister_id) {
556 return Err(BaselinePoolContractError::DuplicateCanisterId { canister_id });
557 }
558 }
559 if unique.is_empty() {
560 return Err(BaselinePoolContractError::EmptyCanisterSet);
561 }
562 Ok(Self {
563 canister_ids: unique.into_iter().collect(),
564 cycle_policy,
565 })
566 }
567
568 pub(super) fn from_restore(
571 snapshots: &ControllerSnapshots,
572 funding: SnapshotRestoreFunding,
573 ) -> Self {
574 Self {
575 canister_ids: snapshots.canister_ids().collect(),
576 cycle_policy: funding.into(),
577 }
578 }
579
580 #[must_use]
582 pub fn canister_ids(&self) -> &[Principal] {
583 &self.canister_ids
584 }
585
586 #[must_use]
588 pub const fn cycle_policy(&self) -> CycleResetPolicy {
589 self.cycle_policy
590 }
591}
592
593impl BaselinePreparationStage {
594 #[must_use]
599 pub fn default_rebuild_reason(self) -> RebuildReason {
600 match self {
601 Self::RestoreCanisters => RebuildReason::SnapshotRestoreFailure,
602 Self::ResetNonSnapshotState => RebuildReason::ResetFailure,
603 Self::DriveToReadiness => RebuildReason::ReadinessFailure,
604 Self::ValidateRestored | Self::ValidateBuilt => {
605 RebuildReason::InvariantValidationFailure
606 }
607 Self::Build => RebuildReason::RecipeClassified {
608 code: "build".to_owned(),
609 },
610 }
611 }
612}
613
614impl BaselinePoolTimings {
615 #[must_use]
617 pub const fn wait(self) -> Duration {
618 self.wait
619 }
620
621 #[must_use]
623 pub const fn build(self) -> Option<Duration> {
624 self.build
625 }
626
627 #[must_use]
629 pub const fn restore(self) -> Option<Duration> {
630 self.restore
631 }
632
633 #[must_use]
635 pub const fn reset(self) -> Option<Duration> {
636 self.reset
637 }
638
639 #[must_use]
641 pub const fn readiness(self) -> Option<Duration> {
642 self.readiness
643 }
644
645 #[must_use]
647 pub const fn validation(self) -> Option<Duration> {
648 self.validation
649 }
650
651 #[must_use]
653 pub const fn stale_teardown(self) -> Option<Duration> {
654 self.stale_teardown
655 }
656
657 #[must_use]
659 pub const fn total(self) -> Duration {
660 self.total
661 }
662}
663
664impl std::fmt::Display for BaselinePoolTimings {
665 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
666 write!(
667 formatter,
668 "total={:?} wait={:?} build={:?} restore={:?} reset={:?} readiness={:?} validation={:?} stale_teardown={:?}",
669 self.total,
670 self.wait,
671 self.build,
672 self.restore,
673 self.reset,
674 self.readiness,
675 self.validation,
676 self.stale_teardown,
677 )
678 }
679}
680
681impl std::fmt::Display for BaselinePoolOutcome {
682 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
683 match &self.disposition {
684 BaselinePoolDisposition::Built => formatter.write_str("built")?,
685 BaselinePoolDisposition::Restored => formatter.write_str("restored")?,
686 BaselinePoolDisposition::Rebuilt(reason) => {
687 write!(
688 formatter,
689 "rebuilt slot={} reason={reason:?} {}",
690 self.slot, self.timings
691 )?;
692 return Ok(());
693 }
694 }
695 write!(formatter, " slot={} {}", self.slot, self.timings)
696 }
697}
698
699impl<E> BaselinePoolError<E> {
700 #[must_use]
702 pub const fn timings(&self) -> BaselinePoolTimings {
703 match self {
704 Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
705 }
706 }
707}
708
709impl<R> CachedPocketIcBaselinePool<R>
710where
711 R: PocketIcBaselineRecipe,
712{
713 #[must_use]
718 pub const fn new(capacity: NonZeroUsize, recipe: R) -> Self {
719 Self {
720 recipe,
721 slots: BoundedSlotPool::new(capacity),
722 }
723 }
724
725 #[must_use]
727 pub fn recipe_id(&self) -> &FixtureRecipeId {
728 self.recipe.id()
729 }
730
731 #[must_use]
733 pub const fn capacity(&self) -> NonZeroUsize {
734 self.slots.capacity()
735 }
736
737 pub fn acquire(
749 &self,
750 ) -> Result<
751 (CachedPocketIcBaselinePoolGuard<'_, R>, BaselinePoolOutcome),
752 BaselinePoolError<R::Error>,
753 > {
754 let total_started = Instant::now();
755 let mut slot = self.slots.acquire();
756 let mut timings = BaselinePoolTimings {
757 wait: slot.wait(),
758 ..BaselinePoolTimings::default()
759 };
760 let slot_index = slot.slot_index();
761
762 let (rebuild_reason, original) = if slot.is_reusable() {
763 match self.prepare_reused(&slot, &mut timings) {
764 Ok(()) => {
765 timings.total = total_started.elapsed();
766 return Ok((
767 CachedPocketIcBaselinePoolGuard { slot },
768 BaselinePoolOutcome {
769 slot: slot_index,
770 timings,
771 disposition: BaselinePoolDisposition::Restored,
772 },
773 ));
774 }
775 Err(original) => match self.failure_disposition(&original) {
776 FailureDisposition::Fatal => {
777 Self::discard_stale_slot(&mut slot, &mut timings);
780 timings.total = total_started.elapsed();
781 return Err(BaselinePoolError::Preparation {
782 error: original,
783 timings: Box::new(timings),
784 });
785 }
786 FailureDisposition::Rebuild(reason) => (Some(reason), Some(original)),
787 },
788 }
789 } else {
790 let reason = if slot.invalidated_by_unwind() {
791 Some(RebuildReason::UnwindWhileLeased)
792 } else {
793 slot.invalidation_reason().cloned()
794 };
795 (reason, None)
796 };
797
798 if slot.is_populated() {
799 Self::discard_stale_slot(&mut slot, &mut timings);
800 }
801 let build = self.build_slot(&mut slot, &mut timings);
802 timings.total = total_started.elapsed();
803 if let Err(error) = build {
804 return Err(match original {
805 Some(original) => BaselinePoolError::RecoveryFailed {
806 original: Box::new(original),
807 rebuild: Box::new(error),
808 timings: Box::new(timings),
809 },
810 None => BaselinePoolError::Preparation {
811 error,
812 timings: Box::new(timings),
813 },
814 });
815 }
816
817 let outcome = BaselinePoolOutcome {
818 slot: slot_index,
819 timings,
820 disposition: rebuild_reason.map_or(
821 BaselinePoolDisposition::Built,
822 BaselinePoolDisposition::Rebuilt,
823 ),
824 };
825 Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
826 }
827
828 fn prepare_reused(
829 &self,
830 slot: &BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
831 timings: &mut BaselinePoolTimings,
832 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
833 let baseline = slot
834 .get()
835 .expect("reusable baseline slot must be populated");
836
837 let started = Instant::now();
838 let restore = self.recipe.restore_canisters(baseline);
839 add_timing(&mut timings.restore, started.elapsed());
840 let canisters = restore.map_err(|source| BaselinePoolPreparationError::Recipe {
841 stage: BaselinePreparationStage::RestoreCanisters,
842 source,
843 })?;
844 if !baseline
845 .snapshot_canister_ids()
846 .eq(canisters.canister_ids().iter().copied())
847 {
848 return Err(BaselinePoolPreparationError::Contract(
849 BaselinePoolContractError::RestoreCanisterSetMismatch {
850 expected: baseline.snapshot_canister_ids().collect(),
851 actual: canisters.canister_ids().to_vec(),
852 },
853 ));
854 }
855
856 let started = Instant::now();
857 let reset_result = self.recipe.reset_non_snapshot_state(baseline);
858 add_timing(&mut timings.reset, started.elapsed());
859 let reset = reset_result.map_err(|source| BaselinePoolPreparationError::Recipe {
860 stage: BaselinePreparationStage::ResetNonSnapshotState,
861 source,
862 })?;
863
864 let started = Instant::now();
865 let readiness_result = self.recipe.drive_to_readiness(baseline);
866 add_timing(&mut timings.readiness, started.elapsed());
867 readiness_result.map_err(|source| BaselinePoolPreparationError::Recipe {
868 stage: BaselinePreparationStage::DriveToReadiness,
869 source,
870 })?;
871 self.recipe
872 .reset_requirements()
873 .verify(&canisters, &reset)
874 .map_err(BaselinePoolPreparationError::Contract)?;
875
876 let preparation = PreparedBaseline::Restored { canisters, reset };
877 self.validate_baseline(
878 baseline,
879 &preparation,
880 BaselinePreparationStage::ValidateRestored,
881 timings,
882 )?;
883 Ok(())
884 }
885
886 fn build_slot(
887 &self,
888 slot: &mut BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
889 timings: &mut BaselinePoolTimings,
890 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
891 let started = Instant::now();
892 let build = self.recipe.build();
893 add_timing(&mut timings.build, started.elapsed());
894 let baseline = build.map_err(|source| BaselinePoolPreparationError::Recipe {
895 stage: BaselinePreparationStage::Build,
896 source,
897 })?;
898
899 if let Err(error) = self.validate_baseline(
900 &baseline,
901 &PreparedBaseline::Built,
902 BaselinePreparationStage::ValidateBuilt,
903 timings,
904 ) {
905 drop_baseline_safely(baseline);
906 return Err(error);
907 }
908 let replaced = slot.replace(baseline);
909 debug_assert!(replaced.is_none());
910 Ok(())
911 }
912
913 fn validate_baseline(
914 &self,
915 baseline: &CachedPocketIcBaseline<R::Metadata>,
916 preparation: &PreparedBaseline,
917 stage: BaselinePreparationStage,
918 timings: &mut BaselinePoolTimings,
919 ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
920 let started = Instant::now();
921 let validation = self.recipe.validate(baseline, preparation);
922 add_timing(&mut timings.validation, started.elapsed());
923 validation.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })
924 }
925
926 fn failure_disposition(
927 &self,
928 error: &BaselinePoolPreparationError<R::Error>,
929 ) -> FailureDisposition {
930 match error {
931 BaselinePoolPreparationError::Recipe { stage, source } => {
932 self.recipe.classify_failure(*stage, source)
933 }
934 BaselinePoolPreparationError::Contract(error) => {
935 FailureDisposition::Rebuild(match error {
936 BaselinePoolContractError::MissingResetDomain { .. }
937 | BaselinePoolContractError::ResetPolicyMismatch { .. }
938 | BaselinePoolContractError::CyclePolicyMismatch { .. }
939 | BaselinePoolContractError::DuplicateResetDomain { .. }
940 | BaselinePoolContractError::RestoreCanisterSetMismatch { .. } => {
941 RebuildReason::ResetCoverageMismatch
942 }
943 _ => RebuildReason::InvariantValidationFailure,
944 })
945 }
946 }
947 }
948
949 fn discard_stale_slot(
950 slot: &mut BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
951 timings: &mut BaselinePoolTimings,
952 ) {
953 let started = Instant::now();
954 if let Some(stale) = slot.take() {
955 drop_baseline_safely(stale);
956 }
957 timings.stale_teardown = Some(started.elapsed());
958 }
959}
960
961impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
962where
963 R: PocketIcBaselineRecipe,
964{
965 #[must_use]
967 pub const fn slot(&self) -> usize {
968 self.slot.slot_index()
969 }
970
971 pub fn invalidate(&mut self, reason: RebuildReason) {
973 self.slot.invalidate(reason);
974 }
975}
976
977impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
978where
979 R: PocketIcBaselineRecipe,
980{
981 type Target = CachedPocketIcBaseline<R::Metadata>;
982
983 fn deref(&self) -> &Self::Target {
984 self.slot
985 .get()
986 .expect("leased baseline pool slot must be populated")
987 }
988}
989
990const fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
991 *total = Some(match *total {
992 None => elapsed,
993 Some(previous) => previous.saturating_add(elapsed),
994 });
995}
996
997fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
998 let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
999}
1000
1001impl std::fmt::Display for FixtureRecipeId {
1002 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1003 formatter.write_str(&self.0)
1004 }
1005}
1006
1007impl std::fmt::Display for BaselinePreparationStage {
1008 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1009 formatter.write_str(match self {
1010 Self::Build => "build",
1011 Self::RestoreCanisters => "canister restore",
1012 Self::ResetNonSnapshotState => "non-snapshot reset",
1013 Self::DriveToReadiness => "readiness",
1014 Self::ValidateBuilt => "built-baseline validation",
1015 Self::ValidateRestored => "restored-baseline validation",
1016 })
1017 }
1018}
1019
1020impl std::fmt::Display for BaselinePoolContractError {
1021 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1022 match self {
1023 Self::EmptyRecipeIdentity => formatter.write_str("fixture recipe identity is empty"),
1024 Self::DuplicateResetDomain { domain } => {
1025 write!(
1026 formatter,
1027 "reset domain {domain:?} was reported more than once"
1028 )
1029 }
1030 Self::DuplicateCanisterId { canister_id } => {
1031 write!(formatter, "restore receipt repeats canister {canister_id}")
1032 }
1033 Self::EmptyCanisterSet => formatter.write_str("restore receipt contains no canisters"),
1034 Self::CyclePolicyMismatch { required, achieved } => write!(
1035 formatter,
1036 "restore cycle policy {achieved:?} does not satisfy {required:?}",
1037 ),
1038 Self::RestoreCanisterSetMismatch { expected, actual } => write!(
1039 formatter,
1040 "restore receipt identified canisters {actual:?}, expected captured set {expected:?}",
1041 ),
1042 Self::MissingResetDomain { domain } => {
1043 write!(
1044 formatter,
1045 "required reset domain {domain:?} was not achieved"
1046 )
1047 }
1048 Self::ResetPolicyMismatch {
1049 requirement,
1050 achievement,
1051 } => write!(
1052 formatter,
1053 "reset achievement {achievement:?} does not satisfy {requirement:?}",
1054 ),
1055 }
1056 }
1057}
1058
1059impl std::error::Error for BaselinePoolContractError {}
1060
1061impl<E> std::fmt::Display for BaselinePoolPreparationError<E>
1062where
1063 E: std::fmt::Display,
1064{
1065 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1066 match self {
1067 Self::Recipe { stage, source } => {
1068 write!(formatter, "baseline {stage} failed: {source}")
1069 }
1070 Self::Contract(error) => write!(formatter, "baseline pool contract failed: {error}"),
1071 }
1072 }
1073}
1074
1075impl<E> std::error::Error for BaselinePoolPreparationError<E>
1076where
1077 E: std::error::Error + 'static,
1078{
1079 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1080 match self {
1081 Self::Recipe { source, .. } => Some(source),
1082 Self::Contract(error) => Some(error),
1083 }
1084 }
1085}
1086
1087impl<E> std::fmt::Display for BaselinePoolError<E>
1088where
1089 E: std::fmt::Display,
1090{
1091 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1092 match self {
1093 Self::Preparation { error, .. } => error.fmt(formatter),
1094 Self::RecoveryFailed {
1095 original, rebuild, ..
1096 } => write!(
1097 formatter,
1098 "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
1099 ),
1100 }
1101 }
1102}
1103
1104impl<E> std::error::Error for BaselinePoolError<E>
1105where
1106 E: std::error::Error + 'static,
1107{
1108 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1109 match self {
1110 Self::Preparation { error, .. } => Some(error),
1111 Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
1112 }
1113 }
1114}
1115
1116#[cfg(test)]
1117mod tests {
1118 use super::{
1119 BaselinePoolContractError, BaselinePoolTimings, CanisterRestoreReceipt, CycleResetPolicy,
1120 ExtraCanisterPolicy, FixtureRecipeId, ResetDomainKind, ResetDomainPolicy, ResetReceipt,
1121 ResetRequirements, StateResetPolicy, TimeResetPolicy, add_timing,
1122 };
1123 use candid::Principal;
1124 use std::time::Duration;
1125
1126 #[test]
1127 fn phase_timing_preserves_absence_zero_and_saturation() {
1128 let mut timings = BaselinePoolTimings::default();
1129 assert_eq!(timings.restore(), None);
1130 add_timing(&mut timings.restore, Duration::ZERO);
1131 assert_eq!(timings.restore(), Some(Duration::ZERO));
1132 add_timing(&mut timings.restore, Duration::from_secs(2));
1133 assert_eq!(timings.restore(), Some(Duration::from_secs(2)));
1134 add_timing(&mut timings.restore, Duration::MAX);
1135 assert_eq!(timings.restore(), Some(Duration::MAX));
1136 assert_eq!(timings.build(), None);
1137 }
1138
1139 #[test]
1140 fn recipe_identity_must_be_nonempty() {
1141 assert!(matches!(
1142 FixtureRecipeId::try_new(" "),
1143 Err(BaselinePoolContractError::EmptyRecipeIdentity)
1144 ));
1145 }
1146
1147 #[test]
1148 fn caller_restore_receipts_validate_and_order_canister_ids() {
1149 let first = Principal::from_slice(&[1]);
1150 let second = Principal::from_slice(&[2]);
1151 let policy = CycleResetPolicy::TopUpTo(123);
1152 let receipt = CanisterRestoreReceipt::try_new([second, first], policy).unwrap();
1153 assert_eq!(receipt.canister_ids(), [first, second]);
1154 assert_eq!(receipt.cycle_policy(), policy);
1155 assert_eq!(
1156 CanisterRestoreReceipt::try_new([], policy),
1157 Err(BaselinePoolContractError::EmptyCanisterSet),
1158 );
1159 assert_eq!(
1160 CanisterRestoreReceipt::try_new([second, first, second], policy),
1161 Err(BaselinePoolContractError::DuplicateCanisterId {
1162 canister_id: second,
1163 }),
1164 );
1165 }
1166
1167 #[test]
1168 fn reset_declarations_and_receipts_reject_duplicate_domains() {
1169 let policies = [
1170 ResetDomainPolicy::PocketIcTime(TimeResetPolicy::PreserveCurrent),
1171 ResetDomainPolicy::PocketIcTime(TimeResetPolicy::RebuildOnMutation),
1172 ];
1173 let error = BaselinePoolContractError::DuplicateResetDomain {
1174 domain: ResetDomainKind::PocketIcTime,
1175 };
1176 assert_eq!(
1177 ResetRequirements::try_new(CycleResetPolicy::PreserveCurrent, policies.clone()),
1178 Err(error.clone()),
1179 );
1180 assert_eq!(ResetReceipt::try_new(policies), Err(error));
1181 }
1182
1183 #[test]
1184 fn required_policy_must_match_achieved_policy() {
1185 let restore = CanisterRestoreReceipt::try_new(
1186 [Principal::anonymous()],
1187 CycleResetPolicy::PreserveCurrent,
1188 )
1189 .unwrap();
1190 for (requirement, achievement) in [
1191 (
1192 ResetDomainPolicy::PocketIcTime(TimeResetPolicy::PreserveCurrent),
1193 ResetDomainPolicy::PocketIcTime(TimeResetPolicy::RebuildOnMutation),
1194 ),
1195 (
1196 ResetDomainPolicy::ExtraCanisters(ExtraCanisterPolicy::RequireBaselineSet),
1197 ResetDomainPolicy::ExtraCanisters(ExtraCanisterPolicy::RemoveTracked),
1198 ),
1199 (
1200 ResetDomainPolicy::PendingMessages(StateResetPolicy::ValidateUnchanged),
1201 ResetDomainPolicy::PendingMessages(StateResetPolicy::IrrelevantByRecipeContract),
1202 ),
1203 (
1204 ResetDomainPolicy::SubnetState(StateResetPolicy::ResetByRecipe),
1205 ResetDomainPolicy::SubnetState(StateResetPolicy::RebuildOnChange),
1206 ),
1207 (
1208 ResetDomainPolicy::ExternalResources(StateResetPolicy::ValidateUnchanged),
1209 ResetDomainPolicy::ExternalResources(StateResetPolicy::ResetByRecipe),
1210 ),
1211 ] {
1212 let requirements = ResetRequirements::try_new(
1213 CycleResetPolicy::PreserveCurrent,
1214 [requirement.clone()],
1215 )
1216 .unwrap();
1217 let receipt = ResetReceipt::try_new([requirement.clone()]).unwrap();
1218 assert_eq!(requirements.verify(&restore, &receipt), Ok(()));
1219 assert_eq!(
1220 requirements.verify(&restore, &ResetReceipt::empty()),
1221 Err(BaselinePoolContractError::MissingResetDomain {
1222 domain: requirement.domain(),
1223 }),
1224 );
1225 let receipt = ResetReceipt::try_new([achievement.clone()]).unwrap();
1226 assert_eq!(
1227 requirements.verify(&restore, &receipt),
1228 Err(BaselinePoolContractError::ResetPolicyMismatch {
1229 requirement,
1230 achievement,
1231 }),
1232 );
1233 }
1234 }
1235
1236 #[test]
1237 fn restore_cycle_policy_must_match_required_policy() {
1238 let requirements =
1239 ResetRequirements::try_new(CycleResetPolicy::RestoreExactBaseline, []).unwrap();
1240 let restore = CanisterRestoreReceipt::try_new(
1241 [Principal::anonymous()],
1242 CycleResetPolicy::PreserveCurrent,
1243 )
1244 .unwrap();
1245 assert!(matches!(
1246 requirements.verify(&restore, &ResetReceipt::empty()),
1247 Err(BaselinePoolContractError::CyclePolicyMismatch {
1248 required: CycleResetPolicy::RestoreExactBaseline,
1249 achieved: CycleResetPolicy::PreserveCurrent,
1250 })
1251 ));
1252 }
1253}