ic_testkit/pic/

standalone_pool.rs

use std::{
    num::NonZeroUsize,
    ops::Deref,
    panic::{AssertUnwindSafe, catch_unwind},
    time::{Duration, Instant},
};

use super::{
    ControllerSnapshotError, ControllerSnapshots, PocketIcSnapshotExt, SnapshotRestoreFunding,
    StandaloneCanisterFixture,
    bounded_pool::{BoundedSlotLease, BoundedSlotPool},
    transport,
};

struct StandaloneFixtureBaseline {
    fixture: StandaloneCanisterFixture,
    snapshots: ControllerSnapshots,
}

impl StandaloneFixtureBaseline {
    fn capture(fixture: StandaloneCanisterFixture) -> Result<Self, ControllerSnapshotError> {
        let canister_id = fixture.canister_id();
        let snapshots = fixture
            .pocket_ic()
            .capture_controller_snapshots(canister_id, [canister_id])?;

        Ok(Self { fixture, snapshots })
    }

    fn restore(&self, funding: SnapshotRestoreFunding) -> Result<(), ControllerSnapshotError> {
        self.fixture
            .pocket_ic()
            .restore_snapshots_with_captured_senders_and_funding(&self.snapshots, funding)
    }
}

/// Timings for one standalone fixture-pool acquisition.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct StandaloneFixturePoolTimings {
    wait: Duration,
    build: Option<Duration>,
    restore: Option<Duration>,
    stale_teardown: Option<Duration>,
    total: Duration,
}

/// Whether a standalone fixture-pool lease was built, restored, or rebuilt.
#[non_exhaustive]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum StandaloneFixturePoolOutcome {
    /// An empty slot was built and its snapshot was captured.
    Built {
        /// Diagnostic slot index.
        slot: usize,
        /// Acquisition phase timings.
        timings: StandaloneFixturePoolTimings,
    },
    /// A populated slot was restored successfully.
    Restored {
        /// Diagnostic slot index.
        slot: usize,
        /// Acquisition phase timings.
        timings: StandaloneFixturePoolTimings,
    },
    /// An invalid or dead slot was rebuilt.
    Rebuilt {
        /// Diagnostic slot index.
        slot: usize,
        /// Reason the previous slot could not be reused.
        reason: StandaloneFixturePoolRebuildReason,
        /// Acquisition phase timings.
        timings: StandaloneFixturePoolTimings,
    },
}

/// Why a standalone fixture-pool slot was rebuilt.
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StandaloneFixturePoolRebuildReason {
    /// PocketIC transport was no longer reachable during snapshot restoration.
    DeadPocketIcTransport,
    /// A previous restore failed after it may have partially changed the slot.
    PreviousRestoreFailure,
    /// A lease was dropped while its thread was unwinding.
    UnwindWhileLeased,
}

/// Standalone fixture preparation stage associated with a snapshot failure.
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StandaloneFixturePoolStage {
    /// Building a fixture and capturing its baseline snapshot.
    Build,
    /// Restoring a populated fixture's baseline snapshot.
    Restore,
}

/// Failure to acquire a standalone fixture-pool lease with structured diagnostics.
#[non_exhaustive]
#[derive(Debug)]
pub enum StandaloneFixturePoolError {
    /// Initial build/capture or snapshot restoration failed.
    Preparation {
        /// Failed lifecycle stage.
        stage: StandaloneFixturePoolStage,
        /// Structured snapshot failure.
        source: Box<ControllerSnapshotError>,
        /// Timings recorded before acquisition failed.
        timings: Box<StandaloneFixturePoolTimings>,
    },
    /// Dead-transport restoration failed and replacement snapshot capture also failed.
    RecoveryFailed {
        /// Original dead-transport restoration failure.
        original: Box<ControllerSnapshotError>,
        /// Snapshot failure while building the replacement slot.
        rebuild: Box<ControllerSnapshotError>,
        /// Combined restore, teardown, and rebuild timings.
        timings: Box<StandaloneFixturePoolTimings>,
    },
}

