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},
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) = if slot.is_reusable() {
614            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        } else {
641            let reason = if slot.invalidated_by_unwind() {
642                Some(RebuildReason::UnwindWhileLeased)
643            } else {
644                slot.invalidation_reason().cloned()
645            };
646            (reason, None)
647        };
648
649        if slot.is_populated() {
650            Self::discard_stale_slot(&mut slot, &mut timings);
651        }
652        let build = self.build_slot(&mut slot, &mut timings);
653        timings.total = total_started.elapsed();
654        if let Err(error) = build {
655            return Err(match original {
656                Some(original) => BaselinePoolError::RecoveryFailed {
657                    original: Box::new(original),
658                    rebuild: Box::new(error),
659                    timings: Box::new(timings),
660                },
661                None => BaselinePoolError::Preparation {
662                    error,
663                    timings: Box::new(timings),
664                },
665            });
666        }
667
668        let outcome = BaselinePoolOutcome {
669            slot: slot_index,
670            timings,
671            disposition: rebuild_reason.map_or(
672                BaselinePoolDisposition::Built,
673                BaselinePoolDisposition::Rebuilt,
674            ),
675        };
676        Ok((CachedPocketIcBaselinePoolGuard { slot }, outcome))
677    }
678
679    fn prepare_reused(
680        &self,
681        slot: &BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
682        timings: &mut BaselinePoolTimings,
683    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
684        let baseline = slot
685            .get()
686            .expect("reusable baseline slot must be populated");
687
688        let started = Instant::now();
689        let restore = self
690            .recipe
691            .prepare_restore(baseline)
692            .map_err(|source| BaselinePoolPreparationError::Recipe {
693                stage: BaselinePreparationStage::RestoreCanisters,
694                source,
695            })
696            .and_then(|()| {
697                baseline
698                    .restore(self.recipe.reset_requirements().funding)
699                    .map_err(BaselinePoolPreparationError::Restore)
700            });
701        add_timing(&mut timings.restore, started.elapsed());
702        restore?;
703
704        let reset = timed_recipe_stage(
705            &mut timings.reset,
706            BaselinePreparationStage::ResetNonSnapshotState,
707            || self.recipe.reset_non_snapshot_state(baseline),
708        )?;
709        timed_recipe_stage(
710            &mut timings.readiness,
711            BaselinePreparationStage::DriveToReadiness,
712            || self.recipe.drive_to_readiness(baseline),
713        )?;
714        self.recipe
715            .reset_requirements()
716            .verify(reset)
717            .map_err(BaselinePoolPreparationError::Contract)?;
718
719        let preparation = PreparedBaseline::Restored { reset };
720        self.validate_baseline(
721            baseline,
722            &preparation,
723            BaselinePreparationStage::ValidateRestored,
724            timings,
725        )?;
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(
740            &baseline,
741            &PreparedBaseline::Built,
742            BaselinePreparationStage::ValidateBuilt,
743            timings,
744        ) {
745            drop_baseline_safely(baseline);
746            return Err(error);
747        }
748        let replaced = slot.replace(baseline);
749        debug_assert!(replaced.is_none());
750        Ok(())
751    }
752
753    fn validate_baseline(
754        &self,
755        baseline: &CachedPocketIcBaseline<R::Metadata>,
756        preparation: &PreparedBaseline,
757        stage: BaselinePreparationStage,
758        timings: &mut BaselinePoolTimings,
759    ) -> Result<(), BaselinePoolPreparationError<R::Error>> {
760        timed_recipe_stage(&mut timings.validation, stage, || {
761            self.recipe.validate(baseline, preparation)
762        })
763    }
764
765    fn failure_disposition(
766        &self,
767        error: &BaselinePoolPreparationError<R::Error>,
768    ) -> FailureDisposition {
769        match error {
770            BaselinePoolPreparationError::Recipe { stage, source } => {
771                self.recipe.classify_failure(*stage, source)
772            }
773            BaselinePoolPreparationError::Restore(source) => self
774                .recipe
775                .classify_failure(BaselinePreparationStage::RestoreCanisters, source),
776            BaselinePoolPreparationError::Contract(_) => {
777                FailureDisposition::Rebuild(RebuildReason::ResetCoverageMismatch)
778            }
779        }
780    }
781
782    fn discard_stale_slot(
783        slot: &mut BoundedSlotLease<'_, CachedPocketIcBaseline<R::Metadata>, RebuildReason>,
784        timings: &mut BaselinePoolTimings,
785    ) {
786        let started = Instant::now();
787        if let Some(stale) = slot.take() {
788            drop_baseline_safely(stale);
789        }
790        timings.stale_teardown = Some(started.elapsed());
791    }
792}
793
794impl<R> CachedPocketIcBaselinePoolGuard<'_, R>
795where
796    R: PocketIcBaselineRecipe,
797{
798    /// Diagnostic slot index held by this lease.
799    #[must_use]
800    pub const fn slot(&self) -> usize {
801        self.slot.slot_index()
802    }
803
804    /// Mark this slot non-reusable with a structured rebuild reason.
805    pub fn invalidate(&mut self, reason: RebuildReason) {
806        self.slot.invalidate(reason);
807    }
808}
809
810impl<R> Deref for CachedPocketIcBaselinePoolGuard<'_, R>
811where
812    R: PocketIcBaselineRecipe,
813{
814    type Target = CachedPocketIcBaseline<R::Metadata>;
815
816    fn deref(&self) -> &Self::Target {
817        self.slot
818            .get()
819            .expect("leased baseline pool slot must be populated")
820    }
821}
822
823fn timed_recipe_stage<T, E>(
824    timing: &mut Option<Duration>,
825    stage: BaselinePreparationStage,
826    operation: impl FnOnce() -> Result<T, E>,
827) -> Result<T, BaselinePoolPreparationError<E>> {
828    let started = Instant::now();
829    let result = operation();
830    add_timing(timing, started.elapsed());
831    result.map_err(|source| BaselinePoolPreparationError::Recipe { stage, source })
832}
833
834const fn add_timing(total: &mut Option<Duration>, elapsed: Duration) {
835    *total = Some(match *total {
836        None => elapsed,
837        Some(previous) => previous.saturating_add(elapsed),
838    });
839}
840
841fn drop_baseline_safely<M>(baseline: CachedPocketIcBaseline<M>) {
842    let _ = catch_unwind(AssertUnwindSafe(|| drop(baseline)));
843}
844
845impl std::fmt::Display for BaselinePreparationStage {
846    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
847        formatter.write_str(match self {
848            Self::Build => "build",
849            Self::RestoreCanisters => "canister restore",
850            Self::ResetNonSnapshotState => "non-snapshot reset",
851            Self::DriveToReadiness => "readiness",
852            Self::ValidateBuilt => "built-baseline validation",
853            Self::ValidateRestored => "restored-baseline validation",
854        })
855    }
856}
857
858impl std::fmt::Display for BaselinePoolContractError {
859    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
860        match self {
861            Self::MissingResetDomain { domain } => {
862                write!(
863                    formatter,
864                    "required reset domain {domain:?} was not achieved"
865                )
866            }
867            Self::ResetPolicyMismatch {
868                domain,
869                requirement,
870                achievement,
871            } => write!(
872                formatter,
873                "reset achievement {achievement:?} does not satisfy {requirement:?} for {domain:?}",
874            ),
875        }
876    }
877}
878
879impl std::error::Error for BaselinePoolContractError {}
880
881impl<E> std::fmt::Display for BaselinePoolPreparationError<E>
882where
883    E: std::fmt::Display,
884{
885    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
886        match self {
887            Self::Recipe { stage, source } => {
888                write!(formatter, "baseline {stage} failed: {source}")
889            }
890            Self::Restore(source) => {
891                write!(formatter, "baseline canister restore failed: {source}")
892            }
893            Self::Contract(error) => write!(formatter, "baseline pool contract failed: {error}"),
894        }
895    }
896}
897
898impl<E> std::error::Error for BaselinePoolPreparationError<E>
899where
900    E: std::error::Error + 'static,
901{
902    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
903        match self {
904            Self::Recipe { source, .. } => Some(source),
905            Self::Restore(source) => Some(source),
906            Self::Contract(error) => Some(error),
907        }
908    }
909}
910
911impl<E> std::fmt::Display for BaselinePoolError<E>
912where
913    E: std::fmt::Display,
914{
915    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
916        match self {
917            Self::Preparation { error, .. } => error.fmt(formatter),
918            Self::RecoveryFailed {
919                original, rebuild, ..
920            } => write!(
921                formatter,
922                "baseline preparation failed ({original}); rebuilding the slot also failed: {rebuild}",
923            ),
924        }
925    }
926}
927
928impl<E> std::error::Error for BaselinePoolError<E>
929where
930    E: std::error::Error + 'static,
931{
932    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
933        match self {
934            Self::Preparation { error, .. } => Some(error),
935            Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
936        }
937    }
938}
939
940#[cfg(test)]
941mod tests {
942    use super::{
943        BaselinePoolContractError, BaselinePoolTimings, ExtraCanisterPolicy, ResetDomainKind,
944        ResetPolicies, ResetReceipt, ResetRequirements, SnapshotRestoreFunding, StateResetPolicy,
945        TimeResetPolicy, add_timing,
946    };
947    use std::time::Duration;
948
949    #[test]
950    fn phase_timing_preserves_absence_zero_and_saturation() {
951        let mut timings = BaselinePoolTimings::default();
952        assert_eq!(timings.restore(), None);
953        add_timing(&mut timings.restore, Duration::ZERO);
954        assert_eq!(timings.restore(), Some(Duration::ZERO));
955        add_timing(&mut timings.restore, Duration::from_secs(2));
956        assert_eq!(timings.restore(), Some(Duration::from_secs(2)));
957        add_timing(&mut timings.restore, Duration::MAX);
958        assert_eq!(timings.restore(), Some(Duration::MAX));
959        assert_eq!(timings.build(), None);
960    }
961
962    #[test]
963    fn required_policy_must_match_achieved_policy() {
964        for (domain, requirement, achievement) in [
965            (
966                ResetDomainKind::PocketIcTime,
967                ResetPolicies {
968                    pocket_ic_time: Some(TimeResetPolicy::PreserveCurrent),
969                    ..ResetPolicies::default()
970                },
971                ResetPolicies {
972                    pocket_ic_time: Some(TimeResetPolicy::RebuildOnMutation),
973                    ..ResetPolicies::default()
974                },
975            ),
976            (
977                ResetDomainKind::ExtraCanisters,
978                ResetPolicies {
979                    extra_canisters: Some(ExtraCanisterPolicy::RequireBaselineSet),
980                    ..ResetPolicies::default()
981                },
982                ResetPolicies {
983                    extra_canisters: Some(ExtraCanisterPolicy::RemoveTracked),
984                    ..ResetPolicies::default()
985                },
986            ),
987            (
988                ResetDomainKind::PendingMessages,
989                ResetPolicies {
990                    pending_messages: Some(StateResetPolicy::ValidateUnchanged),
991                    ..ResetPolicies::default()
992                },
993                ResetPolicies {
994                    pending_messages: Some(StateResetPolicy::IrrelevantByRecipeContract),
995                    ..ResetPolicies::default()
996                },
997            ),
998            (
999                ResetDomainKind::SubnetState,
1000                ResetPolicies {
1001                    subnet_state: Some(StateResetPolicy::ResetByRecipe),
1002                    ..ResetPolicies::default()
1003                },
1004                ResetPolicies {
1005                    subnet_state: Some(StateResetPolicy::RebuildOnChange),
1006                    ..ResetPolicies::default()
1007                },
1008            ),
1009            (
1010                ResetDomainKind::ExternalResources,
1011                ResetPolicies {
1012                    external_resources: Some(StateResetPolicy::ValidateUnchanged),
1013                    ..ResetPolicies::default()
1014                },
1015                ResetPolicies {
1016                    external_resources: Some(StateResetPolicy::ResetByRecipe),
1017                    ..ResetPolicies::default()
1018                },
1019            ),
1020        ] {
1021            let requirements = ResetRequirements {
1022                funding: SnapshotRestoreFunding::Preserve,
1023                policies: requirement,
1024            };
1025            assert_eq!(
1026                requirements.verify(ResetReceipt {
1027                    policies: requirement
1028                }),
1029                Ok(())
1030            );
1031            assert_eq!(
1032                requirements.verify(ResetReceipt::default()),
1033                Err(BaselinePoolContractError::MissingResetDomain { domain })
1034            );
1035            assert_eq!(
1036                requirements.verify(ResetReceipt {
1037                    policies: achievement
1038                }),
1039                Err(BaselinePoolContractError::ResetPolicyMismatch {
1040                    domain,
1041                    requirement,
1042                    achievement
1043                })
1044            );
1045            let unconstrained = ResetRequirements {
1046                funding: SnapshotRestoreFunding::Preserve,
1047                policies: ResetPolicies::default(),
1048            };
1049            assert_eq!(
1050                unconstrained.verify(ResetReceipt {
1051                    policies: achievement
1052                }),
1053                Ok(())
1054            );
1055        }
1056    }
1057
1058    #[test]
1059    fn reset_coverage_preserves_absence_and_first_failure() {
1060        let requirements = ResetRequirements {
1061            funding: SnapshotRestoreFunding::Preserve,
1062            policies: ResetPolicies {
1063                pocket_ic_time: Some(TimeResetPolicy::PreserveCurrent),
1064                pending_messages: Some(StateResetPolicy::IrrelevantByRecipeContract),
1065                ..ResetPolicies::default()
1066            },
1067        };
1068        let mut receipt = ResetReceipt::default();
1069        assert_eq!(
1070            requirements.verify(receipt),
1071            Err(BaselinePoolContractError::MissingResetDomain {
1072                domain: ResetDomainKind::PocketIcTime,
1073            })
1074        );
1075        receipt.policies.pocket_ic_time = Some(TimeResetPolicy::RebuildOnMutation);
1076        assert_eq!(
1077            requirements.verify(receipt),
1078            Err(BaselinePoolContractError::ResetPolicyMismatch {
1079                domain: ResetDomainKind::PocketIcTime,
1080                requirement: requirements.policies,
1081                achievement: receipt.policies,
1082            })
1083        );
1084        receipt.policies.pocket_ic_time = Some(TimeResetPolicy::PreserveCurrent);
1085        assert_eq!(
1086            requirements.verify(receipt),
1087            Err(BaselinePoolContractError::MissingResetDomain {
1088                domain: ResetDomainKind::PendingMessages,
1089            })
1090        );
1091        receipt.policies.pending_messages = Some(StateResetPolicy::IrrelevantByRecipeContract);
1092        assert_eq!(requirements.verify(receipt), Ok(()));
1093    }
1094}
ic_testkit/pic/baseline_pool.rs

ic_testkit/pic/
baseline_pool.rs
