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

ic_testkit/pic/
baseline_pool.rs
