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

ic_testkit/pic/
baseline_pool.rs