/// Caller-owned runtime-capacity pool of independently restorable standalone fixtures.
///
/// Each slot owns one PocketIC instance, one installed canister, and one
/// captured baseline snapshot. Acquiring a populated slot restores that
/// snapshot before returning it. At most [`Self::capacity`] leases can overlap;
/// a caller waits only when every slot is in use.
///
/// One pool owns one fixture builder. The builder must produce the same Wasm,
/// init arguments, topology, and seeded baseline each time it runs. It is not
/// evaluated on a cache hit; use a separate pool for each recipe.
///
/// Snapshot restoration rewinds the installed canister, not the surrounding
/// PocketIC instance. Instance time, other canisters, and cycle changes not
/// covered by the selected [`SnapshotRestoreFunding`] policy may persist.
///
/// The pool contains no process-global state. Downstream suites select a
/// capacity that fits their host and keep lifecycle-sensitive tests on fresh
/// [`StandaloneCanisterFixture`] values when snapshot restoration is not the
/// intended isolation boundary.
pub struct CachedStandaloneCanisterFixturePool<B = fn() -> StandaloneCanisterFixture> {
    build: B,
    slots: BoundedSlotPool<StandaloneFixtureBaseline>,
    restore_funding: SnapshotRestoreFunding,
}

/// Exclusive lease of one independently restored standalone fixture.
///
/// The lease dereferences to [`StandaloneCanisterFixture`], so existing call
/// helpers can borrow it without adding a second fixture API.
pub struct CachedStandaloneCanisterFixtureGuard<'a> {
    slot: BoundedSlotLease<'a, StandaloneFixtureBaseline>,
}

