ic_testkit/pic/
snapshot.rs

1use std::collections::BTreeMap;
2
3use ic_platform::pocket_ic::{PocketIc, RejectResponse};
4use ic_platform::principal::ExternalPrincipal as Principal;
5
6use super::{PocketIcCallFailure, transport};
7
8#[derive(Clone, Debug, Eq, PartialEq)]
9struct ControllerSnapshot {
10    snapshot_id: Vec<u8>,
11    sender: Option<Principal>,
12}
13
14/// Deterministically ordered snapshots retaining each successful capture sender.
15#[derive(Clone, Debug, Eq, PartialEq)]
16pub struct ControllerSnapshots(BTreeMap<Principal, ControllerSnapshot>);
17
18/// One canister and exact management-call sender used for snapshot capture.
19#[derive(Clone, Copy, Debug, Eq, PartialEq)]
20pub struct CanisterSnapshotTarget {
21    /// Canister whose snapshot will be captured.
22    pub canister_id: Principal,
23    /// Exact management-call sender; None selects `PocketIC`'s default.
24    pub sender: Option<Principal>,
25}
26
27/// Failure to remove a snapshot while rolling back a partial capture.
28#[non_exhaustive]
29#[derive(Debug)]
30pub struct SnapshotCleanupFailure {
31    /// Canister whose snapshot could not be removed.
32    pub canister_id: Principal,
33    /// Exact sender selected for cleanup.
34    pub sender: Option<Principal>,
35    /// Complete rejection or operation failure, retaining transport classification.
36    pub source: Box<PocketIcCallFailure>,
37}
38
39/// Caller-selected cycle funding applied immediately before snapshot restore.
40#[non_exhaustive]
41#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42pub enum SnapshotRestoreFunding {
43    /// Do not add cycles before the restore operation.
44    Preserve,
45    /// Add cycles only when needed to reach the given minimum balance.
46    TopUpTo {
47        /// Minimum balance established immediately before each restore.
48        minimum_cycles: u128,
49    },
50}
51
52/// Structured controller-snapshot failure.
53#[non_exhaustive]
54#[derive(Debug, thiserror::Error)]
55pub enum ControllerSnapshotError {
56    /// A cached baseline requires at least one canister.
57    #[error("a cached baseline requires at least one canister")]
58    EmptyCanisterSet,
59    /// Input contained the same canister more than once.
60    #[error("duplicate canister id in snapshot set: {canister_id}")]
61    DuplicateCanisterId {
62        /// Repeated canister id.
63        canister_id: Principal,
64    },
65    /// Snapshot capture failed with the exact selected sender.
66    #[error("failed to capture snapshot for {canister_id} with sender {sender:?}: {source}; {} partial snapshots could not be cleaned up", .cleanup_failures.len())]
67    CaptureFailed {
68        /// Canister whose capture failed.
69        canister_id: Principal,
70        /// Exact sender selected for capture.
71        sender: Option<Principal>,
72        /// Complete rejection or operation failure.
73        source: Box<PocketIcCallFailure>,
74        /// Failures while deleting snapshots captured earlier in the set.
75        cleanup_failures: Vec<SnapshotCleanupFailure>,
76    },
77    /// Snapshot restore failed with the retained capture sender.
78    #[error("failed to restore snapshot for {canister_id} with sender {sender:?}: {source}")]
79    RestoreFailed {
80        /// Canister whose restore failed.
81        canister_id: Principal,
82        /// Exact sender selected for restore.
83        sender: Option<Principal>,
84        /// Complete rejection or operation failure.
85        source: Box<PocketIcCallFailure>,
86    },
87}
88
89/// Exact-sender capture and restore of related canister snapshots.
90pub trait PocketIcSnapshotExt {
91    /// Capture snapshots with one explicit sender per canister.
92    ///
93    /// Input is duplicate-checked before any effects and captured in deterministic
94    /// canister-id order. A later failure rolls back snapshots captured earlier
95    /// in the set. No other sender is attempted after a rejection.
96    ///
97    /// # Errors
98    /// Returns a target-admission or snapshot-capture error, retaining completed work.
99    fn capture_snapshots<I>(
100        &self,
101        targets: I,
102    ) -> Result<ControllerSnapshots, ControllerSnapshotError>
103    where
104        I: IntoIterator<Item = CanisterSnapshotTarget>;
105
106    /// Restore each snapshot with its exact captured sender and selected funding.
107    ///
108    /// `Preserve` never adds cycles; `TopUpTo` is evaluated before each restore.
109    /// `PocketIC` may charge cycles for the restore itself. Rejection or panic stops
110    /// the set without retrying another sender or rewinding earlier restores.
111    ///
112    /// # Errors
113    /// Returns a funding or restoration error, retaining partial progress.
114    fn restore_snapshots(
115        &self,
116        snapshots: &ControllerSnapshots,
117        funding: SnapshotRestoreFunding,
118    ) -> Result<(), ControllerSnapshotError>;
119}
120
121impl PocketIcSnapshotExt for PocketIc {
122    fn capture_snapshots<I>(
123        &self,
124        targets: I,
125    ) -> Result<ControllerSnapshots, ControllerSnapshotError>
126    where
127        I: IntoIterator<Item = CanisterSnapshotTarget>,
128    {
129        capture_snapshot_set(self, targets)
130    }
131
132    fn restore_snapshots(
133        &self,
134        snapshots: &ControllerSnapshots,
135        funding: SnapshotRestoreFunding,
136    ) -> Result<(), ControllerSnapshotError> {
137        for (canister_id, snapshot_id, sender) in snapshots.iter() {
138            restore_controller_snapshot(self, canister_id, snapshot_id, funding, sender)?;
139        }
140        Ok(())
141    }
142}
143
144fn capture_snapshot_set<I>(
145    pocket_ic: &PocketIc,
146    targets: I,
147) -> Result<ControllerSnapshots, ControllerSnapshotError>
148where
149    I: IntoIterator<Item = CanisterSnapshotTarget>,
150{
151    // Validate the entire input before issuing any management calls. The map
152    // also owns the deterministic canister order.
153    let mut ordered_targets = BTreeMap::new();
154    for target in targets {
155        if ordered_targets
156            .insert(target.canister_id, target.sender)
157            .is_some()
158        {
159            return Err(ControllerSnapshotError::DuplicateCanisterId {
160                canister_id: target.canister_id,
161            });
162        }
163    }
164    let mut snapshots = BTreeMap::new();
165    for (canister_id, sender) in ordered_targets {
166        let failure = match try_take_snapshot(pocket_ic, canister_id, sender) {
167            Ok(snapshot) => {
168                snapshots.insert(canister_id, snapshot);
169                continue;
170            }
171            Err(failure) => failure,
172        };
173        let cleanup_failures =
174            cleanup_captured_snapshots(snapshots, |canister_id, sender, snapshot_id| {
175                pocket_ic.delete_canister_snapshot(canister_id, sender, snapshot_id)
176            });
177        return Err(ControllerSnapshotError::CaptureFailed {
178            canister_id,
179            sender,
180            source: Box::new(failure),
181            cleanup_failures,
182        });
183    }
184    Ok(ControllerSnapshots(snapshots))
185}
186
187impl ControllerSnapshots {
188    /// Return the number of captured canisters.
189    #[must_use]
190    pub fn len(&self) -> usize {
191        self.0.len()
192    }
193
194    /// Report whether the set contains no snapshots.
195    #[must_use]
196    pub fn is_empty(&self) -> bool {
197        self.0.is_empty()
198    }
199
200    /// Iterate over captured canister ids in deterministic principal order.
201    pub fn canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
202        self.0.keys().copied()
203    }
204
205    pub(super) fn iter(&self) -> impl Iterator<Item = (Principal, &[u8], Option<Principal>)> + '_ {
206        self.0.iter().map(|(canister_id, snapshot)| {
207            (
208                *canister_id,
209                snapshot.snapshot_id.as_slice(),
210                snapshot.sender,
211            )
212        })
213    }
214}
215
216fn try_take_snapshot(
217    pocket_ic: &PocketIc,
218    canister_id: Principal,
219    sender: Option<Principal>,
220) -> Result<ControllerSnapshot, PocketIcCallFailure> {
221    let snapshot =
222        transport::capture_call(|| pocket_ic.take_canister_snapshot(canister_id, sender, None))?;
223    Ok(ControllerSnapshot {
224        snapshot_id: snapshot.id,
225        sender,
226    })
227}
228
229fn cleanup_captured_snapshots(
230    snapshots: BTreeMap<Principal, ControllerSnapshot>,
231    mut remove: impl FnMut(Principal, Option<Principal>, Vec<u8>) -> Result<(), RejectResponse>,
232) -> Vec<SnapshotCleanupFailure> {
233    let mut failures = Vec::new();
234    for (canister_id, snapshot) in snapshots {
235        let cleanup =
236            transport::capture_call(|| remove(canister_id, snapshot.sender, snapshot.snapshot_id));
237        if let Err(source) = cleanup {
238            failures.push(SnapshotCleanupFailure {
239                canister_id,
240                sender: snapshot.sender,
241                source: Box::new(source),
242            });
243        }
244    }
245    failures
246}
247
248fn restore_controller_snapshot(
249    pocket_ic: &PocketIc,
250    canister_id: Principal,
251    snapshot_id: &[u8],
252    funding: SnapshotRestoreFunding,
253    sender: Option<Principal>,
254) -> Result<(), ControllerSnapshotError> {
255    transport::capture_call(|| {
256        apply_snapshot_restore_funding(pocket_ic, canister_id, funding);
257        pocket_ic.load_canister_snapshot(canister_id, sender, snapshot_id.to_vec())
258    })
259    .map_err(|source| ControllerSnapshotError::RestoreFailed {
260        canister_id,
261        sender,
262        source: Box::new(source),
263    })
264}
265
266fn apply_snapshot_restore_funding(
267    pocket_ic: &PocketIc,
268    canister_id: Principal,
269    funding: SnapshotRestoreFunding,
270) {
271    match funding {
272        SnapshotRestoreFunding::Preserve => {}
273        SnapshotRestoreFunding::TopUpTo { minimum_cycles } => {
274            let top_up = minimum_cycles.saturating_sub(pocket_ic.cycle_balance(canister_id));
275            if top_up > 0 {
276                let _ = pocket_ic.add_cycles(canister_id, top_up);
277            }
278        }
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    #[test]
285    fn cleanup_preserves_causes_and_senders_and_continues_after_each_failure() {
286        use super::{
287            ControllerSnapshot, PocketIcCallFailure, SnapshotCleanupFailure,
288            cleanup_captured_snapshots,
289        };
290        use ic_platform::pocket_ic::{ErrorCode, RejectCode, RejectResponse};
291        use ic_platform::principal::ExternalPrincipal as Principal;
292
293        let ids = [1, 2, 3, 4].map(|byte| Principal::from_slice(&[byte]));
294        let sender = Some(Principal::anonymous());
295        let snapshots = ids
296            .map(|id| {
297                (
298                    id,
299                    ControllerSnapshot {
300                        snapshot_id: vec![7],
301                        sender,
302                    },
303                )
304            })
305            .into_iter()
306            .collect();
307        let rejection = RejectResponse {
308            reject_code: RejectCode::CanisterReject,
309            reject_message: "cleanup rejected".into(),
310            error_code: ErrorCode::CanisterCalledTrap,
311            certified: false,
312        };
313        let transport = "reqwest::Error { kind: Request, url: \"http://127.0.0.1:1234/instances/0/update/delete_canister_snapshot\", source: hyper::Error(IncompleteMessage) }";
314        let mut calls = Vec::new();
315        let failures = cleanup_captured_snapshots(snapshots, |id, selected_sender, snapshot_id| {
316            assert_eq!(snapshot_id, [7]);
317            calls.push((id, selected_sender));
318            if id == ids[1] {
319                return Err(rejection.clone());
320            }
321            if id == ids[2] {
322                std::panic::panic_any(transport.to_owned());
323            }
324            Ok(())
325        });
326        assert_eq!(calls, ids.map(|id| (id, sender)));
327        let [
328            SnapshotCleanupFailure {
329                canister_id: rejected_id,
330                sender: rejected_sender,
331                source: rejected,
332            },
333            SnapshotCleanupFailure {
334                canister_id: panicked_id,
335                sender: panicked_sender,
336                source: panicked,
337            },
338        ] = failures.as_slice()
339        else {
340            panic!("expected both cleanup failures: {failures:?}");
341        };
342        assert_eq!(*rejected_id, ids[1]);
343        assert_eq!(*rejected_sender, sender);
344        let PocketIcCallFailure::Rejected(response) = rejected.as_ref() else {
345            panic!("expected rejection: {rejected:?}");
346        };
347        assert_eq!(response, &rejection);
348        assert_eq!(*panicked_id, ids[2]);
349        assert_eq!(*panicked_sender, sender);
350        let PocketIcCallFailure::Operation(source) = panicked.as_ref() else {
351            panic!("expected operation failure: {panicked:?}");
352        };
353        assert!(source.is_transport());
354        assert_eq!(source.message(), transport);
355    }
356}
ic_testkit/pic/snapshot.rs

ic_testkit/pic/
snapshot.rs
