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

ic_testkit/pic/
baseline_pool.rs