impl<B> CachedStandaloneCanisterFixturePool<B>
where
    B: Fn() -> StandaloneCanisterFixture,
{
    /// Create a runtime-capacity pool that owns its fixture builder.
    ///
    /// Construction is const and lazy: slots are allocated on first acquisition,
    /// and the builder runs only when a slot needs a fixture.
    ///
    /// ```no_run
    /// use std::num::NonZeroUsize;
    /// use ic_testkit::pic::{CachedStandaloneCanisterFixturePool, StandaloneCanisterFixture};
    /// # fn build_fixture() -> StandaloneCanisterFixture { unimplemented!() }
    ///
    /// static POOL: CachedStandaloneCanisterFixturePool =
    ///     CachedStandaloneCanisterFixturePool::new(
    ///         NonZeroUsize::new(2).unwrap(),
    ///         build_fixture,
    ///     );
    ///
    /// // Local pools can choose their capacity at runtime.
    /// let capacity = std::thread::available_parallelism()?;
    /// let pool = CachedStandaloneCanisterFixturePool::new(capacity, build_fixture);
    /// # Ok::<(), std::io::Error>(())
    /// ```
    #[must_use]
    pub const fn new(capacity: NonZeroUsize, build: B) -> Self {
        Self {
            build,
            slots: BoundedSlotPool::new(capacity),
            restore_funding: SnapshotRestoreFunding::Preserve,
        }
    }

    /// Maximum number of simultaneously leased standalone fixtures.
    #[must_use]
    pub const fn capacity(&self) -> NonZeroUsize {
        self.slots.capacity()
    }

    /// Select the cycle-funding policy applied immediately before each
    /// snapshot restore.
    #[must_use]
    pub const fn with_restore_funding(mut self, funding: SnapshotRestoreFunding) -> Self {
        self.restore_funding = funding;
        self
    }

    /// Acquire one isolated fixture, building a slot on first use and restoring
    /// its captured snapshot on later uses.
    ///
    /// The pool's builder runs only when an empty or invalid slot is populated
    /// or a dead PocketIC instance must be replaced.
    ///
    /// A recognized dead-instance transport failure evicts and rebuilds only
    /// the affected slot. Other snapshot failures are returned unchanged and
    /// invalidate the possibly partially restored slot for the next lease.
    ///
    /// Acquire one fixture with a structured lifecycle outcome and phase timings.
    ///
    /// # Errors
    ///
    /// Returns the failed preparation stage, the structured snapshot error,
    /// and all phase timings completed before failure. If dead-transport
    /// restoration and replacement capture both fail, both errors are retained.
    pub fn acquire(
        &self,
    ) -> Result<
        (
            CachedStandaloneCanisterFixtureGuard<'_>,
            StandaloneFixturePoolOutcome,
        ),
        StandaloneFixturePoolError,
    > {
        let total_started = Instant::now();
        let mut slot = self.slots.acquire();
        let slot_index = slot.slot_index();
        let mut timings = StandaloneFixturePoolTimings {
            wait: slot.wait(),
            ..StandaloneFixturePoolTimings::default()
        };

        let (rebuild_reason, original) = if slot.is_reusable() {
            let restore_started = Instant::now();
            let restore = slot
                .get()
                .expect("populated fixture pool slot must remain present")
                .restore(self.restore_funding);
            timings.restore = Some(restore_started.elapsed());
            match restore {
                Ok(()) => {
                    timings.total = total_started.elapsed();
                    return Ok((
                        CachedStandaloneCanisterFixtureGuard { slot },
                        StandaloneFixturePoolOutcome::Restored {
                            slot: slot_index,
                            timings,
                        },
                    ));
                }
                Err(error) if transport::is_dead_pocket_ic_transport_error(&error) => (
                    Some(StandaloneFixturePoolRebuildReason::DeadPocketIcTransport),
                    Some(error),
                ),
                Err(source) => {
                    // Restoration may have partially changed the canister.
                    // Preserve the failure and prevent reuse on the next lease.
                    slot.invalidate();
                    timings.total = total_started.elapsed();
                    return Err(StandaloneFixturePoolError::Preparation {
                        stage: StandaloneFixturePoolStage::Restore,
                        source: Box::new(source),
                        timings: Box::new(timings),
                    });
                }
            }
        } else {
            (Self::rebuild_reason_for_invalid_slot(&slot), None)
        };

        Self::discard_stale_slot(&mut slot, &mut timings);
        let baseline = match self.build_slot(&mut timings) {
            Ok(baseline) => baseline,
            Err(source) => {
                timings.total = total_started.elapsed();
                return Err(match original {
                    Some(original) => StandaloneFixturePoolError::RecoveryFailed {
                        original: Box::new(original),
                        rebuild: Box::new(source),
                        timings: Box::new(timings),
                    },
                    None => StandaloneFixturePoolError::Preparation {
                        stage: StandaloneFixturePoolStage::Build,
                        source: Box::new(source),
                        timings: Box::new(timings),
                    },
                });
            }
        };
        slot.replace(baseline);
        timings.total = total_started.elapsed();
        let outcome = rebuild_reason.map_or_else(
            || StandaloneFixturePoolOutcome::Built {
                slot: slot_index,
                timings,
            },
            |reason| StandaloneFixturePoolOutcome::Rebuilt {
                slot: slot_index,
                reason,
                timings,
            },
        );
        Ok((CachedStandaloneCanisterFixtureGuard { slot }, outcome))
    }

    fn rebuild_reason_for_invalid_slot(
        slot: &BoundedSlotLease<'_, StandaloneFixtureBaseline>,
    ) -> Option<StandaloneFixturePoolRebuildReason> {
        if slot.invalidated_by_unwind() {
            Some(StandaloneFixturePoolRebuildReason::UnwindWhileLeased)
        } else {
            // Non-unwind invalidation only follows a failed restore.
            slot.is_populated()
                .then_some(StandaloneFixturePoolRebuildReason::PreviousRestoreFailure)
        }
    }

    fn build_slot(
        &self,
        timings: &mut StandaloneFixturePoolTimings,
    ) -> Result<StandaloneFixtureBaseline, ControllerSnapshotError> {
        let started = Instant::now();
        let result = StandaloneFixtureBaseline::capture((self.build)());
        timings.build = Some(started.elapsed());
        result
    }

    fn discard_stale_slot(
        slot: &mut BoundedSlotLease<'_, StandaloneFixtureBaseline>,
        timings: &mut StandaloneFixturePoolTimings,
    ) {
        if !slot.is_populated() {
            return;
        }
        let started = Instant::now();
        if let Some(stale) = slot.take() {
            let _ = catch_unwind(AssertUnwindSafe(|| drop(stale)));
        }
        timings.stale_teardown = Some(started.elapsed());
    }
}

impl Deref for CachedStandaloneCanisterFixtureGuard<'_> {
    type Target = StandaloneCanisterFixture;

    fn deref(&self) -> &Self::Target {
        &self
            .slot
            .get()
            .expect("leased fixture pool slot must remain populated")
            .fixture
    }
}

