ic_testkit/pic/
baseline_pool.rs

1use std::{
2    num::NonZeroUsize,
3    ops::Deref,
4    panic::{AssertUnwindSafe, catch_unwind},
5    time::{Duration, Instant},
6};
7
8use super::{
9    CachedPocketIcBaseline, ControllerSnapshotError, SnapshotRestoreFunding,
10    bounded_pool::{BoundedSlotLease, BoundedSlotPool, Slot},
11};
12
13/// Reset domain whose handling is declared by a pooled baseline recipe.
14#[non_exhaustive]
15#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
16pub enum ResetDomainKind {
17    /// `PocketIC` simulated time.
18    PocketIcTime,
19    /// Canisters outside the captured baseline set.
20    ExtraCanisters,
21    /// Pending ingress, timers, or cross-canister messages.
22    PendingMessages,
23    /// Subnet metrics, routing, allocation, or other subnet-global state.
24    SubnetState,
25    /// Files, processes, services, or other caller-owned resources.
26    ExternalResources,
27}
28
29/// `PocketIC` time handling required or achieved by one reset.
30#[non_exhaustive]
31#[derive(Clone, Copy, Debug, Eq, PartialEq)]
32pub enum TimeResetPolicy {
33    /// Preserve the current simulator time rather than claiming it was reset.
34    PreserveCurrent,
35    /// Restore the exact time recorded by the recipe baseline.
36    RestoreBaseline,
37    /// Treat any relevant time mutation as requiring slot reconstruction.
38    RebuildOnMutation,
39}
40
41/// Extra-canister handling required or achieved by one reset.
42#[non_exhaustive]
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
44pub enum ExtraCanisterPolicy {
45    /// Validate that the baseline canister set is unchanged.
46    RequireBaselineSet,
47    /// Remove canisters explicitly tracked by the recipe.
48    RemoveTracked,
49    /// Treat any extra-canister change as requiring slot reconstruction.
50    RebuildOnChange,
51}
52
53/// Generic handling for reset domains without a more specific policy.
54#[non_exhaustive]
55#[derive(Clone, Copy, Debug, Eq, PartialEq)]
56pub enum StateResetPolicy {
57    /// Reset the domain through caller-owned recipe logic.
58    ResetByRecipe,
59    /// Validate that the domain remained unchanged.
60    ValidateUnchanged,
61    /// Explicitly declare the domain irrelevant to this recipe's guarantees.
62    IrrelevantByRecipeContract,
63    /// Treat any relevant change as requiring slot reconstruction.
64    RebuildOnChange,
65}
66
67/// Fixed non-snapshot reset policies, supplied independently by a recipe and its reset.
68///
69/// A missing field makes no claim for that domain. In particular, absence differs
70/// from explicitly reporting [`StateResetPolicy::IrrelevantByRecipeContract`].
71/// Each domain has one typed field, so duplicate declarations cannot be formed.
72///
73/// ```
74/// use ic_testkit::pic::{
75///     ResetPolicies, ResetReceipt, ResetRequirements, SnapshotRestoreFunding, TimeResetPolicy,
76/// };
77///
78/// let requirements = ResetRequirements {
79///     funding: SnapshotRestoreFunding::Preserve,
80///     policies: ResetPolicies {
81///         pocket_ic_time: Some(TimeResetPolicy::PreserveCurrent),
82///         ..ResetPolicies::default()
83///     },
84/// };
85/// // The reset independently reports that it preserved simulated time.
86/// let receipt = ResetReceipt {
87///     policies: ResetPolicies {
88///         pocket_ic_time: Some(TimeResetPolicy::PreserveCurrent),
89///         ..ResetPolicies::default()
90///     },
91/// };
92/// assert_eq!(requirements.policies.pocket_ic_time, receipt.policies.pocket_ic_time);
93/// ```
94#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
95pub struct ResetPolicies {
96    /// `PocketIC` simulated time.
97    pub pocket_ic_time: Option<TimeResetPolicy>,
98    /// Canisters outside the captured baseline set.
99    pub extra_canisters: Option<ExtraCanisterPolicy>,
100    /// Pending ingress, timers, or cross-canister messages.
101    pub pending_messages: Option<StateResetPolicy>,
102    /// Subnet-global state outside canister snapshots.
103    pub subnet_state: Option<StateResetPolicy>,
104    /// Caller-owned files, processes, services or other external state.
105    pub external_resources: Option<StateResetPolicy>,
106}
107
108/// Typed reset guarantees required by one fixture recipe.
109#[derive(Clone, Copy, Debug, Eq, PartialEq)]
110pub struct ResetRequirements {
111    /// Funding selected once for the pool's canonical snapshot restore.
112    pub funding: SnapshotRestoreFunding,
113    /// Required non-snapshot policies, independently checked against each reset.
114    pub policies: ResetPolicies,
115}
116
117/// Typed reset guarantees actually achieved by one preparation pass.
118///
119/// Default reports no achievements. Recipes must derive these fields from their
120/// reset work, rather than copying the requirements into an assumed receipt.
121#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
122pub struct ResetReceipt {
123    /// Non-snapshot policies achieved by the reset operation.
124    pub policies: ResetPolicies,
125}
126
127/// Contract failure while verifying recipe reset evidence.
128#[non_exhaustive]
129#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
130pub enum BaselinePoolContractError {
131    /// A required reset domain had no matching achievement.
132    #[error("required reset domain {domain:?} was not achieved")]
133    MissingResetDomain {
134        /// Reset domain that did not satisfy the contract.
135        domain: ResetDomainKind,
136    },
137    /// The first mismatched domain, with both complete policy sets for inspection.
138    #[error("reset achievement {achievement:?} does not satisfy {requirement:?} for {domain:?}")]
139    ResetPolicyMismatch {
140        /// Reset domain that did not satisfy the contract.
141        domain: ResetDomainKind,
142        /// Policies required by the recipe.
143        requirement: ResetPolicies,
144        /// Policies reported by the reset.
145        achievement: ResetPolicies,
146    },
147}
148
149/// Whether validation is observing a newly built or restored baseline.
150#[non_exhaustive]
151#[derive(Clone, Debug, Eq, PartialEq)]
152pub enum PreparedBaseline {
153    /// The recipe just built this baseline.
154    Built,
155    /// The pool restored an existing baseline and the recipe reset its other state.
156    Restored {
157        /// Typed non-snapshot reset receipt.
158        reset: ResetReceipt,
159    },
160}
161
162/// Recipe stage associated with a structured preparation failure.
163#[non_exhaustive]
164#[derive(Clone, Copy, Debug, Eq, PartialEq)]
165pub enum BaselinePreparationStage {
166    /// Constructing a new baseline.
167    Build,
168    /// Restoring captured canisters.
169    RestoreCanisters,
170    /// Resetting state outside the snapshots.
171    ResetNonSnapshotState,
172    /// Driving the restored topology to readiness.
173    DriveToReadiness,
174    /// Validating a newly built baseline.
175    ValidateBuilt,
176    /// Validating a restored baseline.
177    ValidateRestored,
178}
179
180/// Why an invalid or failed slot was reconstructed.
181#[non_exhaustive]
182#[derive(Clone, Debug, Eq, PartialEq)]
183pub enum RebuildReason {
184    /// `PocketIC` transport was no longer reachable.
185    DeadPocketIcTransport,
186    /// Captured snapshot restoration failed.
187    SnapshotRestoreFailure,
188    /// Non-snapshot reset failed.
189    ResetFailure,
190    /// Readiness or quiescence could not be established.
191    ReadinessFailure,
192    /// Required and achieved reset domains did not match.
193    ResetCoverageMismatch,
194    /// Final invariant validation failed.
195    InvariantValidationFailure,
196    /// A caller explicitly invalidated its lease.
197    ExplicitLeaseInvalidation,
198    /// A lease was dropped while its thread was unwinding.
199    UnwindWhileLeased,
200    /// Recipe-specific structured reason.
201    RecipeClassified {
202        /// Caller-defined rebuild reason.
203        code: String,
204    },
205}
206
207/// Recipe decision for a failed restored-slot preparation stage.
208#[non_exhaustive]
209#[derive(Clone, Debug, Eq, PartialEq)]
210pub enum FailureDisposition {
211    /// Return the failure without rebuilding during this acquisition.
212    Fatal,
213    /// Invalidate and rebuild the slot once.
214    Rebuild(RebuildReason),
215}
216
217/// Timings for one baseline-pool acquisition.
218#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
219pub struct BaselinePoolTimings {
220    wait: Duration,
221    build: Option<Duration>,
222    restore: Option<Duration>,
223    reset: Option<Duration>,
224    readiness: Option<Duration>,
225    validation: Option<Duration>,
226    stale_teardown: Option<Duration>,
227    total: Duration,
228}
229
230/// One successful acquisition's slot, timings and preparation disposition.
231#[derive(Clone, Debug, Eq, PartialEq)]
232pub struct BaselinePoolOutcome {
233    /// Diagnostic slot index.
234    pub slot: usize,
235    /// Acquisition phase timings.
236    pub timings: BaselinePoolTimings,
237    /// Whether this acquisition built, restored or rebuilt the baseline.
238    pub disposition: BaselinePoolDisposition,
239}
240
241/// Whether a baseline-pool lease was built, restored, or rebuilt.
242#[non_exhaustive]
243#[derive(Clone, Debug, Eq, PartialEq)]
244pub enum BaselinePoolDisposition {
245    /// An empty slot was constructed and validated.
246    Built,
247    /// An existing slot was restored, reset, and validated.
248    Restored,
249    /// An invalid or failed slot was reconstructed and validated.
250    Rebuilt(RebuildReason),
251}
252
253/// One failed recipe or contract stage while preparing a baseline slot.
254#[non_exhaustive]
255#[derive(Debug)]
256pub enum BaselinePoolPreparationError<E> {
257    /// Caller-owned recipe logic returned an error.
258    Recipe {
259        /// Failed lifecycle stage.
260        stage: BaselinePreparationStage,
261        /// Caller-owned structured source error.
262        source: E,
263    },
264    /// Canonical restoration of the captured snapshot set failed.
265    Restore(ControllerSnapshotError),
266    /// Typed non-snapshot reset evidence violated the pool contract.
267    Contract(BaselinePoolContractError),
268}
269
270/// Failure to acquire a validated baseline-pool lease.
271#[non_exhaustive]
272#[derive(Debug)]
273pub enum BaselinePoolError<E> {
274    /// Initial construction or a non-rebuilt preparation failed.
275    Preparation {
276        /// Recipe or contract failure that stopped acquisition.
277        error: BaselinePoolPreparationError<E>,
278        /// Phase timings recorded before acquisition failed.
279        timings: Box<BaselinePoolTimings>,
280    },
281    /// Reused-slot preparation failed and its one rebuild attempt also failed.
282    RecoveryFailed {
283        /// Original restore/reset/readiness/validation failure.
284        original: Box<BaselinePoolPreparationError<E>>,
285        /// Failure while rebuilding or validating the replacement.
286        rebuild: Box<BaselinePoolPreparationError<E>>,
287        /// Combined timings for preparation, stale teardown, and rebuilding.
288        timings: Box<BaselinePoolTimings>,
289    },
290}
291
292/// Caller-owned build, preflight and non-snapshot reset policy for one pooled baseline.
293pub trait PocketIcBaselineRecipe: Send + Sync + 'static {
294    /// Metadata retained beside every baseline owned by this recipe.
295    type Metadata: Send + 'static;
296    /// Structured caller error shared by recipe lifecycle stages.
297    type Error: std::error::Error + Send + Sync + 'static;
298
299    /// Reset guarantees required before an existing slot may be reused.
300    fn reset_requirements(&self) -> &ResetRequirements;
301
302    /// Construct and capture one complete baseline.
303    ///
304    /// # Errors
305    /// Returns the recipe error if construction or baseline capture fails.
306    fn build(&self) -> Result<CachedPocketIcBaseline<Self::Metadata>, Self::Error>;
307
308    /// Admit recipe-specific preconditions before the pool restores snapshots.
309    ///
310    /// Most recipes need no preflight. Checks such as exact current cycle-balance
311    /// guards run here, before restoration can charge cycles. This cannot replace
312    /// the pool's canonical restore or manufacture its completion evidence.
313    ///
314    /// # Errors
315    /// Returns the recipe error when a precondition forbids snapshot restoration.
316    fn prepare_restore(
317        &self,
318        _baseline: &CachedPocketIcBaseline<Self::Metadata>,
319    ) -> Result<(), Self::Error> {
320        Ok(())
321    }
322
323    /// Reset state not covered by canister snapshots.
324    ///
325    /// # Errors
326    /// Returns the recipe error if the required non-snapshot reset fails.
327    fn reset_non_snapshot_state(
328        &self,
329        baseline: &CachedPocketIcBaseline<Self::Metadata>,
330    ) -> Result<ResetReceipt, Self::Error>;
331
332    /// Drive the topology to the recipe's readiness boundary.
333    /// Successful completion permits final invariant validation.
334    ///
335    /// # Errors
336    /// Returns the recipe error if the caller-defined readiness boundary is not reached.
337    fn drive_to_readiness(
338        &self,
339        baseline: &CachedPocketIcBaseline<Self::Metadata>,
340    ) -> Result<(), Self::Error>;
341
342    /// Validate the same baseline invariants after build and restore.
343    /// Success admits this pool-owned recipe's baseline for leasing.
344    ///
345    /// # Errors
346    /// Returns the recipe error when baseline invariants are not satisfied.
347    fn validate(
348        &self,
349        baseline: &CachedPocketIcBaseline<Self::Metadata>,
350        preparation: &PreparedBaseline,
351    ) -> Result<(), Self::Error>;
352
353    /// Classify a restored-slot recipe failure as fatal or rebuildable.
354    fn classify_failure(
355        &self,
356        stage: BaselinePreparationStage,
357        _error: &(dyn std::error::Error + 'static),
358    ) -> FailureDisposition {
359        FailureDisposition::Rebuild(stage.default_rebuild_reason())
360    }
361}
362
363/// Caller-owned runtime-capacity pool of independently restorable `PocketIC` baselines.
364///
365/// One pool structurally owns one [`PocketIcBaselineRecipe`]. A warm
366/// acquisition restores the complete captured canister set, applies the
367/// recipe's non-snapshot reset, reaches its readiness boundary, checks typed
368/// reset coverage, and validates final invariants before exposing a lease.
369/// Each capacity slot owns an independent `PocketIC` instance.
370///
371/// Snapshot reuse is not a complete `PocketIC` rollback. The recipe must account
372/// for time, extra canisters, pending messages, subnet state, cycles, and
373/// external resources when those domains matter to its tests.
374pub struct CachedPocketIcBaselinePool<R>
375where
376    R: PocketIcBaselineRecipe,
377{
378    recipe: R,
379    slots: BoundedSlotPool<CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
380}
381
382/// Exclusive lease of one validated pooled `PocketIC` baseline.
383pub struct CachedPocketIcBaselinePoolGuard<'a, R>
384where
385    R: PocketIcBaselineRecipe,
386{
387    slot: BoundedSlotLease<'a, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
388}
389
390impl ResetRequirements {
391    fn verify(&self, receipt: ResetReceipt) -> Result<(), BaselinePoolContractError> {
392        self.verify_policy(
393            receipt,
394            ResetDomainKind::PocketIcTime,
395            self.policies.pocket_ic_time,
396            receipt.policies.pocket_ic_time,
397        )?;
398        self.verify_policy(
399            receipt,
400            ResetDomainKind::ExtraCanisters,
401            self.policies.extra_canisters,
402            receipt.policies.extra_canisters,
403        )?;
404        self.verify_policy(
405            receipt,
406            ResetDomainKind::PendingMessages,
407            self.policies.pending_messages,
408            receipt.policies.pending_messages,
409        )?;
410        self.verify_policy(
411            receipt,
412            ResetDomainKind::SubnetState,
413            self.policies.subnet_state,
414            receipt.policies.subnet_state,
415        )?;
416        self.verify_policy(
417            receipt,
418            ResetDomainKind::ExternalResources,
419            self.policies.external_resources,
420            receipt.policies.external_resources,
421        )
422    }
423
424    fn verify_policy<P: PartialEq>(
425        &self,
426        receipt: ResetReceipt,
427        domain: ResetDomainKind,
428        requirement: Option<P>,
429        achievement: Option<P>,
430    ) -> Result<(), BaselinePoolContractError> {
431        let Some(requirement) = requirement else {
432            return Ok(());
433        };
434        let Some(achievement) = achievement else {
435            return Err(BaselinePoolContractError::MissingResetDomain { domain });
436        };
437        if requirement != achievement {
438            return Err(BaselinePoolContractError::ResetPolicyMismatch {
439                domain,
440                requirement: self.policies,
441                achievement: receipt.policies,
442            });
443        }
444        Ok(())
445    }
446}
447
448impl BaselinePreparationStage {
449    /// Default rebuild reason used by [`PocketIcBaselineRecipe::classify_failure`].
450    ///
451    /// Recipes that override classification can use this for their fallback
452    /// after handling a more specific error such as dead `PocketIC` transport.
453    #[must_use]
454    pub fn default_rebuild_reason(self) -> RebuildReason {
455        match self {
456            Self::RestoreCanisters => RebuildReason::SnapshotRestoreFailure,
457            Self::ResetNonSnapshotState => RebuildReason::ResetFailure,
458            Self::DriveToReadiness => RebuildReason::ReadinessFailure,
459            Self::ValidateRestored | Self::ValidateBuilt => {
460                RebuildReason::InvariantValidationFailure
461            }
462            Self::Build => RebuildReason::RecipeClassified {
463                code: "build".to_owned(),
464            },
465        }
466    }
467}
468
469impl BaselinePoolTimings {
470    /// Time spent waiting for a capacity slot.
471    #[must_use]
472    pub const fn wait(self) -> Duration {
473        self.wait
474    }
475
476    /// Time spent constructing a new baseline.
477    #[must_use]
478    pub const fn build(self) -> Option<Duration> {
479        self.build
480    }
481
482    /// Time spent restoring captured canisters.
483    #[must_use]
484    pub const fn restore(self) -> Option<Duration> {
485        self.restore
486    }
487
488    /// Time spent resetting non-snapshot state.
489    #[must_use]
490    pub const fn reset(self) -> Option<Duration> {
491        self.reset
492    }
493
494    /// Time spent driving the topology to readiness.
495    #[must_use]
496    pub const fn readiness(self) -> Option<Duration> {
497        self.readiness
498    }
499
500    /// Time spent validating final baseline invariants.
501    #[must_use]
502    pub const fn validation(self) -> Option<Duration> {
503        self.validation
504    }
505
506    /// Time spent dropping an invalid baseline before rebuilding.
507    #[must_use]
508    pub const fn stale_teardown(self) -> Option<Duration> {
509        self.stale_teardown
510    }
511
512    /// Complete acquisition duration.
513    #[must_use]
514    pub const fn total(self) -> Duration {
515        self.total
516    }
517}
518
519impl std::fmt::Display for BaselinePoolTimings {
520    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
521        write!(
522            formatter,
523            "total={:?} wait={:?} build={:?} restore={:?} reset={:?} readiness={:?} validation={:?} stale_teardown={:?}",
524            self.total,
525            self.wait,
526            self.build,
527            self.restore,
528            self.reset,
529            self.readiness,
530            self.validation,
531            self.stale_teardown,
532        )
533    }
534}
535
536impl std::fmt::Display for BaselinePoolOutcome {
537    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
538        match &self.disposition {
539            BaselinePoolDisposition::Built => formatter.write_str("built")?,
540            BaselinePoolDisposition::Restored => formatter.write_str("restored")?,
541            BaselinePoolDisposition::Rebuilt(reason) => {
542                write!(
543                    formatter,
544                    "rebuilt slot={} reason={reason:?} {}",
545                    self.slot, self.timings
546                )?;
547                return Ok(());
548            }
549        }
550        write!(formatter, " slot={} {}", self.slot, self.timings)
551    }
552}
553
554impl<E> BaselinePoolError<E> {
555    /// Timings recorded before this acquisition failed.
556    #[must_use]
557    pub const fn timings(&self) -> BaselinePoolTimings {
558        match self {
559            Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
560        }
561    }
562}
563
564impl<R> CachedPocketIcBaselinePool<R>
565where
566    R: PocketIcBaselineRecipe,
567{
568    /// Create a runtime-capacity pool that structurally owns one recipe.
569    ///
570    /// Construction is const and lazy: slots are allocated on first acquisition,
571    /// and the recipe builds a baseline only when a slot needs one.
572    #[must_use]
573    pub const fn new(capacity: NonZeroUsize, recipe: R) -> Self {
574        Self {
575            recipe,
576            slots: BoundedSlotPool::new(capacity),
577        }
578    }
579
580    /// Maximum number of simultaneously leased `PocketIC` baselines.
581    #[must_use]
582    pub const fn capacity(&self) -> NonZeroUsize {
583        self.slots.capacity()
584    }
585
586    /// Acquire one fully built or restored and validated baseline lease.
587    ///
588    /// A rebuildable warm-preparation failure discards the stale slot and
589    /// performs at most one build attempt. Recipe and caller panics are never
590    /// converted into cache misses; the lease is invalidated while the panic
591    /// continues unwinding.
592    ///
593    /// # Errors
594    ///
595    /// Returns a stage-specific recipe or contract error. If warm preparation
596    /// and its one recovery build both fail, the error preserves both causes.
597    #[expect(
598        clippy::significant_drop_tightening,
599        reason = "lease and coordinator scope deliberately cover the complete operation"
600    )]
601    pub fn acquire(
602        &self,
603    ) -> Result<
604        (CachedPocketIcBaselinePoolGuard<'_, R>, BaselinePoolOutcome),
605        BaselinePoolError<R::Error>,
606    > {
607        let total_started = Instant::now();
608        let mut slot = self.slots.acquire();
609        let mut timings = BaselinePoolTimings {
610            wait: slot.wait(),
611            ..BaselinePoolTimings::default()
612        };
613        let slot_index = slot.slot_index();
614
615        let (rebuild_reason, original) = match slot.state() {
616            Slot::Reusable(_) => match self.prepare_reused(&slot, &mut timings) {
617                Ok(()) => {
618                    timings.total = total_started.elapsed();
619                    return Ok((
620                        CachedPocketIcBaselinePoolGuard { slot },
621                        BaselinePoolOutcome {
622                            slot: slot_index,
623                            timings,
624                            disposition: BaselinePoolDisposition::Restored,
625                        },
626                    ));
627                }
628                Err(original) => match self.failure_disposition(&original) {
629                    FailureDisposition::Fatal => {
630                        // A failed restore may have partially changed the instance.
631                        // Discard it without treating a fatal failure as a rebuild.
632                        Self::discard_stale_slot(&mut slot, &mut timings);
633                        timings.total = total_started.elapsed();
634                        return Err(BaselinePoolError::Preparation {
635                            error: original,
636                            timings: Box::new(timings),
637                        });
638                    }
639                    FailureDisposition::Rebuild(reason) => (Some(reason), Some(original)),
640                },
641            },
642            Slot::Empty => (None, None),
643            Slot::Invalidated { reason, .. } => (Some(reason.clone()), None),
644            Slot::Unwound(_) => (Some(RebuildReason::UnwindWhileLeased), None),
645        };
646
647        if slot.get().is_some() {
648            Self::discard_stale_slot(&mut slot, &mut timings);
649        }
650        let build = self.build_slot(&mut slot, &mut timings);
651        timings.total = total_started.elapsed();
652        if let Err(error) = build {
653            return Err(match original {
654                Some(original) => BaselinePoolError::RecoveryFailed {
655                    original: Box::new(original),
656                    rebuild: Box::new(error),
657                    timings: Box::new(timings),
658                },
659                None => BaselinePoolError::Preparation {
660                    error,
661                    timings: Box::new(timings),
662                },
663            });
664        }
665
666        let outcome = BaselinePoolOutcome {
667            slot: slot_index,
668            timings,
669            disposition: rebuild_reason.map_or(
670                BaselinePoolDisposition::Built,
671                BaselinePoolDisposition::Rebuilt,
672            ),
673        };
674        Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
675    }
676
677    fn prepare_reused(
678        &self,
679        slot: &BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
680        timings: &mut BaselinePoolTimings,
681    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
682        let baseline = slot
683            .get()
684            .expect("reusable baseline slot must be populated");
685
686        let started = Instant::now();
687        let restore = self
688            .recipe
689            .prepare_restore(baseline)
690            .map_err(|source| BaselinePoolPreparationError::Recipe {
691                stage: BaselinePreparationStage::RestoreCanisters,
692                source,
693            })
694            .and_then(|()| {
695                baseline
696                    .restore(self.recipe.reset_requirements().funding)
697                    .map_err(BaselinePoolPreparationError::Restore)
698            });
699        add_timing(&mut timings.restore, started.elapsed());
700        restore?;
701
702        let reset = timed_recipe_stage(
703            &mut timings.reset,
704            BaselinePreparationStage::ResetNonSnapshotState,
705            || self.recipe.reset_non_snapshot_state(baseline),
706        )?;
707        timed_recipe_stage(
708            &mut timings.readiness,
709            BaselinePreparationStage::DriveToReadiness,
710            || self.recipe.drive_to_readiness(baseline),
711        )?;
712        self.recipe
713            .reset_requirements()
714            .verify(reset)
715            .map_err(BaselinePoolPreparationError::Contract)?;
716
717        let preparation = PreparedBaseline::Restored { reset };
718        self.validate_baseline(baseline, &preparation, timings)?;
719        Ok(())
720    }
721
722    fn build_slot(
723        &self,
724        slot: &mut BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
725        timings: &mut BaselinePoolTimings,
726    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
727        let baseline =
728            timed_recipe_stage(&mut timings.build, BaselinePreparationStage::Build, || {
729                self.recipe.build()
730            })?;
731
732        if let Err(error) = self.validate_baseline(&baseline, &PreparedBaseline::Built, timings) {
733            drop_baseline_safely(baseline);
734            return Err(error);
735        }
736        let replaced = slot.replace(baseline);
737        debug_assert!(replaced.is_none());
738        Ok(())
739    }
740
741    fn validate_baseline(
742        &self,
743        baseline: &CachedPocketIcBaseline<R::Metadata>,
744        preparation: &PreparedBaseline,
745        timings: &mut BaselinePoolTimings,
746    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
747        let stage = match preparation {
748            PreparedBaseline::Built => BaselinePreparationStage::ValidateBuilt,
749            PreparedBaseline::Restored { .. } => BaselinePreparationStage::ValidateRestored,
750        };
751        timed_recipe_stage(&mut timings.validation, stage, || {
752            self.recipe.validate(baseline, preparation)
753        })
754    }
755
756    fn failure_disposition(
757        &self,
758        error: &BaselinePoolPreparationError<R::Error>,
759    ) -> FailureDisposition {
760        match error {
761            BaselinePoolPreparationError::Recipe { stage, source } => {
762                self.recipe.classify_failure(*stage, source)
763            }
764            BaselinePoolPreparationError::Restore(source) => self
765                .recipe
766                .classify_failure(BaselinePreparationStage::RestoreCanisters, source),
767            BaselinePoolPreparationError::Contract(_) => {
768                FailureDisposition::Rebuild(RebuildReason::ResetCoverageMismatch)
769            }
770        }
771    }
772
773    fn discard_stale_slot(
774        slot: &mut BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
775        timings: &mut BaselinePoolTimings,
776    ) {
777        let started = Instant::now();
778        if let Some(stale) = slot.take() {
779            drop_baseline_safely(stale);
780        }
781        timings.stale_teardown = Some(started.elapsed());
782    }
783}
784
785impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
786where
787    R: PocketIcBaselineRecipe,
788{
789    /// Diagnostic slot index held by this lease.
790    #[must_use]
791    pub const fn slot(&self) -> usize {
792        self.slot.slot_index()
793    }
794
795    /// Mark this slot non-reusable with a structured rebuild reason.
796    pub fn invalidate(&mut self, reason: RebuildReason) {
797        self.slot.invalidate(reason);
798    }
799}
800
801impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
802where
803    R: PocketIcBaselineRecipe,
804{
805    type Target = CachedPocketIcBaseline<R::Metadata>;
806
807    fn deref(&self) -> &Self::Target {
808        self.slot
809            .get()
810            .expect("leased baseline pool slot must be populated")
811    }
812}
813
814fn timed_recipe_stage<T, E>(
815    timing: &mut Option<Duration>,
816    stage: BaselinePreparationStage,
817    operation: impl FnOnce() -> Result<T, E>,
818) -> Result<T, BaselinePoolPreparationError<E>> {
819    let started = Instant::now();
820    let result = operation();
821    add_timing(timing, started.elapsed());
822    result.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })
823}
824
825const fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
826    *total = Some(match *total {
827        None => elapsed,
828        Some(previous) => previous.saturating_add(elapsed),
829    });
830}
831
832fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
833    let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
834}
835
836impl std::fmt::Display for BaselinePreparationStage {
837    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
838        formatter.write_str(match self {
839            Self::Build => "build",
840            Self::RestoreCanisters => "canister restore",
841            Self::ResetNonSnapshotState => "non-snapshot reset",
842            Self::DriveToReadiness => "readiness",
843            Self::ValidateBuilt => "built-baseline validation",
844            Self::ValidateRestored => "restored-baseline validation",
845        })
846    }
847}
848
849impl<E> std::fmt::Display for BaselinePoolPreparationError<E>
850where
851    E: std::fmt::Display,
852{
853    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
854        match self {
855            Self::Recipe { stage, source } => {
856                write!(formatter, "baseline {stage} failed: {source}")
857            }
858            Self::Restore(source) => {
859                write!(formatter, "baseline canister restore failed: {source}")
860            }
861            Self::Contract(error) => write!(formatter, "baseline pool contract failed: {error}"),
862        }
863    }
864}
865
866impl<E> std::error::Error for BaselinePoolPreparationError<E>
867where
868    E: std::error::Error + 'static,
869{
870    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
871        match self {
872            Self::Recipe { source, .. } => Some(source),
873            Self::Restore(source) => Some(source),
874            Self::Contract(error) => Some(error),
875        }
876    }
877}
878
879impl<E> std::fmt::Display for BaselinePoolError<E>
880where
881    E: std::fmt::Display,
882{
883    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
884        match self {
885            Self::Preparation { error, .. } => error.fmt(formatter),
886            Self::RecoveryFailed {
887                original, rebuild, ..
888            } => write!(
889                formatter,
890                "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
891            ),
892        }
893    }
894}
895
896impl<E> std::error::Error for BaselinePoolError<E>
897where
898    E: std::error::Error + 'static,
899{
900    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
901        match self {
902            Self::Preparation { error, .. } => Some(error),
903            Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
904        }
905    }
906}
907
908#[cfg(test)]
909mod tests {
910    use super::{
911        BaselinePoolContractError, BaselinePoolTimings, ExtraCanisterPolicy, ResetDomainKind,
912        ResetPolicies, ResetReceipt, ResetRequirements, SnapshotRestoreFunding, StateResetPolicy,
913        TimeResetPolicy, add_timing,
914    };
915    use std::time::Duration;
916
917    #[test]
918    fn phase_timing_preserves_absence_zero_and_saturation() {
919        let mut timings = BaselinePoolTimings::default();
920        assert_eq!(timings.restore(), None);
921        add_timing(&mut timings.restore, Duration::ZERO);
922        assert_eq!(timings.restore(), Some(Duration::ZERO));
923        add_timing(&mut timings.restore, Duration::from_secs(2));
924        assert_eq!(timings.restore(), Some(Duration::from_secs(2)));
925        add_timing(&mut timings.restore, Duration::MAX);
926        assert_eq!(timings.restore(), Some(Duration::MAX));
927        assert_eq!(timings.build(), None);
928    }
929
930    #[test]
931    fn required_policy_must_match_achieved_policy() {
932        for (domain, requirement, achievement) in [
933            (
934                ResetDomainKind::PocketIcTime,
935                ResetPolicies {
936                    pocket_ic_time: Some(TimeResetPolicy::PreserveCurrent),
937                    ..ResetPolicies::default()
938                },
939                ResetPolicies {
940                    pocket_ic_time: Some(TimeResetPolicy::RebuildOnMutation),
941                    ..ResetPolicies::default()
942                },
943            ),
944            (
945                ResetDomainKind::ExtraCanisters,
946                ResetPolicies {
947                    extra_canisters: Some(ExtraCanisterPolicy::RequireBaselineSet),
948                    ..ResetPolicies::default()
949                },
950                ResetPolicies {
951                    extra_canisters: Some(ExtraCanisterPolicy::RemoveTracked),
952                    ..ResetPolicies::default()
953                },
954            ),
955            (
956                ResetDomainKind::PendingMessages,
957                ResetPolicies {
958                    pending_messages: Some(StateResetPolicy::ValidateUnchanged),
959                    ..ResetPolicies::default()
960                },
961                ResetPolicies {
962                    pending_messages: Some(StateResetPolicy::IrrelevantByRecipeContract),
963                    ..ResetPolicies::default()
964                },
965            ),
966            (
967                ResetDomainKind::SubnetState,
968                ResetPolicies {
969                    subnet_state: Some(StateResetPolicy::ResetByRecipe),
970                    ..ResetPolicies::default()
971                },
972                ResetPolicies {
973                    subnet_state: Some(StateResetPolicy::RebuildOnChange),
974                    ..ResetPolicies::default()
975                },
976            ),
977            (
978                ResetDomainKind::ExternalResources,
979                ResetPolicies {
980                    external_resources: Some(StateResetPolicy::ValidateUnchanged),
981                    ..ResetPolicies::default()
982                },
983                ResetPolicies {
984                    external_resources: Some(StateResetPolicy::ResetByRecipe),
985                    ..ResetPolicies::default()
986                },
987            ),
988        ] {
989            let requirements = ResetRequirements {
990                funding: SnapshotRestoreFunding::Preserve,
991                policies: requirement,
992            };
993            assert_eq!(
994                requirements.verify(ResetReceipt {
995                    policies: requirement
996                }),
997                Ok(())
998            );
999            assert_eq!(
1000                requirements.verify(ResetReceipt::default()),
1001                Err(BaselinePoolContractError::MissingResetDomain { domain })
1002            );
1003            assert_eq!(
1004                requirements.verify(ResetReceipt {
1005                    policies: achievement
1006                }),
1007                Err(BaselinePoolContractError::ResetPolicyMismatch {
1008                    domain,
1009                    requirement,
1010                    achievement
1011                })
1012            );
1013            let unconstrained = ResetRequirements {
1014                funding: SnapshotRestoreFunding::Preserve,
1015                policies: ResetPolicies::default(),
1016            };
1017            assert_eq!(
1018                unconstrained.verify(ResetReceipt {
1019                    policies: achievement
1020                }),
1021                Ok(())
1022            );
1023        }
1024    }
1025
1026    #[test]
1027    fn reset_coverage_preserves_absence_and_first_failure() {
1028        let requirements = ResetRequirements {
1029            funding: SnapshotRestoreFunding::Preserve,
1030            policies: ResetPolicies {
1031                pocket_ic_time: Some(TimeResetPolicy::PreserveCurrent),
1032                pending_messages: Some(StateResetPolicy::IrrelevantByRecipeContract),
1033                ..ResetPolicies::default()
1034            },
1035        };
1036        let mut receipt = ResetReceipt::default();
1037        assert_eq!(
1038            requirements.verify(receipt),
1039            Err(BaselinePoolContractError::MissingResetDomain {
1040                domain: ResetDomainKind::PocketIcTime,
1041            })
1042        );
1043        receipt.policies.pocket_ic_time = Some(TimeResetPolicy::RebuildOnMutation);
1044        assert_eq!(
1045            requirements.verify(receipt),
1046            Err(BaselinePoolContractError::ResetPolicyMismatch {
1047                domain: ResetDomainKind::PocketIcTime,
1048                requirement: requirements.policies,
1049                achievement: receipt.policies,
1050            })
1051        );
1052        receipt.policies.pocket_ic_time = Some(TimeResetPolicy::PreserveCurrent);
1053        assert_eq!(
1054            requirements.verify(receipt),
1055            Err(BaselinePoolContractError::MissingResetDomain {
1056                domain: ResetDomainKind::PendingMessages,
1057            })
1058        );
1059        receipt.policies.pending_messages = Some(StateResetPolicy::IrrelevantByRecipeContract);
1060        assert_eq!(requirements.verify(receipt), Ok(()));
1061    }
1062}
ic_testkit/pic/baseline_pool.rs

ic_testkit/pic/
baseline_pool.rs
