Skip to main content

ic_testkit/pic/
baseline_pool.rs

1use candid::Principal;
2use std::{
3    collections::{BTreeMap, BTreeSet},
4    num::NonZeroUsize,
5    ops::Deref,
6    panic::{AssertUnwindSafe, catch_unwind},
7    time::{Duration, Instant},
8};
9
10use super::{
11    CachedPocketIcBaseline,
12    bounded_pool::{BoundedSlotLease, BoundedSlotPool},
13};
14
15/// Caller-owned stable identity for one pooled fixture recipe.
16#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
17pub struct FixtureRecipeId(String);
18
19/// Reset domain whose handling is declared by a pooled baseline recipe.
20#[non_exhaustive]
21#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
22pub enum ResetDomainKind {
23    /// Captured canister snapshots.
24    CanisterSnapshots,
25    /// Canister cycle balances.
26    CanisterCycles,
27    /// PocketIC simulated time.
28    PocketIcTime,
29    /// Canisters outside the captured baseline set.
30    ExtraCanisters,
31    /// Pending ingress, timers, or cross-canister messages.
32    PendingMessages,
33    /// Subnet metrics, routing, allocation, or other subnet-global state.
34    SubnetState,
35    /// Files, processes, services, or other caller-owned resources.
36    ExternalResources,
37}
38
39/// Cycle handling required or achieved by one reset.
40#[non_exhaustive]
41#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42pub enum CycleResetPolicy {
43    /// Do not proactively add or remove cycles before snapshot restoration.
44    ///
45    /// PocketIC may still charge cycles while performing the restore.
46    PreserveCurrent,
47    /// Add cycles as needed to reach this minimum immediately before restore,
48    /// without removing excess.
49    TopUpTo(u128),
50    /// Restore the exact balance recorded by the recipe baseline.
51    RestoreExactBaseline,
52    /// Treat any relevant cycle mutation as requiring slot reconstruction.
53    RebuildOnMutation,
54}
55
56/// PocketIC time handling required or achieved by one reset.
57#[non_exhaustive]
58#[derive(Clone, Copy, Debug, Eq, PartialEq)]
59pub enum TimeResetPolicy {
60    /// Preserve the current simulator time rather than claiming it was reset.
61    PreserveCurrent,
62    /// Restore the exact time recorded by the recipe baseline.
63    RestoreBaseline,
64    /// Treat any relevant time mutation as requiring slot reconstruction.
65    RebuildOnMutation,
66}
67
68/// Extra-canister handling required or achieved by one reset.
69#[non_exhaustive]
70#[derive(Clone, Copy, Debug, Eq, PartialEq)]
71pub enum ExtraCanisterPolicy {
72    /// Validate that the baseline canister set is unchanged.
73    RequireBaselineSet,
74    /// Remove canisters explicitly tracked by the recipe.
75    RemoveTracked,
76    /// Treat any extra-canister change as requiring slot reconstruction.
77    RebuildOnChange,
78}
79
80/// Generic handling for reset domains without a more specific policy.
81#[non_exhaustive]
82#[derive(Clone, Copy, Debug, Eq, PartialEq)]
83pub enum StateResetPolicy {
84    /// Reset the domain through caller-owned recipe logic.
85    ResetByRecipe,
86    /// Validate that the domain remained unchanged.
87    ValidateUnchanged,
88    /// Explicitly declare the domain irrelevant to this recipe's guarantees.
89    IrrelevantByRecipeContract,
90    /// Treat any relevant change as requiring slot reconstruction.
91    RebuildOnChange,
92}
93
94/// One reset guarantee required before a baseline may be reused.
95#[non_exhaustive]
96#[derive(Clone, Debug, Eq, PartialEq)]
97pub enum ResetRequirement {
98    /// Restore every captured canister snapshot.
99    CanisterSnapshots,
100    /// Apply this cycle policy.
101    CanisterCycles(CycleResetPolicy),
102    /// Apply this time policy.
103    PocketIcTime(TimeResetPolicy),
104    /// Apply this extra-canister policy.
105    ExtraCanisters(ExtraCanisterPolicy),
106    /// Apply this pending-message policy.
107    PendingMessages(StateResetPolicy),
108    /// Apply this subnet-state policy.
109    SubnetState(StateResetPolicy),
110    /// Apply this external-resource policy.
111    ExternalResources(StateResetPolicy),
112}
113
114/// One reset guarantee reported as achieved by a recipe.
115#[non_exhaustive]
116#[derive(Clone, Debug, Eq, PartialEq)]
117pub enum ResetAchievement {
118    /// Every captured canister snapshot was restored.
119    CanisterSnapshots,
120    /// This cycle policy was achieved.
121    CanisterCycles(CycleResetPolicy),
122    /// This time policy was achieved.
123    PocketIcTime(TimeResetPolicy),
124    /// This extra-canister policy was achieved.
125    ExtraCanisters(ExtraCanisterPolicy),
126    /// This pending-message policy was achieved.
127    PendingMessages(StateResetPolicy),
128    /// This subnet-state policy was achieved.
129    SubnetState(StateResetPolicy),
130    /// This external-resource policy was achieved.
131    ExternalResources(StateResetPolicy),
132}
133
134/// Typed reset guarantees required by one fixture recipe.
135#[derive(Clone, Debug, Eq, PartialEq)]
136pub struct ResetRequirements(BTreeMap<ResetDomainKind, ResetRequirement>);
137
138/// Typed reset guarantees achieved by one preparation pass.
139#[derive(Clone, Debug, Default, Eq, PartialEq)]
140pub struct ResetReceipt(BTreeMap<ResetDomainKind, ResetAchievement>);
141
142/// Receipt for restoring the recipe's captured canister set.
143#[derive(Clone, Debug, Eq, PartialEq)]
144pub struct CanisterRestoreReceipt {
145    canister_ids: Vec<Principal>,
146    cycle_policy: CycleResetPolicy,
147}
148
149/// Receipt identifying the readiness boundary reached after reset.
150#[derive(Clone, Debug, Eq, PartialEq)]
151pub struct ReadinessReceipt {
152    identity: String,
153}
154
155/// Receipt proving the recipe's final invariant validation ran successfully.
156#[derive(Clone, Debug, Eq, PartialEq)]
157pub struct ValidationReceipt {
158    recipe_id: FixtureRecipeId,
159    invariant_identity: String,
160}
161
162/// Contract failure while constructing or verifying recipe reset evidence.
163#[non_exhaustive]
164#[derive(Clone, Debug, Eq, PartialEq)]
165pub enum BaselinePoolContractError {
166    /// Recipe identity was empty or whitespace-only.
167    EmptyRecipeIdentity,
168    /// A receipt identity was empty or whitespace-only.
169    EmptyReceiptIdentity { receipt: &'static str },
170    /// A reset domain was declared more than once.
171    DuplicateResetDomain { domain: ResetDomainKind },
172    /// A restored canister appeared more than once.
173    DuplicateCanisterId { canister_id: Principal },
174    /// A restore receipt contained no canisters.
175    EmptyCanisterSet,
176    /// A recipe omitted a reset domain required by every pooled baseline.
177    UndeclaredRequiredResetDomain { domain: ResetDomainKind },
178    /// The restored canister receipt did not identify the complete snapshot set.
179    RestoreCanisterSetMismatch {
180        expected: Vec<Principal>,
181        actual: Vec<Principal>,
182    },
183    /// A required reset domain had no matching achievement.
184    MissingResetDomain { domain: ResetDomainKind },
185    /// A reset achievement did not satisfy the required policy.
186    ResetPolicyMismatch {
187        requirement: ResetRequirement,
188        achievement: ResetAchievement,
189    },
190    /// Final validation reported a recipe other than the pool-owned recipe.
191    RecipeIdentityMismatch {
192        expected: FixtureRecipeId,
193        actual: FixtureRecipeId,
194    },
195}
196
197/// Whether validation is observing a newly built or restored baseline.
198#[non_exhaustive]
199#[derive(Clone, Debug, Eq, PartialEq)]
200pub enum PreparedBaseline {
201    /// The recipe just built this baseline.
202    Built,
203    /// The recipe restored and reset an existing baseline.
204    Restored {
205        /// Captured canisters restored by the recipe.
206        canisters: CanisterRestoreReceipt,
207        /// Combined typed reset receipt.
208        reset: ResetReceipt,
209        /// Readiness boundary reached after reset.
210        readiness: ReadinessReceipt,
211    },
212}
213
214/// Recipe stage associated with a structured preparation failure.
215#[non_exhaustive]
216#[derive(Clone, Copy, Debug, Eq, PartialEq)]
217pub enum BaselinePreparationStage {
218    /// Constructing a new baseline.
219    Build,
220    /// Restoring captured canisters.
221    RestoreCanisters,
222    /// Resetting state outside the snapshots.
223    ResetNonSnapshotState,
224    /// Driving the restored topology to readiness.
225    DriveToReadiness,
226    /// Validating a newly built baseline.
227    ValidateBuilt,
228    /// Validating a restored baseline.
229    ValidateRestored,
230}
231
232/// Why an invalid or failed slot was reconstructed.
233#[non_exhaustive]
234#[derive(Clone, Debug, Eq, PartialEq)]
235pub enum RebuildReason {
236    /// PocketIC transport was no longer reachable.
237    DeadPocketIcTransport,
238    /// Captured snapshot restoration failed.
239    SnapshotRestoreFailure,
240    /// Non-snapshot reset failed.
241    ResetFailure,
242    /// Readiness or quiescence could not be established.
243    ReadinessFailure,
244    /// Required and achieved reset domains did not match.
245    ResetCoverageMismatch,
246    /// Final invariant validation failed.
247    InvariantValidationFailure,
248    /// A caller explicitly invalidated its lease.
249    ExplicitLeaseInvalidation,
250    /// A lease was dropped while its thread was unwinding.
251    UnwindWhileLeased,
252    /// Recipe-specific structured reason.
253    RecipeClassified { code: String },
254}
255
256/// Recipe decision for a failed restored-slot preparation stage.
257#[non_exhaustive]
258#[derive(Clone, Debug, Eq, PartialEq)]
259pub enum FailureDisposition {
260    /// Return the failure without rebuilding during this acquisition.
261    Fatal,
262    /// Invalidate and rebuild the slot once.
263    Rebuild(RebuildReason),
264}
265
266/// Timings for one baseline-pool acquisition.
267#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
268pub struct BaselinePoolTimings {
269    wait: Duration,
270    build: Option<Duration>,
271    restore: Option<Duration>,
272    reset: Option<Duration>,
273    readiness: Option<Duration>,
274    validation: Option<Duration>,
275    stale_teardown: Option<Duration>,
276    total: Duration,
277}
278
279/// Whether a baseline-pool lease was built, restored, or rebuilt.
280#[non_exhaustive]
281#[derive(Clone, Debug, Eq, PartialEq)]
282pub enum BaselinePoolOutcome {
283    /// An empty slot was constructed and validated.
284    Built {
285        /// Diagnostic slot index.
286        slot: usize,
287        /// Acquisition phase timings.
288        timings: BaselinePoolTimings,
289    },
290    /// An existing slot was restored, reset, and validated.
291    Restored {
292        /// Diagnostic slot index.
293        slot: usize,
294        /// Acquisition phase timings.
295        timings: BaselinePoolTimings,
296    },
297    /// An invalid or failed slot was reconstructed and validated.
298    Rebuilt {
299        /// Diagnostic slot index.
300        slot: usize,
301        /// Reason the previous slot could not be reused.
302        reason: RebuildReason,
303        /// Acquisition phase timings.
304        timings: BaselinePoolTimings,
305    },
306}
307
308/// One failed recipe or contract stage while preparing a baseline slot.
309#[non_exhaustive]
310#[derive(Debug)]
311pub enum BaselinePoolPreparationError<E> {
312    /// Caller-owned recipe logic returned an error.
313    Recipe {
314        /// Failed lifecycle stage.
315        stage: BaselinePreparationStage,
316        /// Caller-owned structured source error.
317        source: E,
318    },
319    /// Typed reset or recipe evidence violated the pool contract.
320    Contract(BaselinePoolContractError),
321}
322
323/// Failure to acquire a validated baseline-pool lease.
324#[non_exhaustive]
325#[derive(Debug)]
326pub enum BaselinePoolError<E> {
327    /// Initial construction or a non-rebuilt preparation failed.
328    Preparation {
329        /// Recipe or contract failure that stopped acquisition.
330        error: BaselinePoolPreparationError<E>,
331        /// Phase timings recorded before acquisition failed.
332        timings: Box<BaselinePoolTimings>,
333    },
334    /// Reused-slot preparation failed and its one rebuild attempt also failed.
335    RecoveryFailed {
336        /// Original restore/reset/readiness/validation failure.
337        original: Box<BaselinePoolPreparationError<E>>,
338        /// Failure while rebuilding or validating the replacement.
339        rebuild: Box<BaselinePoolPreparationError<E>>,
340        /// Combined timings for preparation, stale teardown, and rebuilding.
341        timings: Box<BaselinePoolTimings>,
342    },
343}
344
345/// Complete caller-owned lifecycle recipe for one pooled PocketIC baseline.
346pub trait PocketIcBaselineRecipe: Send + Sync + 'static {
347    /// Metadata retained beside every baseline owned by this recipe.
348    type Metadata: Send + 'static;
349    /// Structured caller error shared by recipe lifecycle stages.
350    type Error: std::error::Error + Send + Sync + 'static;
351
352    /// Stable caller-owned recipe identity.
353    fn id(&self) -> &FixtureRecipeId;
354
355    /// Reset guarantees required before an existing slot may be reused.
356    fn reset_requirements(&self) -> &ResetRequirements;
357
358    /// Construct and capture one complete baseline.
359    fn build(&self) -> Result<CachedPocketIcBaseline<Self::Metadata>, Self::Error>;
360
361    /// Restore every captured canister and report the cycle policy applied.
362    fn restore_canisters(
363        &self,
364        baseline: &CachedPocketIcBaseline<Self::Metadata>,
365    ) -> Result<CanisterRestoreReceipt, Self::Error>;
366
367    /// Reset state not covered by canister snapshots.
368    fn reset_non_snapshot_state(
369        &self,
370        baseline: &CachedPocketIcBaseline<Self::Metadata>,
371    ) -> Result<ResetReceipt, Self::Error>;
372
373    /// Drive the topology to the recipe's readiness boundary.
374    fn drive_to_readiness(
375        &self,
376        baseline: &CachedPocketIcBaseline<Self::Metadata>,
377    ) -> Result<ReadinessReceipt, Self::Error>;
378
379    /// Validate the same baseline invariants after build and restore.
380    fn validate(
381        &self,
382        baseline: &CachedPocketIcBaseline<Self::Metadata>,
383        preparation: &PreparedBaseline,
384    ) -> Result<ValidationReceipt, Self::Error>;
385
386    /// Classify a restored-slot recipe failure as fatal or rebuildable.
387    fn classify_failure(
388        &self,
389        stage: BaselinePreparationStage,
390        _error: &Self::Error,
391    ) -> FailureDisposition {
392        FailureDisposition::Rebuild(stage.default_rebuild_reason())
393    }
394}
395
396/// Caller-owned runtime-capacity pool of independently restorable PocketIC baselines.
397///
398/// One pool structurally owns one [`PocketIcBaselineRecipe`]. A warm
399/// acquisition restores the complete captured canister set, applies the
400/// recipe's non-snapshot reset, reaches its readiness boundary, checks typed
401/// reset coverage, and validates final invariants before exposing a lease.
402/// Each capacity slot owns an independent PocketIC instance.
403///
404/// Snapshot reuse is not a complete PocketIC rollback. The recipe must account
405/// for time, extra canisters, pending messages, subnet state, cycles, and
406/// external resources when those domains matter to its tests.
407pub struct CachedPocketIcBaselinePool<R>
408where
409    R: PocketIcBaselineRecipe,
410{
411    recipe: R,
412    slots: BoundedSlotPool<BaselineSlot<R::Metadata>>,
413}
414
415struct BaselineSlot<M> {
416    baseline: CachedPocketIcBaseline<M>,
417    invalidation_reason: Option<RebuildReason>,
418}
419
420/// Exclusive lease of one validated pooled PocketIC baseline.
421pub struct CachedPocketIcBaselinePoolGuard<'a, R>
422where
423    R: PocketIcBaselineRecipe,
424{
425    slot: BoundedSlotLease<'a, BaselineSlot<R::Metadata>>,
426}
427
428impl FixtureRecipeId {
429    /// Construct a nonempty caller-owned stable recipe identity.
430    pub fn try_new(identity: impl Into<String>) -> Result<Self, BaselinePoolContractError> {
431        let identity = identity.into();
432        if identity.trim().is_empty() {
433            return Err(BaselinePoolContractError::EmptyRecipeIdentity);
434        }
435        Ok(Self(identity))
436    }
437
438    /// Borrow the recipe identity.
439    #[must_use]
440    pub fn as_str(&self) -> &str {
441        &self.0
442    }
443}
444
445impl ResetRequirement {
446    /// Domain governed by this requirement.
447    #[must_use]
448    pub const fn domain(&self) -> ResetDomainKind {
449        match self {
450            Self::CanisterSnapshots => ResetDomainKind::CanisterSnapshots,
451            Self::CanisterCycles(_) => ResetDomainKind::CanisterCycles,
452            Self::PocketIcTime(_) => ResetDomainKind::PocketIcTime,
453            Self::ExtraCanisters(_) => ResetDomainKind::ExtraCanisters,
454            Self::PendingMessages(_) => ResetDomainKind::PendingMessages,
455            Self::SubnetState(_) => ResetDomainKind::SubnetState,
456            Self::ExternalResources(_) => ResetDomainKind::ExternalResources,
457        }
458    }
459}
460
461impl ResetAchievement {
462    /// Domain governed by this achievement.
463    #[must_use]
464    pub const fn domain(&self) -> ResetDomainKind {
465        match self {
466            Self::CanisterSnapshots => ResetDomainKind::CanisterSnapshots,
467            Self::CanisterCycles(_) => ResetDomainKind::CanisterCycles,
468            Self::PocketIcTime(_) => ResetDomainKind::PocketIcTime,
469            Self::ExtraCanisters(_) => ResetDomainKind::ExtraCanisters,
470            Self::PendingMessages(_) => ResetDomainKind::PendingMessages,
471            Self::SubnetState(_) => ResetDomainKind::SubnetState,
472            Self::ExternalResources(_) => ResetDomainKind::ExternalResources,
473        }
474    }
475
476    fn satisfies(&self, requirement: &ResetRequirement) -> bool {
477        match (requirement, self) {
478            (ResetRequirement::CanisterSnapshots, Self::CanisterSnapshots) => true,
479            (ResetRequirement::CanisterCycles(left), Self::CanisterCycles(right)) => left == right,
480            (ResetRequirement::PocketIcTime(left), Self::PocketIcTime(right)) => left == right,
481            (ResetRequirement::ExtraCanisters(left), Self::ExtraCanisters(right)) => left == right,
482            (ResetRequirement::PendingMessages(left), Self::PendingMessages(right))
483            | (ResetRequirement::SubnetState(left), Self::SubnetState(right))
484            | (ResetRequirement::ExternalResources(left), Self::ExternalResources(right)) => {
485                left == right
486            }
487            _ => false,
488        }
489    }
490}
491
492impl ResetRequirements {
493    /// Construct a duplicate-checked reset requirement set.
494    ///
495    /// Snapshot restoration and its cycle policy are mandatory for every
496    /// pooled baseline recipe.
497    pub fn try_new<I>(requirements: I) -> Result<Self, BaselinePoolContractError>
498    where
499        I: IntoIterator<Item = ResetRequirement>,
500    {
501        let mut domains = BTreeMap::new();
502        for requirement in requirements {
503            let domain = requirement.domain();
504            if domains.insert(domain, requirement).is_some() {
505                return Err(BaselinePoolContractError::DuplicateResetDomain { domain });
506            }
507        }
508        for domain in [
509            ResetDomainKind::CanisterSnapshots,
510            ResetDomainKind::CanisterCycles,
511        ] {
512            if !domains.contains_key(&domain) {
513                return Err(BaselinePoolContractError::UndeclaredRequiredResetDomain { domain });
514            }
515        }
516        Ok(Self(domains))
517    }
518
519    /// Read the requirement for one domain.
520    #[must_use]
521    pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetRequirement> {
522        self.0.get(&domain)
523    }
524
525    /// Iterate over requirements in deterministic domain order.
526    pub fn iter(&self) -> impl Iterator<Item = &ResetRequirement> {
527        self.0.values()
528    }
529
530    fn verify(&self, receipt: &ResetReceipt) -> Result<(), BaselinePoolContractError> {
531        for (domain, requirement) in &self.0 {
532            let Some(achievement) = receipt.0.get(domain) else {
533                return Err(BaselinePoolContractError::MissingResetDomain { domain: *domain });
534            };
535            if !achievement.satisfies(requirement) {
536                return Err(BaselinePoolContractError::ResetPolicyMismatch {
537                    requirement: requirement.clone(),
538                    achievement: achievement.clone(),
539                });
540            }
541        }
542        Ok(())
543    }
544}
545
546impl ResetReceipt {
547    /// Construct a duplicate-checked non-snapshot reset achievement set.
548    ///
549    /// Omit snapshot and cycle achievements from the receipt returned by
550    /// [`PocketIcBaselineRecipe::reset_non_snapshot_state`]; the pool derives
551    /// those from [`CanisterRestoreReceipt`].
552    pub fn try_new<I>(achievements: I) -> Result<Self, BaselinePoolContractError>
553    where
554        I: IntoIterator<Item = ResetAchievement>,
555    {
556        let mut domains = BTreeMap::new();
557        for achievement in achievements {
558            let domain = achievement.domain();
559            if domains.insert(domain, achievement).is_some() {
560                return Err(BaselinePoolContractError::DuplicateResetDomain { domain });
561            }
562        }
563        Ok(Self(domains))
564    }
565
566    /// Create an empty receipt for recipes with no non-snapshot reset achievements.
567    #[must_use]
568    pub const fn empty() -> Self {
569        Self(BTreeMap::new())
570    }
571
572    /// Read the achievement for one domain.
573    #[must_use]
574    pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetAchievement> {
575        self.0.get(&domain)
576    }
577
578    /// Iterate over achievements in deterministic domain order.
579    pub fn iter(&self) -> impl Iterator<Item = &ResetAchievement> {
580        self.0.values()
581    }
582
583    fn include_restore(
584        &mut self,
585        restore: &CanisterRestoreReceipt,
586    ) -> Result<(), BaselinePoolContractError> {
587        self.insert(ResetAchievement::CanisterSnapshots)?;
588        self.insert(ResetAchievement::CanisterCycles(restore.cycle_policy))
589    }
590
591    fn insert(&mut self, achievement: ResetAchievement) -> Result<(), BaselinePoolContractError> {
592        let domain = achievement.domain();
593        if self.0.insert(domain, achievement).is_some() {
594            return Err(BaselinePoolContractError::DuplicateResetDomain { domain });
595        }
596        Ok(())
597    }
598}
599
600impl CanisterRestoreReceipt {
601    /// Construct a deterministic, duplicate-checked canister restore receipt.
602    pub fn try_new<I>(
603        canister_ids: I,
604        cycle_policy: CycleResetPolicy,
605    ) -> Result<Self, BaselinePoolContractError>
606    where
607        I: IntoIterator<Item = Principal>,
608    {
609        let mut unique = BTreeSet::new();
610        for canister_id in canister_ids {
611            if !unique.insert(canister_id) {
612                return Err(BaselinePoolContractError::DuplicateCanisterId { canister_id });
613            }
614        }
615        if unique.is_empty() {
616            return Err(BaselinePoolContractError::EmptyCanisterSet);
617        }
618        Ok(Self {
619            canister_ids: unique.into_iter().collect(),
620            cycle_policy,
621        })
622    }
623
624    /// Construct a restore receipt for the exact canister set captured by a baseline.
625    ///
626    /// This is the preferred constructor after successfully calling
627    /// [`CachedPocketIcBaseline::restore`] or
628    /// [`CachedPocketIcBaseline::restore_with_funding`]. Deriving the set from
629    /// the baseline avoids duplicating canister ids in recipe metadata solely
630    /// to satisfy the pool contract.
631    pub fn try_from_baseline<M>(
632        baseline: &CachedPocketIcBaseline<M>,
633        cycle_policy: CycleResetPolicy,
634    ) -> Result<Self, BaselinePoolContractError> {
635        Self::try_new(baseline.snapshot_canister_ids(), cycle_policy)
636    }
637
638    /// Restored canister ids in deterministic order.
639    #[must_use]
640    pub fn canister_ids(&self) -> &[Principal] {
641        &self.canister_ids
642    }
643
644    /// Cycle policy applied while restoring canisters.
645    #[must_use]
646    pub const fn cycle_policy(&self) -> CycleResetPolicy {
647        self.cycle_policy
648    }
649}
650
651impl ReadinessReceipt {
652    /// Construct a nonempty caller-owned readiness identity.
653    pub fn try_new(identity: impl Into<String>) -> Result<Self, BaselinePoolContractError> {
654        Ok(Self {
655            identity: nonempty_receipt_identity("readiness", identity.into())?,
656        })
657    }
658
659    /// Borrow the readiness identity.
660    #[must_use]
661    pub fn identity(&self) -> &str {
662        &self.identity
663    }
664}
665
666impl ValidationReceipt {
667    /// Construct final validation evidence for one recipe.
668    pub fn try_new(
669        recipe_id: FixtureRecipeId,
670        invariant_identity: impl Into<String>,
671    ) -> Result<Self, BaselinePoolContractError> {
672        Ok(Self {
673            recipe_id,
674            invariant_identity: nonempty_receipt_identity("validation", invariant_identity.into())?,
675        })
676    }
677
678    /// Recipe identity validated by this receipt.
679    #[must_use]
680    pub const fn recipe_id(&self) -> &FixtureRecipeId {
681        &self.recipe_id
682    }
683
684    /// Borrow the caller-owned invariant identity.
685    #[must_use]
686    pub fn invariant_identity(&self) -> &str {
687        &self.invariant_identity
688    }
689}
690
691impl BaselinePreparationStage {
692    /// Default rebuild reason used by [`PocketIcBaselineRecipe::classify_failure`].
693    ///
694    /// Recipes that override classification can use this for their fallback
695    /// after handling a more specific error such as dead PocketIC transport.
696    #[must_use]
697    pub fn default_rebuild_reason(self) -> RebuildReason {
698        match self {
699            Self::RestoreCanisters => RebuildReason::SnapshotRestoreFailure,
700            Self::ResetNonSnapshotState => RebuildReason::ResetFailure,
701            Self::DriveToReadiness => RebuildReason::ReadinessFailure,
702            Self::ValidateRestored | Self::ValidateBuilt => {
703                RebuildReason::InvariantValidationFailure
704            }
705            Self::Build => RebuildReason::RecipeClassified {
706                code: "build".to_owned(),
707            },
708        }
709    }
710}
711
712impl BaselinePoolTimings {
713    /// Time spent waiting for a capacity slot.
714    #[must_use]
715    pub const fn wait(self) -> Duration {
716        self.wait
717    }
718
719    /// Time spent constructing a new baseline.
720    #[must_use]
721    pub const fn build(self) -> Option<Duration> {
722        self.build
723    }
724
725    /// Time spent restoring captured canisters.
726    #[must_use]
727    pub const fn restore(self) -> Option<Duration> {
728        self.restore
729    }
730
731    /// Time spent resetting non-snapshot state.
732    #[must_use]
733    pub const fn reset(self) -> Option<Duration> {
734        self.reset
735    }
736
737    /// Time spent driving the topology to readiness.
738    #[must_use]
739    pub const fn readiness(self) -> Option<Duration> {
740        self.readiness
741    }
742
743    /// Time spent validating final baseline invariants.
744    #[must_use]
745    pub const fn validation(self) -> Option<Duration> {
746        self.validation
747    }
748
749    /// Time spent dropping an invalid baseline before rebuilding.
750    #[must_use]
751    pub const fn stale_teardown(self) -> Option<Duration> {
752        self.stale_teardown
753    }
754
755    /// Complete acquisition duration.
756    #[must_use]
757    pub const fn total(self) -> Duration {
758        self.total
759    }
760}
761
762impl BaselinePoolOutcome {
763    /// Diagnostic slot index used by this acquisition.
764    #[must_use]
765    pub const fn slot(&self) -> usize {
766        match self {
767            Self::Built { slot, .. } | Self::Restored { slot, .. } | Self::Rebuilt { slot, .. } => {
768                *slot
769            }
770        }
771    }
772
773    /// Acquisition phase timings.
774    #[must_use]
775    pub const fn timings(&self) -> BaselinePoolTimings {
776        match self {
777            Self::Built { timings, .. }
778            | Self::Restored { timings, .. }
779            | Self::Rebuilt { timings, .. } => *timings,
780        }
781    }
782}
783
784impl<E> BaselinePoolError<E> {
785    /// Timings recorded before this acquisition failed.
786    #[must_use]
787    pub const fn timings(&self) -> BaselinePoolTimings {
788        match self {
789            Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
790        }
791    }
792}
793
794impl<R> CachedPocketIcBaselinePool<R>
795where
796    R: PocketIcBaselineRecipe,
797{
798    /// Create a runtime-capacity pool that structurally owns one recipe.
799    #[must_use]
800    pub fn new(capacity: NonZeroUsize, recipe: R) -> Self {
801        Self {
802            recipe,
803            slots: BoundedSlotPool::new(capacity),
804        }
805    }
806
807    /// Borrow the caller-owned identity of this pool's only recipe.
808    #[must_use]
809    pub fn recipe_id(&self) -> &FixtureRecipeId {
810        self.recipe.id()
811    }
812
813    /// Maximum number of simultaneously leased PocketIC baselines.
814    #[must_use]
815    pub fn capacity(&self) -> NonZeroUsize {
816        self.slots.capacity()
817    }
818
819    /// Acquire one fully built or restored and validated baseline lease.
820    ///
821    /// A rebuildable warm-preparation failure discards the stale slot and
822    /// performs at most one build attempt. Recipe and caller panics are never
823    /// converted into cache misses; the lease is invalidated while the panic
824    /// continues unwinding.
825    ///
826    /// # Errors
827    ///
828    /// Returns a stage-specific recipe or contract error. If warm preparation
829    /// and its one recovery build both fail, the error preserves both causes.
830    pub fn acquire(
831        &self,
832    ) -> Result<
833        (CachedPocketIcBaselinePoolGuard<'_, R>, BaselinePoolOutcome),
834        BaselinePoolError<R::Error>,
835    > {
836        let total_started = Instant::now();
837        let mut slot = self.slots.acquire();
838        let mut timings = BaselinePoolTimings {
839            wait: slot.wait(),
840            ..BaselinePoolTimings::default()
841        };
842        let slot_index = slot.slot_index();
843
844        if slot.is_reusable() {
845            match self.prepare_reused(&mut slot, &mut timings) {
846                Ok(()) => {
847                    timings.total = total_started.elapsed();
848                    return Ok((
849                        CachedPocketIcBaselinePoolGuard { slot },
850                        BaselinePoolOutcome::Restored {
851                            slot: slot_index,
852                            timings,
853                        },
854                    ));
855                }
856                Err(original) => {
857                    let disposition = self.failure_disposition(&original);
858                    match disposition {
859                        FailureDisposition::Fatal => {
860                            // A failed restore may have partially changed the
861                            // instance. Discard it, but do not reinterpret a
862                            // caller-declared fatal error as a rebuild reason.
863                            Self::discard_stale_slot(&mut slot, &mut timings);
864                            timings.total = total_started.elapsed();
865                            return Err(BaselinePoolError::Preparation {
866                                error: original,
867                                timings: Box::new(timings),
868                            });
869                        }
870                        FailureDisposition::Rebuild(reason) => {
871                            Self::discard_stale_slot(&mut slot, &mut timings);
872                            if let Err(rebuild) = self.build_slot(&mut slot, &mut timings) {
873                                timings.total = total_started.elapsed();
874                                return Err(BaselinePoolError::RecoveryFailed {
875                                    original: Box::new(original),
876                                    rebuild: Box::new(rebuild),
877                                    timings: Box::new(timings),
878                                });
879                            }
880                            timings.total = total_started.elapsed();
881                            return Ok((
882                                CachedPocketIcBaselinePoolGuard { slot },
883                                BaselinePoolOutcome::Rebuilt {
884                                    slot: slot_index,
885                                    reason,
886                                    timings,
887                                },
888                            ));
889                        }
890                    }
891                }
892            }
893        }
894
895        let rebuild_reason = if slot.invalidated_by_unwind() {
896            Some(RebuildReason::UnwindWhileLeased)
897        } else {
898            slot.get()
899                .and_then(|slot| slot.invalidation_reason.clone())
900                .or_else(|| {
901                    slot.is_populated()
902                        .then_some(RebuildReason::ExplicitLeaseInvalidation)
903                })
904        };
905        if slot.is_populated() {
906            Self::discard_stale_slot(&mut slot, &mut timings);
907        }
908        if let Err(error) = self.build_slot(&mut slot, &mut timings) {
909            timings.total = total_started.elapsed();
910            return Err(BaselinePoolError::Preparation {
911                error,
912                timings: Box::new(timings),
913            });
914        }
915        timings.total = total_started.elapsed();
916
917        let outcome = rebuild_reason.map_or_else(
918            || BaselinePoolOutcome::Built {
919                slot: slot_index,
920                timings,
921            },
922            |reason| BaselinePoolOutcome::Rebuilt {
923                slot: slot_index,
924                reason,
925                timings,
926            },
927        );
928        Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
929    }
930
931    fn prepare_reused(
932        &self,
933        slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
934        timings: &mut BaselinePoolTimings,
935    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
936        let baseline = &slot
937            .get()
938            .expect("reusable baseline slot must be populated")
939            .baseline;
940
941        let started = Instant::now();
942        let restore = self.recipe.restore_canisters(baseline);
943        add_timing(&mut timings.restore, started.elapsed());
944        let canisters = restore.map_err(|source| BaselinePoolPreparationError::Recipe {
945            stage: BaselinePreparationStage::RestoreCanisters,
946            source,
947        })?;
948        let expected_canisters = baseline.snapshot_canister_ids().collect::<Vec<_>>();
949        if canisters.canister_ids() != expected_canisters {
950            return Err(BaselinePoolPreparationError::Contract(
951                BaselinePoolContractError::RestoreCanisterSetMismatch {
952                    expected: expected_canisters,
953                    actual: canisters.canister_ids().to_vec(),
954                },
955            ));
956        }
957
958        let started = Instant::now();
959        let reset_result = self.recipe.reset_non_snapshot_state(baseline);
960        add_timing(&mut timings.reset, started.elapsed());
961        let mut reset = reset_result.map_err(|source| BaselinePoolPreparationError::Recipe {
962            stage: BaselinePreparationStage::ResetNonSnapshotState,
963            source,
964        })?;
965        reset
966            .include_restore(&canisters)
967            .map_err(BaselinePoolPreparationError::Contract)?;
968
969        let started = Instant::now();
970        let readiness_result = self.recipe.drive_to_readiness(baseline);
971        add_timing(&mut timings.readiness, started.elapsed());
972        let readiness =
973            readiness_result.map_err(|source| BaselinePoolPreparationError::Recipe {
974                stage: BaselinePreparationStage::DriveToReadiness,
975                source,
976            })?;
977        self.recipe
978            .reset_requirements()
979            .verify(&reset)
980            .map_err(BaselinePoolPreparationError::Contract)?;
981
982        let preparation = PreparedBaseline::Restored {
983            canisters,
984            reset,
985            readiness,
986        };
987        self.validate_baseline(
988            baseline,
989            &preparation,
990            BaselinePreparationStage::ValidateRestored,
991            timings,
992        )?;
993        slot.get_mut()
994            .expect("validated baseline slot must remain populated")
995            .invalidation_reason = None;
996        Ok(())
997    }
998
999    fn build_slot(
1000        &self,
1001        slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1002        timings: &mut BaselinePoolTimings,
1003    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1004        let started = Instant::now();
1005        let build = self.recipe.build();
1006        add_timing(&mut timings.build, started.elapsed());
1007        let baseline = build.map_err(|source| BaselinePoolPreparationError::Recipe {
1008            stage: BaselinePreparationStage::Build,
1009            source,
1010        })?;
1011
1012        if let Err(error) = self.validate_baseline(
1013            &baseline,
1014            &PreparedBaseline::Built,
1015            BaselinePreparationStage::ValidateBuilt,
1016            timings,
1017        ) {
1018            drop_baseline_safely(baseline);
1019            return Err(error);
1020        }
1021        let replaced = slot.replace(BaselineSlot {
1022            baseline,
1023            invalidation_reason: None,
1024        });
1025        debug_assert!(replaced.is_none());
1026        Ok(())
1027    }
1028
1029    fn validate_baseline(
1030        &self,
1031        baseline: &CachedPocketIcBaseline<R::Metadata>,
1032        preparation: &PreparedBaseline,
1033        stage: BaselinePreparationStage,
1034        timings: &mut BaselinePoolTimings,
1035    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
1036        let started = Instant::now();
1037        let validation = self.recipe.validate(baseline, preparation);
1038        add_timing(&mut timings.validation, started.elapsed());
1039        let receipt =
1040            validation.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })?;
1041        if receipt.recipe_id() != self.recipe.id() {
1042            return Err(BaselinePoolPreparationError::Contract(
1043                BaselinePoolContractError::RecipeIdentityMismatch {
1044                    expected: self.recipe.id().clone(),
1045                    actual: receipt.recipe_id().clone(),
1046                },
1047            ));
1048        }
1049        Ok(())
1050    }
1051
1052    fn failure_disposition(
1053        &self,
1054        error: &BaselinePoolPreparationError<R::Error>,
1055    ) -> FailureDisposition {
1056        match error {
1057            BaselinePoolPreparationError::Recipe { stage, source } => {
1058                self.recipe.classify_failure(*stage, source)
1059            }
1060            BaselinePoolPreparationError::Contract(
1061                BaselinePoolContractError::RecipeIdentityMismatch { .. },
1062            ) => FailureDisposition::Fatal,
1063            BaselinePoolPreparationError::Contract(_) => {
1064                FailureDisposition::Rebuild(rebuild_reason_for_error(error))
1065            }
1066        }
1067    }
1068
1069    fn discard_stale_slot(
1070        slot: &mut BoundedSlotLease<'_, BaselineSlot<R::Metadata>>,
1071        timings: &mut BaselinePoolTimings,
1072    ) {
1073        let started = Instant::now();
1074        if let Some(stale) = slot.take() {
1075            drop_baseline_safely(stale.baseline);
1076        }
1077        timings.stale_teardown = Some(started.elapsed());
1078    }
1079}
1080
1081impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
1082where
1083    R: PocketIcBaselineRecipe,
1084{
1085    /// Diagnostic slot index held by this lease.
1086    #[must_use]
1087    pub const fn slot(&self) -> usize {
1088        self.slot.slot_index()
1089    }
1090
1091    /// Mark this slot non-reusable with a structured rebuild reason.
1092    pub fn invalidate(&mut self, reason: RebuildReason) {
1093        if let Some(slot) = self.slot.get_mut() {
1094            slot.invalidation_reason = Some(reason);
1095        }
1096        self.slot.invalidate();
1097    }
1098}
1099
1100impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
1101where
1102    R: PocketIcBaselineRecipe,
1103{
1104    type Target = CachedPocketIcBaseline<R::Metadata>;
1105
1106    fn deref(&self) -> &Self::Target {
1107        &self
1108            .slot
1109            .get()
1110            .expect("leased baseline pool slot must be populated")
1111            .baseline
1112    }
1113}
1114
1115fn nonempty_receipt_identity(
1116    receipt: &'static str,
1117    identity: String,
1118) -> Result<String, BaselinePoolContractError> {
1119    if identity.trim().is_empty() {
1120        return Err(BaselinePoolContractError::EmptyReceiptIdentity { receipt });
1121    }
1122    Ok(identity)
1123}
1124
1125fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
1126    *total = Some(total.unwrap_or_default().saturating_add(elapsed));
1127}
1128
1129fn rebuild_reason_for_error<E>(error: &BaselinePoolPreparationError<E>) -> RebuildReason {
1130    match error {
1131        BaselinePoolPreparationError::Recipe { stage, .. } => stage.default_rebuild_reason(),
1132        BaselinePoolPreparationError::Contract(
1133            BaselinePoolContractError::MissingResetDomain { .. }
1134            | BaselinePoolContractError::ResetPolicyMismatch { .. }
1135            | BaselinePoolContractError::DuplicateResetDomain { .. }
1136            | BaselinePoolContractError::RestoreCanisterSetMismatch { .. },
1137        ) => RebuildReason::ResetCoverageMismatch,
1138        BaselinePoolPreparationError::Contract(_) => RebuildReason::InvariantValidationFailure,
1139    }
1140}
1141
1142fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
1143    let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
1144}
1145
1146impl std::fmt::Display for FixtureRecipeId {
1147    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1148        formatter.write_str(&self.0)
1149    }
1150}
1151
1152impl std::fmt::Display for BaselinePreparationStage {
1153    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1154        formatter.write_str(match self {
1155            Self::Build => "build",
1156            Self::RestoreCanisters => "canister restore",
1157            Self::ResetNonSnapshotState => "non-snapshot reset",
1158            Self::DriveToReadiness => "readiness",
1159            Self::ValidateBuilt => "built-baseline validation",
1160            Self::ValidateRestored => "restored-baseline validation",
1161        })
1162    }
1163}
1164
1165impl std::fmt::Display for BaselinePoolContractError {
1166    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1167        match self {
1168            Self::EmptyRecipeIdentity => formatter.write_str("fixture recipe identity is empty"),
1169            Self::EmptyReceiptIdentity { receipt } => {
1170                write!(formatter, "{receipt} receipt identity is empty")
1171            }
1172            Self::DuplicateResetDomain { domain } => {
1173                write!(
1174                    formatter,
1175                    "reset domain {domain:?} was reported more than once"
1176                )
1177            }
1178            Self::DuplicateCanisterId { canister_id } => {
1179                write!(formatter, "restore receipt repeats canister {canister_id}")
1180            }
1181            Self::EmptyCanisterSet => formatter.write_str("restore receipt contains no canisters"),
1182            Self::UndeclaredRequiredResetDomain { domain } => write!(
1183                formatter,
1184                "baseline recipe does not declare required reset domain {domain:?}",
1185            ),
1186            Self::RestoreCanisterSetMismatch { expected, actual } => write!(
1187                formatter,
1188                "restore receipt identified canisters {actual:?}, expected captured set {expected:?}",
1189            ),
1190            Self::MissingResetDomain { domain } => {
1191                write!(
1192                    formatter,
1193                    "required reset domain {domain:?} was not achieved"
1194                )
1195            }
1196            Self::ResetPolicyMismatch {
1197                requirement,
1198                achievement,
1199            } => write!(
1200                formatter,
1201                "reset achievement {achievement:?} does not satisfy {requirement:?}",
1202            ),
1203            Self::RecipeIdentityMismatch { expected, actual } => write!(
1204                formatter,
1205                "validation receipt used recipe `{actual}` instead of `{expected}`",
1206            ),
1207        }
1208    }
1209}
1210
1211impl std::error::Error for BaselinePoolContractError {}
1212
1213impl<E> std::fmt::Display for BaselinePoolPreparationError<E>
1214where
1215    E: std::fmt::Display,
1216{
1217    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1218        match self {
1219            Self::Recipe { stage, source } => {
1220                write!(formatter, "baseline {stage} failed: {source}")
1221            }
1222            Self::Contract(error) => write!(formatter, "baseline pool contract failed: {error}"),
1223        }
1224    }
1225}
1226
1227impl<E> std::error::Error for BaselinePoolPreparationError<E>
1228where
1229    E: std::error::Error + 'static,
1230{
1231    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1232        match self {
1233            Self::Recipe { source, .. } => Some(source),
1234            Self::Contract(error) => Some(error),
1235        }
1236    }
1237}
1238
1239impl<E> std::fmt::Display for BaselinePoolError<E>
1240where
1241    E: std::fmt::Display,
1242{
1243    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1244        match self {
1245            Self::Preparation { error, .. } => error.fmt(formatter),
1246            Self::RecoveryFailed {
1247                original, rebuild, ..
1248            } => write!(
1249                formatter,
1250                "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
1251            ),
1252        }
1253    }
1254}
1255
1256impl<E> std::error::Error for BaselinePoolError<E>
1257where
1258    E: std::error::Error + 'static,
1259{
1260    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1261        match self {
1262            Self::Preparation { error, .. } => Some(error),
1263            Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
1264        }
1265    }
1266}
1267
1268#[cfg(test)]
1269mod tests {
1270    use super::{
1271        BaselinePoolContractError, CycleResetPolicy, FixtureRecipeId, ResetAchievement,
1272        ResetDomainKind, ResetReceipt, ResetRequirement, ResetRequirements,
1273    };
1274
1275    #[test]
1276    fn recipe_identity_must_be_nonempty() {
1277        assert!(matches!(
1278            FixtureRecipeId::try_new("  "),
1279            Err(BaselinePoolContractError::EmptyRecipeIdentity)
1280        ));
1281    }
1282
1283    #[test]
1284    fn reset_requirements_reject_duplicate_domains() {
1285        let result = ResetRequirements::try_new([
1286            ResetRequirement::CanisterCycles(CycleResetPolicy::PreserveCurrent),
1287            ResetRequirement::CanisterCycles(CycleResetPolicy::RestoreExactBaseline),
1288        ]);
1289        assert!(matches!(
1290            result,
1291            Err(BaselinePoolContractError::DuplicateResetDomain {
1292                domain: ResetDomainKind::CanisterCycles,
1293            })
1294        ));
1295    }
1296
1297    #[test]
1298    fn required_policy_must_match_achieved_policy() {
1299        let requirements = ResetRequirements::try_new([
1300            ResetRequirement::CanisterSnapshots,
1301            ResetRequirement::CanisterCycles(CycleResetPolicy::RestoreExactBaseline),
1302        ])
1303        .unwrap();
1304        let receipt = ResetReceipt::try_new([
1305            ResetAchievement::CanisterSnapshots,
1306            ResetAchievement::CanisterCycles(CycleResetPolicy::PreserveCurrent),
1307        ])
1308        .unwrap();
1309
1310        assert!(matches!(
1311            requirements.verify(&receipt),
1312            Err(BaselinePoolContractError::ResetPolicyMismatch { .. })
1313        ));
1314    }
1315}