impl StandaloneFixturePoolOutcome {
    /// Diagnostic slot index used by this acquisition.
    #[must_use]
    pub const fn slot(&self) -> usize {
        match self {
            Self::Built { slot, .. } | Self::Restored { slot, .. } | Self::Rebuilt { slot, .. } => {
                *slot
            }
        }
    }

    /// Acquisition phase timings.
    #[must_use]
    pub const fn timings(&self) -> StandaloneFixturePoolTimings {
        match self {
            Self::Built { timings, .. }
            | Self::Restored { timings, .. }
            | Self::Rebuilt { timings, .. } => *timings,
        }
    }

    /// Report whether an existing slot was restored without reconstruction.
    #[must_use]
    pub const fn is_reused(&self) -> bool {
        matches!(self, Self::Restored { .. })
    }
}

impl StandaloneFixturePoolTimings {
    /// Time spent waiting for a capacity slot.
    #[must_use]
    pub const fn wait(self) -> Duration {
        self.wait
    }

    /// Time spent building a fixture and capturing its baseline snapshot.
    #[must_use]
    pub const fn build(self) -> Option<Duration> {
        self.build
    }

    /// Time spent restoring a populated slot.
    #[must_use]
    pub const fn restore(self) -> Option<Duration> {
        self.restore
    }

    /// Time spent dropping an invalid or dead slot before rebuilding.
    #[must_use]
    pub const fn stale_teardown(self) -> Option<Duration> {
        self.stale_teardown
    }

    /// Complete acquisition duration.
    #[must_use]
    pub const fn total(self) -> Duration {
        self.total
    }
}

impl std::fmt::Display for StandaloneFixturePoolTimings {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            formatter,
            "total={:?} wait={:?} build={:?} restore={:?} stale_teardown={:?}",
            self.total, self.wait, self.build, self.restore, self.stale_teardown,
        )
    }
}

impl std::fmt::Display for StandaloneFixturePoolOutcome {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Built { slot, timings } => write!(formatter, "built slot={slot} {timings}"),
            Self::Restored { slot, timings } => {
                write!(formatter, "restored slot={slot} {timings}")
            }
            Self::Rebuilt {
                slot,
                reason,
                timings,
            } => write!(formatter, "rebuilt slot={slot} reason={reason:?} {timings}"),
        }
    }
}

impl StandaloneFixturePoolError {
    /// Timings recorded before this acquisition failed.
    #[must_use]
    pub const fn timings(&self) -> StandaloneFixturePoolTimings {
        match self {
            Self::Preparation { timings, .. } | Self::RecoveryFailed { timings, .. } => **timings,
        }
    }
}

impl std::fmt::Display for StandaloneFixturePoolStage {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(match self {
            Self::Build => "fixture build and snapshot capture",
            Self::Restore => "fixture snapshot restore",
        })
    }
}

impl std::fmt::Display for StandaloneFixturePoolError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Preparation { stage, source, .. } => {
                write!(formatter, "standalone {stage} failed: {source}")
            }
            Self::RecoveryFailed {
                original, rebuild, ..
            } => write!(
                formatter,
                "standalone fixture restore failed ({original}); rebuilding the slot also failed: {rebuild}",
            ),
        }
    }
}

impl std::error::Error for StandaloneFixturePoolError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Preparation { source, .. } => Some(source.as_ref()),
            Self::RecoveryFailed { original, .. } => Some(original.as_ref()),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::num::NonZeroUsize;

    use super::{
        CachedStandaloneCanisterFixturePool, SnapshotRestoreFunding, StandaloneCanisterFixture,
    };

    const _: CachedStandaloneCanisterFixturePool = {
        let pool: CachedStandaloneCanisterFixturePool =
            CachedStandaloneCanisterFixturePool::new(NonZeroUsize::new(1).unwrap(), build_fixture);
        pool.with_restore_funding(SnapshotRestoreFunding::TopUpTo { minimum_cycles: 1 })
    };

    fn build_fixture() -> StandaloneCanisterFixture {
        panic!("constructing an empty pool must not invoke its builder")
    }

    #[test]
    fn runtime_capacity_preserves_lazy_fixture_construction() {
        for capacity in [1, 2, 4, 8].map(|value| NonZeroUsize::new(value).unwrap()) {
            let pool = CachedStandaloneCanisterFixturePool::new(capacity, build_fixture);
            assert_eq!(pool.capacity(), capacity);
        }
    }
}