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

ic_testkit/pic/
baseline_pool.rs
