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, ControllerSnapshots, SnapshotRestoreFunding,
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
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/// Reset domain whose handling is declared by a pooled baseline recipe.
29#[non_exhaustive]
30#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
31pub enum ResetDomainKind {
32    /// PocketIC simulated time.
33    PocketIcTime,
34    /// Canisters outside the captured baseline set.
35    ExtraCanisters,
36    /// Pending ingress, timers, or cross-canister messages.
37    PendingMessages,
38    /// Subnet metrics, routing, allocation, or other subnet-global state.
39    SubnetState,
40    /// Files, processes, services, or other caller-owned resources.
41    ExternalResources,
42}
43
44/// Cycle handling required or achieved by one reset.
45#[non_exhaustive]
46#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum CycleResetPolicy {
48    /// Do not proactively add or remove cycles before snapshot restoration.
49    ///
50    /// PocketIC may still charge cycles while performing the restore.
51    PreserveCurrent,
52    /// Add cycles as needed to reach this minimum immediately before restore,
53    /// without removing excess.
54    TopUpTo(u128),
55    /// Restore the exact balance recorded by the recipe baseline.
56    RestoreExactBaseline,
57    /// Treat any relevant cycle mutation as requiring slot reconstruction.
58    RebuildOnMutation,
59}
60
61/// PocketIC time handling required or achieved by one reset.
62#[non_exhaustive]
63#[derive(Clone, Copy, Debug, Eq, PartialEq)]
64pub enum TimeResetPolicy {
65    /// Preserve the current simulator time rather than claiming it was reset.
66    PreserveCurrent,
67    /// Restore the exact time recorded by the recipe baseline.
68    RestoreBaseline,
69    /// Treat any relevant time mutation as requiring slot reconstruction.
70    RebuildOnMutation,
71}
72
73/// Extra-canister handling required or achieved by one reset.
74#[non_exhaustive]
75#[derive(Clone, Copy, Debug, Eq, PartialEq)]
76pub enum ExtraCanisterPolicy {
77    /// Validate that the baseline canister set is unchanged.
78    RequireBaselineSet,
79    /// Remove canisters explicitly tracked by the recipe.
80    RemoveTracked,
81    /// Treat any extra-canister change as requiring slot reconstruction.
82    RebuildOnChange,
83}
84
85/// Generic handling for reset domains without a more specific policy.
86#[non_exhaustive]
87#[derive(Clone, Copy, Debug, Eq, PartialEq)]
88pub enum StateResetPolicy {
89    /// Reset the domain through caller-owned recipe logic.
90    ResetByRecipe,
91    /// Validate that the domain remained unchanged.
92    ValidateUnchanged,
93    /// Explicitly declare the domain irrelevant to this recipe's guarantees.
94    IrrelevantByRecipeContract,
95    /// Treat any relevant change as requiring slot reconstruction.
96    RebuildOnChange,
97}
98
99/// Policy for one non-snapshot reset domain.
100///
101/// [`ResetRequirements`] declares policies that a recipe requires;
102/// [`ResetReceipt`] reports policies achieved by its reset operation. The pool
103/// checks every required domain against the reported policy before reuse.
104///
105/// ```
106/// use ic_testkit::pic::{
107///     CycleResetPolicy, ResetDomainKind, ResetDomainPolicy, ResetReceipt,
108///     ResetRequirements, TimeResetPolicy,
109/// };
110///
111/// let requirements = ResetRequirements::try_new(
112///     CycleResetPolicy::PreserveCurrent,
113///     [ResetDomainPolicy::PocketIcTime(TimeResetPolicy::PreserveCurrent)],
114/// )?;
115/// // The reset operation deliberately preserves simulated time.
116/// let receipt = ResetReceipt::try_new([
117///     ResetDomainPolicy::PocketIcTime(TimeResetPolicy::PreserveCurrent),
118/// ])?;
119/// assert_eq!(
120///     requirements.get(ResetDomainKind::PocketIcTime),
121///     receipt.get(ResetDomainKind::PocketIcTime),
122/// );
123/// # Ok::<(), ic_testkit::pic::BaselinePoolContractError>(())
124/// ```
125#[non_exhaustive]
126#[derive(Clone, Debug, Eq, PartialEq)]
127pub enum ResetDomainPolicy {
128    /// PocketIC time policy.
129    PocketIcTime(TimeResetPolicy),
130    /// Extra-canister policy.
131    ExtraCanisters(ExtraCanisterPolicy),
132    /// Pending-message policy.
133    PendingMessages(StateResetPolicy),
134    /// Subnet-state policy.
135    SubnetState(StateResetPolicy),
136    /// External-resource policy.
137    ExternalResources(StateResetPolicy),
138}
139
140/// Typed reset guarantees required by one fixture recipe.
141#[derive(Clone, Debug, Eq, PartialEq)]
142pub struct ResetRequirements {
143    cycle_policy: CycleResetPolicy,
144    domains: BTreeMap<ResetDomainKind, ResetDomainPolicy>,
145}
146
147/// Typed reset guarantees achieved by one preparation pass.
148#[derive(Clone, Debug, Default, Eq, PartialEq)]
149pub struct ResetReceipt(BTreeMap<ResetDomainKind, ResetDomainPolicy>);
150
151/// Receipt for restoring the recipe's captured canister set.
152#[derive(Clone, Debug, Eq, PartialEq)]
153pub struct CanisterRestoreReceipt {
154    canister_ids: Vec<Principal>,
155    cycle_policy: CycleResetPolicy,
156}
157
158/// Contract failure while constructing or verifying recipe reset evidence.
159#[non_exhaustive]
160#[derive(Clone, Debug, Eq, PartialEq)]
161pub enum BaselinePoolContractError {
162    /// Recipe identity was empty or whitespace-only.
163    EmptyRecipeIdentity,
164    /// A reset domain was declared more than once.
165    DuplicateResetDomain { domain: ResetDomainKind },
166    /// A restored canister appeared more than once.
167    DuplicateCanisterId { canister_id: Principal },
168    /// A restore receipt contained no canisters.
169    EmptyCanisterSet,
170    /// The restore receipt did not satisfy the required cycle policy.
171    CyclePolicyMismatch {
172        required: CycleResetPolicy,
173        achieved: CycleResetPolicy,
174    },
175    /// The restored canister receipt did not identify the complete snapshot set.
176    RestoreCanisterSetMismatch {
177        expected: Vec<Principal>,
178        actual: Vec<Principal>,
179    },
180    /// A required reset domain had no matching achievement.
181    MissingResetDomain { domain: ResetDomainKind },
182    /// A reset achievement did not satisfy the required policy.
183    ResetPolicyMismatch {
184        requirement: ResetDomainPolicy,
185        achievement: ResetDomainPolicy,
186    },
187}
188
189/// Whether validation is observing a newly built or restored baseline.
190#[non_exhaustive]
191#[derive(Clone, Debug, Eq, PartialEq)]
192pub enum PreparedBaseline {
193    /// The recipe just built this baseline.
194    Built,
195    /// The recipe restored and reset an existing baseline.
196    Restored {
197        /// Captured canisters restored by the recipe.
198        canisters: CanisterRestoreReceipt,
199        /// Typed non-snapshot reset receipt.
200        reset: ResetReceipt,
201    },
202}
203
204/// Recipe stage associated with a structured preparation failure.
205#[non_exhaustive]
206#[derive(Clone, Copy, Debug, Eq, PartialEq)]
207pub enum BaselinePreparationStage {
208    /// Constructing a new baseline.
209    Build,
210    /// Restoring captured canisters.
211    RestoreCanisters,
212    /// Resetting state outside the snapshots.
213    ResetNonSnapshotState,
214    /// Driving the restored topology to readiness.
215    DriveToReadiness,
216    /// Validating a newly built baseline.
217    ValidateBuilt,
218    /// Validating a restored baseline.
219    ValidateRestored,
220}
221
222/// Why an invalid or failed slot was reconstructed.
223#[non_exhaustive]
224#[derive(Clone, Debug, Eq, PartialEq)]
225pub enum RebuildReason {
226    /// PocketIC transport was no longer reachable.
227    DeadPocketIcTransport,
228    /// Captured snapshot restoration failed.
229    SnapshotRestoreFailure,
230    /// Non-snapshot reset failed.
231    ResetFailure,
232    /// Readiness or quiescence could not be established.
233    ReadinessFailure,
234    /// Required and achieved reset domains did not match.
235    ResetCoverageMismatch,
236    /// Final invariant validation failed.
237    InvariantValidationFailure,
238    /// A caller explicitly invalidated its lease.
239    ExplicitLeaseInvalidation,
240    /// A lease was dropped while its thread was unwinding.
241    UnwindWhileLeased,
242    /// Recipe-specific structured reason.
243    RecipeClassified { code: String },
244}
245
246/// Recipe decision for a failed restored-slot preparation stage.
247#[non_exhaustive]
248#[derive(Clone, Debug, Eq, PartialEq)]
249pub enum FailureDisposition {
250    /// Return the failure without rebuilding during this acquisition.
251    Fatal,
252    /// Invalidate and rebuild the slot once.
253    Rebuild(RebuildReason),
254}
255
256/// Timings for one baseline-pool acquisition.
257#[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/// One successful acquisition's slot, timings and preparation disposition.
270#[derive(Clone, Debug, Eq, PartialEq)]
271pub struct BaselinePoolOutcome {
272    /// Diagnostic slot index.
273    pub slot: usize,
274    /// Acquisition phase timings.
275    pub timings: BaselinePoolTimings,
276    /// Whether this acquisition built, restored or rebuilt the baseline.
277    pub disposition: BaselinePoolDisposition,
278}
279
280/// Whether a baseline-pool lease was built, restored, or rebuilt.
281#[non_exhaustive]
282#[derive(Clone, Debug, Eq, PartialEq)]
283pub enum BaselinePoolDisposition {
284    /// An empty slot was constructed and validated.
285    Built,
286    /// An existing slot was restored, reset, and validated.
287    Restored,
288    /// An invalid or failed slot was reconstructed and validated.
289    Rebuilt(RebuildReason),
290}
291
292/// One failed recipe or contract stage while preparing a baseline slot.
293#[non_exhaustive]
294#[derive(Debug)]
295pub enum BaselinePoolPreparationError<E> {
296    /// Caller-owned recipe logic returned an error.
297    Recipe {
298        /// Failed lifecycle stage.
299        stage: BaselinePreparationStage,
300        /// Caller-owned structured source error.
301        source: E,
302    },
303    /// Typed reset or recipe evidence violated the pool contract.
304    Contract(BaselinePoolContractError),
305}
306
307/// Failure to acquire a validated baseline-pool lease.
308#[non_exhaustive]
309#[derive(Debug)]
310pub enum BaselinePoolError<E> {
311    /// Initial construction or a non-rebuilt preparation failed.
312    Preparation {
313        /// Recipe or contract failure that stopped acquisition.
314        error: BaselinePoolPreparationError<E>,
315        /// Phase timings recorded before acquisition failed.
316        timings: Box<BaselinePoolTimings>,
317    },
318    /// Reused-slot preparation failed and its one rebuild attempt also failed.
319    RecoveryFailed {
320        /// Original restore/reset/readiness/validation failure.
321        original: Box<BaselinePoolPreparationError<E>>,
322        /// Failure while rebuilding or validating the replacement.
323        rebuild: Box<BaselinePoolPreparationError<E>>,
324        /// Combined timings for preparation, stale teardown, and rebuilding.
325        timings: Box<BaselinePoolTimings>,
326    },
327}
328
329/// Complete caller-owned lifecycle recipe for one pooled PocketIC baseline.
330pub trait PocketIcBaselineRecipe: Send + Sync + 'static {
331    /// Metadata retained beside every baseline owned by this recipe.
332    type Metadata: Send + 'static;
333    /// Structured caller error shared by recipe lifecycle stages.
334    type Error: std::error::Error + Send + Sync + 'static;
335
336    /// Stable caller-owned recipe identity.
337    fn id(&self) -> &FixtureRecipeId;
338
339    /// Reset guarantees required before an existing slot may be reused.
340    fn reset_requirements(&self) -> &ResetRequirements;
341
342    /// Construct and capture one complete baseline.
343    fn build(&self) -> Result<CachedPocketIcBaseline<Self::Metadata>, Self::Error>;
344
345    /// Restore every captured canister and report the cycle policy applied.
346    fn restore_canisters(
347        &self,
348        baseline: &CachedPocketIcBaseline<Self::Metadata>,
349    ) -> Result<CanisterRestoreReceipt, Self::Error>;
350
351    /// Reset state not covered by canister snapshots.
352    fn reset_non_snapshot_state(
353        &self,
354        baseline: &CachedPocketIcBaseline<Self::Metadata>,
355    ) -> Result<ResetReceipt, Self::Error>;
356
357    /// Drive the topology to the recipe's readiness boundary.
358    /// Successful completion permits final invariant validation.
359    fn drive_to_readiness(
360        &self,
361        baseline: &CachedPocketIcBaseline<Self::Metadata>,
362    ) -> Result<(), Self::Error>;
363
364    /// Validate the same baseline invariants after build and restore.
365    /// Success admits this pool-owned recipe's baseline for leasing.
366    fn validate(
367        &self,
368        baseline: &CachedPocketIcBaseline<Self::Metadata>,
369        preparation: &PreparedBaseline,
370    ) -> Result<(), Self::Error>;
371
372    /// Classify a restored-slot recipe failure as fatal or rebuildable.
373    fn classify_failure(
374        &self,
375        stage: BaselinePreparationStage,
376        _error: &Self::Error,
377    ) -> FailureDisposition {
378        FailureDisposition::Rebuild(stage.default_rebuild_reason())
379    }
380}
381
382/// Caller-owned runtime-capacity pool of independently restorable PocketIC baselines.
383///
384/// One pool structurally owns one [`PocketIcBaselineRecipe`]. A warm
385/// acquisition restores the complete captured canister set, applies the
386/// recipe's non-snapshot reset, reaches its readiness boundary, checks typed
387/// reset coverage, and validates final invariants before exposing a lease.
388/// Each capacity slot owns an independent PocketIC instance.
389///
390/// Snapshot reuse is not a complete PocketIC rollback. The recipe must account
391/// for time, extra canisters, pending messages, subnet state, cycles, and
392/// external resources when those domains matter to its tests.
393pub struct CachedPocketIcBaselinePool<R>
394where
395    R: PocketIcBaselineRecipe,
396{
397    recipe: R,
398    slots: BoundedSlotPool<CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
399}
400
401/// Exclusive lease of one validated pooled PocketIC baseline.
402pub struct CachedPocketIcBaselinePoolGuard<'a, R>
403where
404    R: PocketIcBaselineRecipe,
405{
406    slot: BoundedSlotLease<'a, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
407}
408
409impl FixtureRecipeId {
410    /// Construct a nonempty caller-owned stable recipe identity.
411    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    /// Borrow the recipe identity.
420    #[must_use]
421    pub fn as_str(&self) -> &str {
422        &self.0
423    }
424}
425
426impl ResetDomainPolicy {
427    /// Domain governed by this policy.
428    #[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    /// Construct a duplicate-checked reset requirement set.
455    ///
456    /// Every recipe restores its complete snapshot set. The required cycle
457    /// policy is explicit; `requirements` describe only non-snapshot domains.
458    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    /// Cycle policy required of the canister restore receipt.
472    #[must_use]
473    pub const fn cycle_policy(&self) -> CycleResetPolicy {
474        self.cycle_policy
475    }
476
477    /// Read the requirement for one domain.
478    #[must_use]
479    pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetDomainPolicy> {
480        self.domains.get(&domain)
481    }
482
483    /// Iterate over requirements in deterministic domain order.
484    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    /// Construct a duplicate-checked non-snapshot reset achievement set.
516    ///
517    /// Snapshot restoration and cycle policy are verified separately through
518    /// [`CanisterRestoreReceipt`].
519    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    /// Create an empty receipt for recipes with no non-snapshot reset achievements.
527    #[must_use]
528    pub const fn empty() -> Self {
529        Self(BTreeMap::new())
530    }
531
532    /// Read the achievement for one domain.
533    #[must_use]
534    pub fn get(&self, domain: ResetDomainKind) -> Option<&ResetDomainPolicy> {
535        self.0.get(&domain)
536    }
537
538    /// Iterate over achievements in deterministic domain order.
539    pub fn iter(&self) -> impl Iterator<Item = &ResetDomainPolicy> {
540        self.0.values()
541    }
542}
543
544impl CanisterRestoreReceipt {
545    /// Construct a deterministic, duplicate-checked canister restore receipt.
546    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    // Only a completed baseline restore produces this receipt. Capture already
569    // admitted the nonempty set, duplicate validation and deterministic order.
570    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    /// Restored canister ids in deterministic order.
581    #[must_use]
582    pub fn canister_ids(&self) -> &[Principal] {
583        &self.canister_ids
584    }
585
586    /// Cycle policy applied while restoring canisters.
587    #[must_use]
588    pub const fn cycle_policy(&self) -> CycleResetPolicy {
589        self.cycle_policy
590    }
591}
592
593impl BaselinePreparationStage {
594    /// Default rebuild reason used by [`PocketIcBaselineRecipe::classify_failure`].
595    ///
596    /// Recipes that override classification can use this for their fallback
597    /// after handling a more specific error such as dead PocketIC transport.
598    #[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    /// Time spent waiting for a capacity slot.
616    #[must_use]
617    pub const fn wait(self) -> Duration {
618        self.wait
619    }
620
621    /// Time spent constructing a new baseline.
622    #[must_use]
623    pub const fn build(self) -> Option<Duration> {
624        self.build
625    }
626
627    /// Time spent restoring captured canisters.
628    #[must_use]
629    pub const fn restore(self) -> Option<Duration> {
630        self.restore
631    }
632
633    /// Time spent resetting non-snapshot state.
634    #[must_use]
635    pub const fn reset(self) -> Option<Duration> {
636        self.reset
637    }
638
639    /// Time spent driving the topology to readiness.
640    #[must_use]
641    pub const fn readiness(self) -> Option<Duration> {
642        self.readiness
643    }
644
645    /// Time spent validating final baseline invariants.
646    #[must_use]
647    pub const fn validation(self) -> Option<Duration> {
648        self.validation
649    }
650
651    /// Time spent dropping an invalid baseline before rebuilding.
652    #[must_use]
653    pub const fn stale_teardown(self) -> Option<Duration> {
654        self.stale_teardown
655    }
656
657    /// Complete acquisition duration.
658    #[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    /// Timings recorded before this acquisition failed.
701    #[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    /// Create a runtime-capacity pool that structurally owns one recipe.
714    ///
715    /// Construction is const and lazy: slots are allocated on first acquisition,
716    /// and the recipe builds a baseline only when a slot needs one.
717    #[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    /// Borrow the caller-owned identity of this pool's only recipe.
726    #[must_use]
727    pub fn recipe_id(&self) -> &FixtureRecipeId {
728        self.recipe.id()
729    }
730
731    /// Maximum number of simultaneously leased PocketIC baselines.
732    #[must_use]
733    pub const fn capacity(&self) -> NonZeroUsize {
734        self.slots.capacity()
735    }
736
737    /// Acquire one fully built or restored and validated baseline lease.
738    ///
739    /// A rebuildable warm-preparation failure discards the stale slot and
740    /// performs at most one build attempt. Recipe and caller panics are never
741    /// converted into cache misses; the lease is invalidated while the panic
742    /// continues unwinding.
743    ///
744    /// # Errors
745    ///
746    /// Returns a stage-specific recipe or contract error. If warm preparation
747    /// and its one recovery build both fail, the error preserves both causes.
748    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                        // A failed restore may have partially changed the instance.
778                        // Discard it without treating a fatal failure as a rebuild.
779                        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    /// Diagnostic slot index held by this lease.
966    #[must_use]
967    pub const fn slot(&self) -> usize {
968        self.slot.slot_index()
969    }
970
971    /// Mark this slot non-reusable with a structured rebuild reason.
972    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}