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, thiserror::Error)]
256pub enum BaselinePoolPreparationError<E> {
257    /// Caller-owned recipe logic returned an error.
258    #[error("baseline {stage} failed: {source}")]
259    Recipe {
260        /// Failed lifecycle stage.
261        stage: BaselinePreparationStage,
262        /// Caller-owned structured source error.
263        source: E,
264    },
265    /// Canonical restoration of the captured snapshot set failed.
266    #[error("baseline canister restore failed: {0}")]
267    Restore(#[source] ControllerSnapshotError),
268    /// Typed non-snapshot reset evidence violated the pool contract.
269    #[error("baseline pool contract failed: {0}")]
270    Contract(#[source] BaselinePoolContractError),
271}
272
273/// Failure to acquire a validated baseline-pool lease.
274#[non_exhaustive]
275#[derive(Debug)]
276pub enum BaselinePoolError<E> {
277    /// Initial construction or a non-rebuilt preparation failed.
278    Preparation {
279        /// Recipe or contract failure that stopped acquisition.
280        error: BaselinePoolPreparationError<E>,
281        /// Phase timings recorded before acquisition failed.
282        timings: Box<BaselinePoolTimings>,
283    },
284    /// Reused-slot preparation failed and its one rebuild attempt also failed.
285    RecoveryFailed {
286        /// Original restore/reset/readiness/validation failure.
287        original: Box<BaselinePoolPreparationError<E>>,
288        /// Failure while rebuilding or validating the replacement.
289        rebuild: Box<BaselinePoolPreparationError<E>>,
290        /// Combined timings for preparation, stale teardown, and rebuilding.
291        timings: Box<BaselinePoolTimings>,
292    },
293}
294
295/// Caller-owned build, preflight and non-snapshot reset policy for one pooled baseline.
296pub trait PocketIcBaselineRecipe: Send + Sync + 'static {
297    /// Metadata retained beside every baseline owned by this recipe.
298    type Metadata: Send + 'static;
299    /// Structured caller error shared by recipe lifecycle stages.
300    type Error: std::error::Error + Send + Sync + 'static;
301
302    /// Reset guarantees required before an existing slot may be reused.
303    fn reset_requirements(&self) -> &ResetRequirements;
304
305    /// Construct and capture one complete baseline.
306    ///
307    /// # Errors
308    /// Returns the recipe error if construction or baseline capture fails.
309    fn build(&self) -> Result<CachedPocketIcBaseline<Self::Metadata>, Self::Error>;
310
311    /// Admit recipe-specific preconditions before the pool restores snapshots.
312    ///
313    /// Most recipes need no preflight. Checks such as exact current cycle-balance
314    /// guards run here, before restoration can charge cycles. This cannot replace
315    /// the pool's canonical restore or manufacture its completion evidence.
316    ///
317    /// # Errors
318    /// Returns the recipe error when a precondition forbids snapshot restoration.
319    fn prepare_restore(
320        &self,
321        _baseline: &CachedPocketIcBaseline<Self::Metadata>,
322    ) -> Result<(), Self::Error> {
323        Ok(())
324    }
325
326    /// Reset state not covered by canister snapshots.
327    ///
328    /// # Errors
329    /// Returns the recipe error if the required non-snapshot reset fails.
330    fn reset_non_snapshot_state(
331        &self,
332        baseline: &CachedPocketIcBaseline<Self::Metadata>,
333    ) -> Result<ResetReceipt, Self::Error>;
334
335    /// Drive the topology to the recipe's readiness boundary.
336    /// Successful completion permits final invariant validation.
337    ///
338    /// # Errors
339    /// Returns the recipe error if the caller-defined readiness boundary is not reached.
340    fn drive_to_readiness(
341        &self,
342        baseline: &CachedPocketIcBaseline<Self::Metadata>,
343    ) -> Result<(), Self::Error>;
344
345    /// Validate the same baseline invariants after build and restore.
346    /// Success admits this pool-owned recipe's baseline for leasing.
347    ///
348    /// # Errors
349    /// Returns the recipe error when baseline invariants are not satisfied.
350    fn validate(
351        &self,
352        baseline: &CachedPocketIcBaseline<Self::Metadata>,
353        preparation: &PreparedBaseline,
354    ) -> Result<(), Self::Error>;
355
356    /// Classify a restored-slot recipe failure as fatal or rebuildable.
357    fn classify_failure(
358        &self,
359        stage: BaselinePreparationStage,
360        _error: &(dyn std::error::Error + 'static),
361    ) -> FailureDisposition {
362        FailureDisposition::Rebuild(stage.default_rebuild_reason())
363    }
364}
365
366/// Caller-owned runtime-capacity pool of independently restorable `PocketIC` baselines.
367///
368/// One pool structurally owns one [`PocketIcBaselineRecipe`]. A warm
369/// acquisition restores the complete captured canister set, applies the
370/// recipe's non-snapshot reset, reaches its readiness boundary, checks typed
371/// reset coverage, and validates final invariants before exposing a lease.
372/// Each capacity slot owns an independent `PocketIC` instance.
373///
374/// Snapshot reuse is not a complete `PocketIC` rollback. The recipe must account
375/// for time, extra canisters, pending messages, subnet state, cycles, and
376/// external resources when those domains matter to its tests.
377pub struct CachedPocketIcBaselinePool<R>
378where
379    R: PocketIcBaselineRecipe,
380{
381    recipe: R,
382    slots: BoundedSlotPool<CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
383}
384
385/// Exclusive lease of one validated pooled `PocketIC` baseline.
386pub struct CachedPocketIcBaselinePoolGuard<'a, R>
387where
388    R: PocketIcBaselineRecipe,
389{
390    slot: BoundedSlotLease<'a, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
391}
392
393impl ResetRequirements {
394    fn verify(&self, receipt: ResetReceipt) -> Result<(), BaselinePoolContractError> {
395        self.verify_policy(
396            receipt,
397            ResetDomainKind::PocketIcTime,
398            self.policies.pocket_ic_time,
399            receipt.policies.pocket_ic_time,
400        )?;
401        self.verify_policy(
402            receipt,
403            ResetDomainKind::ExtraCanisters,
404            self.policies.extra_canisters,
405            receipt.policies.extra_canisters,
406        )?;
407        self.verify_policy(
408            receipt,
409            ResetDomainKind::PendingMessages,
410            self.policies.pending_messages,
411            receipt.policies.pending_messages,
412        )?;
413        self.verify_policy(
414            receipt,
415            ResetDomainKind::SubnetState,
416            self.policies.subnet_state,
417            receipt.policies.subnet_state,
418        )?;
419        self.verify_policy(
420            receipt,
421            ResetDomainKind::ExternalResources,
422            self.policies.external_resources,
423            receipt.policies.external_resources,
424        )
425    }
426
427    fn verify_policy<P: PartialEq>(
428        &self,
429        receipt: ResetReceipt,
430        domain: ResetDomainKind,
431        requirement: Option<P>,
432        achievement: Option<P>,
433    ) -> Result<(), BaselinePoolContractError> {
434        let Some(requirement) = requirement else {
435            return Ok(());
436        };
437        let Some(achievement) = achievement else {
438            return Err(BaselinePoolContractError::MissingResetDomain { domain });
439        };
440        if requirement != achievement {
441            return Err(BaselinePoolContractError::ResetPolicyMismatch {
442                domain,
443                requirement: self.policies,
444                achievement: receipt.policies,
445            });
446        }
447        Ok(())
448    }
449}
450
451impl BaselinePreparationStage {
452    /// Default rebuild reason used by [`PocketIcBaselineRecipe::classify_failure`].
453    ///
454    /// Recipes that override classification can use this for their fallback
455    /// after handling a more specific error such as dead `PocketIC` transport.
456    #[must_use]
457    pub fn default_rebuild_reason(self) -> RebuildReason {
458        match self {
459            Self::RestoreCanisters => RebuildReason::SnapshotRestoreFailure,
460            Self::ResetNonSnapshotState => RebuildReason::ResetFailure,
461            Self::DriveToReadiness => RebuildReason::ReadinessFailure,
462            Self::ValidateRestored | Self::ValidateBuilt => {
463                RebuildReason::InvariantValidationFailure
464            }
465            Self::Build => RebuildReason::RecipeClassified {
466                code: "build".to_owned(),
467            },
468        }
469    }
470}
471
472impl BaselinePoolTimings {
473    /// Time spent waiting for a capacity slot.
474    #[must_use]
475    pub const fn wait(self) -> Duration {
476        self.wait
477    }
478
479    /// Time spent constructing a new baseline.
480    #[must_use]
481    pub const fn build(self) -> Option<Duration> {
482        self.build
483    }
484
485    /// Time spent restoring captured canisters.
486    #[must_use]
487    pub const fn restore(self) -> Option<Duration> {
488        self.restore
489    }
490
491    /// Time spent resetting non-snapshot state.
492    #[must_use]
493    pub const fn reset(self) -> Option<Duration> {
494        self.reset
495    }
496
497    /// Time spent driving the topology to readiness.
498    #[must_use]
499    pub const fn readiness(self) -> Option<Duration> {
500        self.readiness
501    }
502
503    /// Time spent validating final baseline invariants.
504    #[must_use]
505    pub const fn validation(self) -> Option<Duration> {
506        self.validation
507    }
508
509    /// Time spent dropping an invalid baseline before rebuilding.
510    #[must_use]
511    pub const fn stale_teardown(self) -> Option<Duration> {
512        self.stale_teardown
513    }
514
515    /// Complete acquisition duration.
516    #[must_use]
517    pub const fn total(self) -> Duration {
518        self.total
519    }
520}
521
522impl std::fmt::Display for BaselinePoolTimings {
523    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
524        write!(
525            formatter,
526            "total={:?} wait={:?} build={:?} restore={:?} reset={:?} readiness={:?} validation={:?} stale_teardown={:?}",
527            self.total,
528            self.wait,
529            self.build,
530            self.restore,
531            self.reset,
532            self.readiness,
533            self.validation,
534            self.stale_teardown,
535        )
536    }
537}
538
539impl std::fmt::Display for BaselinePoolOutcome {
540    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
541        match &self.disposition {
542            BaselinePoolDisposition::Built => formatter.write_str("built")?,
543            BaselinePoolDisposition::Restored => formatter.write_str("restored")?,
544            BaselinePoolDisposition::Rebuilt(reason) => {
545                write!(
546                    formatter,
547                    "rebuilt slot={} reason={reason:?} {}",
548                    self.slot, self.timings
549                )?;
550                return Ok(());
551            }
552        }
553        write!(formatter, " slot={} {}", self.slot, self.timings)
554    }
555}
556
557impl<E> BaselinePoolError<E> {
558    /// Timings recorded before this acquisition failed.
559    #[must_use]
560    pub const fn timings(&self) -> BaselinePoolTimings {
561        match self {
562            Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
563        }
564    }
565}
566
567impl<R> CachedPocketIcBaselinePool<R>
568where
569    R: PocketIcBaselineRecipe,
570{
571    /// Create a runtime-capacity pool that structurally owns one recipe.
572    ///
573    /// Construction is const and lazy: slots are allocated on first acquisition,
574    /// and the recipe builds a baseline only when a slot needs one.
575    #[must_use]
576    pub const fn new(capacity: NonZeroUsize, recipe: R) -> Self {
577        Self {
578            recipe,
579            slots: BoundedSlotPool::new(capacity),
580        }
581    }
582
583    /// Maximum number of simultaneously leased `PocketIC` baselines.
584    #[must_use]
585    pub const fn capacity(&self) -> NonZeroUsize {
586        self.slots.capacity()
587    }
588
589    /// Acquire one fully built or restored and validated baseline lease.
590    ///
591    /// A rebuildable warm-preparation failure discards the stale slot and
592    /// performs at most one build attempt. Recipe and caller panics are never
593    /// converted into cache misses; the lease is invalidated while the panic
594    /// continues unwinding.
595    ///
596    /// # Errors
597    ///
598    /// Returns a stage-specific recipe or contract error. If warm preparation
599    /// and its one recovery build both fail, the error preserves both causes.
600    #[expect(
601        clippy::significant_drop_tightening,
602        reason = "lease and coordinator scope deliberately cover the complete operation"
603    )]
604    pub fn acquire(
605        &self,
606    ) -> Result<
607        (CachedPocketIcBaselinePoolGuard<'_, R>, BaselinePoolOutcome),
608        BaselinePoolError<R::Error>,
609    > {
610        let total_started = Instant::now();
611        let mut slot = self.slots.acquire();
612        let mut timings = BaselinePoolTimings {
613            wait: slot.wait(),
614            ..BaselinePoolTimings::default()
615        };
616        let slot_index = slot.slot_index();
617
618        let (disposition, original) = match slot.state() {
619            Slot::Reusable(_) => match self.prepare_reused(&slot, &mut timings) {
620                Ok(()) => {
621                    timings.total = total_started.elapsed();
622                    return Ok((
623                        CachedPocketIcBaselinePoolGuard { slot },
624                        BaselinePoolOutcome {
625                            slot: slot_index,
626                            timings,
627                            disposition: BaselinePoolDisposition::Restored,
628                        },
629                    ));
630                }
631                Err(original) => match self.failure_disposition(&original) {
632                    FailureDisposition::Fatal => {
633                        // A failed restore may have partially changed the instance.
634                        // Discard it without treating a fatal failure as a rebuild.
635                        Self::discard_stale_slot(&mut slot, &mut timings);
636                        timings.total = total_started.elapsed();
637                        return Err(BaselinePoolError::Preparation {
638                            error: original,
639                            timings: Box::new(timings),
640                        });
641                    }
642                    FailureDisposition::Rebuild(reason) => {
643                        (BaselinePoolDisposition::Rebuilt(reason), Some(original))
644                    }
645                },
646            },
647            Slot::Empty => (BaselinePoolDisposition::Built, None),
648            Slot::Invalidated { reason, .. } => {
649                (BaselinePoolDisposition::Rebuilt(reason.clone()), None)
650            }
651            Slot::Unwound(_) => (
652                BaselinePoolDisposition::Rebuilt(RebuildReason::UnwindWhileLeased),
653                None,
654            ),
655        };
656
657        if slot.get().is_some() {
658            Self::discard_stale_slot(&mut slot, &mut timings);
659        }
660        let build = self.build_slot(&mut slot, &mut timings);
661        timings.total = total_started.elapsed();
662        if let Err(error) = build {
663            return Err(match original {
664                Some(original) => BaselinePoolError::RecoveryFailed {
665                    original: Box::new(original),
666                    rebuild: Box::new(error),
667                    timings: Box::new(timings),
668                },
669                None => BaselinePoolError::Preparation {
670                    error,
671                    timings: Box::new(timings),
672                },
673            });
674        }
675
676        let outcome = BaselinePoolOutcome {
677            slot: slot_index,
678            timings,
679            disposition,
680        };
681        Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
682    }
683
684    fn prepare_reused(
685        &self,
686        slot: &BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
687        timings: &mut BaselinePoolTimings,
688    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
689        let baseline = slot
690            .get()
691            .expect("reusable baseline slot must be populated");
692
693        let started = Instant::now();
694        let restore = self
695            .recipe
696            .prepare_restore(baseline)
697            .map_err(|source| BaselinePoolPreparationError::Recipe {
698                stage: BaselinePreparationStage::RestoreCanisters,
699                source,
700            })
701            .and_then(|()| {
702                baseline
703                    .restore(self.recipe.reset_requirements().funding)
704                    .map_err(BaselinePoolPreparationError::Restore)
705            });
706        add_timing(&mut timings.restore, started.elapsed());
707        restore?;
708
709        let reset = timed_recipe_stage(
710            &mut timings.reset,
711            BaselinePreparationStage::ResetNonSnapshotState,
712            || self.recipe.reset_non_snapshot_state(baseline),
713        )?;
714        timed_recipe_stage(
715            &mut timings.readiness,
716            BaselinePreparationStage::DriveToReadiness,
717            || self.recipe.drive_to_readiness(baseline),
718        )?;
719        self.recipe
720            .reset_requirements()
721            .verify(reset)
722            .map_err(BaselinePoolPreparationError::Contract)?;
723
724        let preparation = PreparedBaseline::Restored { reset };
725        self.validate_baseline(baseline, &preparation, timings)?;
726        Ok(())
727    }
728
729    fn build_slot(
730        &self,
731        slot: &mut BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
732        timings: &mut BaselinePoolTimings,
733    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
734        let baseline =
735            timed_recipe_stage(&mut timings.build, BaselinePreparationStage::Build, || {
736                self.recipe.build()
737            })?;
738
739        if let Err(error) = self.validate_baseline(&baseline, &PreparedBaseline::Built, timings) {
740            drop_baseline_safely(baseline);
741            return Err(error);
742        }
743        let replaced = slot.replace(baseline);
744        debug_assert!(replaced.is_none());
745        Ok(())
746    }
747
748    fn validate_baseline(
749        &self,
750        baseline: &CachedPocketIcBaseline<R::Metadata>,
751        preparation: &PreparedBaseline,
752        timings: &mut BaselinePoolTimings,
753    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
754        let stage = match preparation {
755            PreparedBaseline::Built => BaselinePreparationStage::ValidateBuilt,
756            PreparedBaseline::Restored { .. } => BaselinePreparationStage::ValidateRestored,
757        };
758        timed_recipe_stage(&mut timings.validation, stage, || {
759            self.recipe.validate(baseline, preparation)
760        })
761    }
762
763    fn failure_disposition(
764        &self,
765        error: &BaselinePoolPreparationError<R::Error>,
766    ) -> FailureDisposition {
767        match error {
768            BaselinePoolPreparationError::Recipe { stage, source } => {
769                self.recipe.classify_failure(*stage, source)
770            }
771            BaselinePoolPreparationError::Restore(source) => self
772                .recipe
773                .classify_failure(BaselinePreparationStage::RestoreCanisters, source),
774            BaselinePoolPreparationError::Contract(_) => {
775                FailureDisposition::Rebuild(RebuildReason::ResetCoverageMismatch)
776            }
777        }
778    }
779
780    fn discard_stale_slot(
781        slot: &mut BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
782        timings: &mut BaselinePoolTimings,
783    ) {
784        let started = Instant::now();
785        if let Some(stale) = slot.take() {
786            drop_baseline_safely(stale);
787        }
788        timings.stale_teardown = Some(started.elapsed());
789    }
790}
791
792impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
793where
794    R: PocketIcBaselineRecipe,
795{
796    /// Diagnostic slot index held by this lease.
797    #[must_use]
798    pub const fn slot(&self) -> usize {
799        self.slot.slot_index()
800    }
801
802    /// Mark this slot non-reusable with a structured rebuild reason.
803    pub fn invalidate(&mut self, reason: RebuildReason) {
804        self.slot.invalidate(reason);
805    }
806}
807
808impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
809where
810    R: PocketIcBaselineRecipe,
811{
812    type Target = CachedPocketIcBaseline<R::Metadata>;
813
814    fn deref(&self) -> &Self::Target {
815        self.slot
816            .get()
817            .expect("leased baseline pool slot must be populated")
818    }
819}
820
821fn timed_recipe_stage<T, E>(
822    timing: &mut Option<Duration>,
823    stage: BaselinePreparationStage,
824    operation: impl FnOnce() -> Result<T, E>,
825) -> Result<T, BaselinePoolPreparationError<E>> {
826    let started = Instant::now();
827    let result = operation();
828    add_timing(timing, started.elapsed());
829    result.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })
830}
831
832const fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
833    *total = Some(match *total {
834        None => elapsed,
835        Some(previous) => previous.saturating_add(elapsed),
836    });
837}
838
839fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
840    let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
841}
842
843impl std::fmt::Display for BaselinePreparationStage {
844    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
845        formatter.write_str(match self {
846            Self::Build => "build",
847            Self::RestoreCanisters => "canister restore",
848            Self::ResetNonSnapshotState => "non-snapshot reset",
849            Self::DriveToReadiness => "readiness",
850            Self::ValidateBuilt => "built-baseline validation",
851            Self::ValidateRestored => "restored-baseline validation",
852        })
853    }
854}
855
856impl<E> std::fmt::Display for BaselinePoolError<E>
857where
858    E: std::fmt::Display,
859{
860    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
861        match self {
862            Self::Preparation { error, .. } => error.fmt(formatter),
863            Self::RecoveryFailed {
864                original, rebuild, ..
865            } => write!(
866                formatter,
867                "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
868            ),
869        }
870    }
871}
872
873impl<E> std::error::Error for BaselinePoolError<E>
874where
875    E: std::error::Error + 'static,
876{
877    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
878        match self {
879            Self::Preparation { error, .. } => Some(error),
880            Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
881        }
882    }
883}
884
885#[cfg(test)]
886mod tests {
887    use super::{
888        BaselinePoolContractError, BaselinePoolTimings, ExtraCanisterPolicy, ResetDomainKind,
889        ResetPolicies, ResetReceipt, ResetRequirements, SnapshotRestoreFunding, StateResetPolicy,
890        TimeResetPolicy, add_timing,
891    };
892    use std::time::Duration;
893
894    #[test]
895    fn preparation_and_recovery_keep_display_and_unboxed_primary_sources() {
896        use std::{error::Error as _, io};
897        let preparation = |message| super::BaselinePoolPreparationError::Recipe {
898            stage: super::BaselinePreparationStage::Build,
899            source: io::Error::other(message),
900        };
901        let original = preparation("original");
902        assert_eq!(original.to_string(), "baseline build failed: original");
903        assert!(original.source().unwrap().is::<io::Error>());
904        let error = super::BaselinePoolError::RecoveryFailed {
905            original: Box::new(original),
906            rebuild: Box::new(preparation("replacement")),
907            timings: Box::default(),
908        };
909        assert_eq!(
910            error.to_string(),
911            "baseline preparation failed (baseline build failed: original); rebuilding the slot also failed: baseline build failed: replacement"
912        );
913        assert!(
914            error
915                .source()
916                .unwrap()
917                .is::<super::BaselinePoolPreparationError<io::Error>>()
918        );
919    }
920
921    #[test]
922    fn phase_timing_preserves_absence_zero_and_saturation() {
923        let mut timings = BaselinePoolTimings::default();
924        assert_eq!(timings.restore(), None);
925        add_timing(&mut timings.restore, Duration::ZERO);
926        assert_eq!(timings.restore(), Some(Duration::ZERO));
927        add_timing(&mut timings.restore, Duration::from_secs(2));
928        assert_eq!(timings.restore(), Some(Duration::from_secs(2)));
929        add_timing(&mut timings.restore, Duration::MAX);
930        assert_eq!(timings.restore(), Some(Duration::MAX));
931        assert_eq!(timings.build(), None);
932    }
933
934    #[test]
935    fn required_policy_must_match_achieved_policy() {
936        for (domain, requirement, achievement) in [
937            (
938                ResetDomainKind::PocketIcTime,
939                ResetPolicies {
940                    pocket_ic_time: Some(TimeResetPolicy::PreserveCurrent),
941                    ..ResetPolicies::default()
942                },
943                ResetPolicies {
944                    pocket_ic_time: Some(TimeResetPolicy::RebuildOnMutation),
945                    ..ResetPolicies::default()
946                },
947            ),
948            (
949                ResetDomainKind::ExtraCanisters,
950                ResetPolicies {
951                    extra_canisters: Some(ExtraCanisterPolicy::RequireBaselineSet),
952                    ..ResetPolicies::default()
953                },
954                ResetPolicies {
955                    extra_canisters: Some(ExtraCanisterPolicy::RemoveTracked),
956                    ..ResetPolicies::default()
957                },
958            ),
959            (
960                ResetDomainKind::PendingMessages,
961                ResetPolicies {
962                    pending_messages: Some(StateResetPolicy::ValidateUnchanged),
963                    ..ResetPolicies::default()
964                },
965                ResetPolicies {
966                    pending_messages: Some(StateResetPolicy::IrrelevantByRecipeContract),
967                    ..ResetPolicies::default()
968                },
969            ),
970            (
971                ResetDomainKind::SubnetState,
972                ResetPolicies {
973                    subnet_state: Some(StateResetPolicy::ResetByRecipe),
974                    ..ResetPolicies::default()
975                },
976                ResetPolicies {
977                    subnet_state: Some(StateResetPolicy::RebuildOnChange),
978                    ..ResetPolicies::default()
979                },
980            ),
981            (
982                ResetDomainKind::ExternalResources,
983                ResetPolicies {
984                    external_resources: Some(StateResetPolicy::ValidateUnchanged),
985                    ..ResetPolicies::default()
986                },
987                ResetPolicies {
988                    external_resources: Some(StateResetPolicy::ResetByRecipe),
989                    ..ResetPolicies::default()
990                },
991            ),
992        ] {
993            let requirements = ResetRequirements {
994                funding: SnapshotRestoreFunding::Preserve,
995                policies: requirement,
996            };
997            assert_eq!(
998                requirements.verify(ResetReceipt {
999                    policies: requirement
1000                }),
1001                Ok(())
1002            );
1003            assert_eq!(
1004                requirements.verify(ResetReceipt::default()),
1005                Err(BaselinePoolContractError::MissingResetDomain { domain })
1006            );
1007            assert_eq!(
1008                requirements.verify(ResetReceipt {
1009                    policies: achievement
1010                }),
1011                Err(BaselinePoolContractError::ResetPolicyMismatch {
1012                    domain,
1013                    requirement,
1014                    achievement
1015                })
1016            );
1017            let unconstrained = ResetRequirements {
1018                funding: SnapshotRestoreFunding::Preserve,
1019                policies: ResetPolicies::default(),
1020            };
1021            assert_eq!(
1022                unconstrained.verify(ResetReceipt {
1023                    policies: achievement
1024                }),
1025                Ok(())
1026            );
1027        }
1028    }
1029
1030    #[test]
1031    fn reset_coverage_preserves_absence_and_first_failure() {
1032        let requirements = ResetRequirements {
1033            funding: SnapshotRestoreFunding::Preserve,
1034            policies: ResetPolicies {
1035                pocket_ic_time: Some(TimeResetPolicy::PreserveCurrent),
1036                pending_messages: Some(StateResetPolicy::IrrelevantByRecipeContract),
1037                ..ResetPolicies::default()
1038            },
1039        };
1040        let mut receipt = ResetReceipt::default();
1041        assert_eq!(
1042            requirements.verify(receipt),
1043            Err(BaselinePoolContractError::MissingResetDomain {
1044                domain: ResetDomainKind::PocketIcTime,
1045            })
1046        );
1047        receipt.policies.pocket_ic_time = Some(TimeResetPolicy::RebuildOnMutation);
1048        assert_eq!(
1049            requirements.verify(receipt),
1050            Err(BaselinePoolContractError::ResetPolicyMismatch {
1051                domain: ResetDomainKind::PocketIcTime,
1052                requirement: requirements.policies,
1053                achievement: receipt.policies,
1054            })
1055        );
1056        receipt.policies.pocket_ic_time = Some(TimeResetPolicy::PreserveCurrent);
1057        assert_eq!(
1058            requirements.verify(receipt),
1059            Err(BaselinePoolContractError::MissingResetDomain {
1060                domain: ResetDomainKind::PendingMessages,
1061            })
1062        );
1063        receipt.policies.pending_messages = Some(StateResetPolicy::IrrelevantByRecipeContract);
1064        assert_eq!(requirements.verify(receipt), Ok(()));
1065    }
1066}
ic_testkit/pic/baseline_pool.rs

ic_testkit/pic/
baseline_pool.rs
