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

ic_testkit/pic/
baseline_pool.rs
