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

ic_testkit/pic/
baseline_pool.rs
