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

ic_testkit/pic/
baseline_pool.rs
