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

ic_testkit/pic/
baseline_pool.rs
