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
8mod transfer;
9
10#[derive(Clone, Debug, Eq, PartialEq)]
11struct ControllerSnapshot {
12    snapshot_id: Vec<u8>,
13    sender: Option<Principal>,
14}
15
16/// Deterministically ordered snapshots retaining each successful capture sender.
17#[derive(Clone, Debug, Eq, PartialEq)]
18pub struct ControllerSnapshots(BTreeMap<Principal, ControllerSnapshot>);
19
20/// One canister and exact management-call sender used for snapshot capture.
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
22pub struct CanisterSnapshotTarget {
23    /// Canister whose snapshot will be captured.
24    pub canister_id: Principal,
25    /// Exact management-call sender; None selects `PocketIC`'s default.
26    pub sender: Option<Principal>,
27}
28
29/// Failure to remove a snapshot while rolling back a partial capture.
30#[non_exhaustive]
31#[derive(Debug)]
32pub struct SnapshotCleanupFailure {
33    /// Canister whose snapshot could not be removed.
34    pub canister_id: Principal,
35    /// Exact sender selected for cleanup.
36    pub sender: Option<Principal>,
37    /// Exact known snapshot whose deletion failed; no other ID was substituted.
38    pub snapshot_id: Vec<u8>,
39    /// Complete rejection or operation failure, retaining transport classification.
40    pub source: Box<PocketIcCallFailure>,
41}
42
43/// Caller-selected cycle funding applied immediately before snapshot restore.
44#[non_exhaustive]
45#[derive(Clone, Copy, Debug, Eq, PartialEq)]
46pub enum SnapshotRestoreFunding {
47    /// Do not add cycles before the restore operation.
48    Preserve,
49    /// Add cycles only when needed to reach the given minimum balance.
50    TopUpTo {
51        /// Minimum balance established immediately before each restore.
52        minimum_cycles: u128,
53    },
54}
55
56/// Structured controller-snapshot failure.
57#[non_exhaustive]
58#[derive(Debug, thiserror::Error)]
59pub enum ControllerSnapshotError {
60    /// A cached baseline requires at least one canister.
61    #[error("a cached baseline requires at least one canister")]
62    EmptyCanisterSet,
63    /// Input contained the same canister more than once.
64    #[error("duplicate canister id in snapshot set: {canister_id}")]
65    DuplicateCanisterId {
66        /// Repeated canister id.
67        canister_id: Principal,
68    },
69    /// Snapshot transfer metadata or complete data admission failed.
70    #[error("invalid snapshot for {canister_id}: {message}; {} partial snapshots could not be cleaned up", .cleanup_failures.len())]
71    InvalidSnapshot {
72        /// Exact canister selected for capture.
73        canister_id: Principal,
74        /// Refused snapshot property, excluding source bytes.
75        message: &'static str,
76        /// Known snapshot deletions that failed during rollback.
77        cleanup_failures: Vec<SnapshotCleanupFailure>,
78    },
79    /// The external management snapshot message could not be encoded or decoded.
80    #[error("snapshot codec failed for {canister_id}: {source}")]
81    SnapshotCodec {
82        /// Exact canister selected for capture.
83        canister_id: Principal,
84        /// Original external Candid codec failure.
85        source: ic_platform::candid::contract::Error,
86        /// Known snapshot deletions that failed during rollback.
87        cleanup_failures: Vec<SnapshotCleanupFailure>,
88    },
89    /// Snapshot capture failed with the exact selected sender.
90    #[error("failed to capture snapshot for {canister_id} with sender {sender:?}: {source}; {} partial snapshots could not be cleaned up", .cleanup_failures.len())]
91    CaptureFailed {
92        /// Canister whose capture failed.
93        canister_id: Principal,
94        /// Exact sender selected for capture.
95        sender: Option<Principal>,
96        /// Complete rejection or operation failure.
97        source: Box<PocketIcCallFailure>,
98        /// Failures while deleting snapshots captured earlier in the set.
99        cleanup_failures: Vec<SnapshotCleanupFailure>,
100    },
101    /// Snapshot restore failed with the retained capture sender.
102    #[error("failed to restore snapshot for {canister_id} with sender {sender:?}: {source}")]
103    RestoreFailed {
104        /// Canister whose restore failed.
105        canister_id: Principal,
106        /// Exact sender selected for restore.
107        sender: Option<Principal>,
108        /// Complete rejection or operation failure.
109        source: Box<PocketIcCallFailure>,
110    },
111}
112
113/// Exact-sender capture and restore of related canister snapshots.
114pub trait PocketIcSnapshotExt {
115    /// Capture snapshots with one explicit sender per canister.
116    ///
117    /// The caller must stop and drain the complete target set before capture,
118    /// then restart it only after the complete set succeeds. This helper does
119    /// not stop/start canisters or rewind simulated time. Disable automatic
120    /// progress and exclude concurrent instance operations throughout the set.
121    /// Input is duplicate-checked before any effects and captured in deterministic
122    /// canister-id order. A later failure rolls back snapshots captured earlier
123    /// in the set, including known temporary transfer snapshots. No other sender
124    /// is attempted after a rejection. Upload semantics preserve the captured
125    /// global timer and low-memory hook alongside the module, globals, heap,
126    /// stable memory, certified data and chunk store. Transfer uses bounded
127    /// management messages; total work follows the captured state size.
128    ///
129    /// # Errors
130    /// Returns a target-admission or snapshot-capture error, retaining completed work.
131    fn capture_snapshots<I>(
132        &self,
133        targets: I,
134    ) -> Result<ControllerSnapshots, ControllerSnapshotError>
135    where
136        I: IntoIterator<Item = CanisterSnapshotTarget>;
137
138    /// Restore each snapshot with its exact captured sender and selected funding.
139    ///
140    /// The caller must stop/drain the entire captured set before restoring and
141    /// restart only after every restore and non-snapshot reset succeeds. Automatic
142    /// progress and concurrent instance operations must remain disabled. Absolute
143    /// timer deadlines are restored, not shifted; an expired timer may run when
144    /// its canister restarts. Simulated time and lifecycle remain caller-owned.
145    ///
146    /// `Preserve` never adds cycles; `TopUpTo` is evaluated before each restore.
147    /// `PocketIC` may charge cycles for the restore itself. Rejection or panic stops
148    /// the set without retrying another sender or rewinding earlier restores.
149    ///
150    /// # Errors
151    /// Returns a funding or restoration error, retaining partial progress.
152    fn restore_snapshots(
153        &self,
154        snapshots: &ControllerSnapshots,
155        funding: SnapshotRestoreFunding,
156    ) -> Result<(), ControllerSnapshotError>;
157}
158
159impl PocketIcSnapshotExt for PocketIc {
160    fn capture_snapshots<I>(
161        &self,
162        targets: I,
163    ) -> Result<ControllerSnapshots, ControllerSnapshotError>
164    where
165        I: IntoIterator<Item = CanisterSnapshotTarget>,
166    {
167        capture_snapshot_set(self, targets)
168    }
169
170    fn restore_snapshots(
171        &self,
172        snapshots: &ControllerSnapshots,
173        funding: SnapshotRestoreFunding,
174    ) -> Result<(), ControllerSnapshotError> {
175        for (canister_id, snapshot_id, sender) in snapshots.iter() {
176            restore_controller_snapshot(self, canister_id, snapshot_id, funding, sender)?;
177        }
178        Ok(())
179    }
180}
181
182fn capture_snapshot_set<I>(
183    pocket_ic: &PocketIc,
184    targets: I,
185) -> Result<ControllerSnapshots, ControllerSnapshotError>
186where
187    I: IntoIterator<Item = CanisterSnapshotTarget>,
188{
189    // Validate the entire input before issuing any management calls. The map
190    // also owns the deterministic canister order.
191    let mut ordered_targets = BTreeMap::new();
192    for target in targets {
193        if ordered_targets
194            .insert(target.canister_id, target.sender)
195            .is_some()
196        {
197            return Err(ControllerSnapshotError::DuplicateCanisterId {
198                canister_id: target.canister_id,
199            });
200        }
201    }
202    let mut snapshots = BTreeMap::new();
203    for (canister_id, sender) in ordered_targets {
204        let mut pending = Vec::new();
205        let failure = match try_take_snapshot(pocket_ic, canister_id, sender, &mut pending) {
206            Ok(snapshot) => {
207                snapshots.insert(canister_id, snapshot);
208                continue;
209            }
210            Err(failure) => failure,
211        };
212        let cleanup_failures = cleanup_captured_snapshots(
213            pending
214                .into_iter()
215                .map(|snapshot_id| {
216                    (
217                        canister_id,
218                        ControllerSnapshot {
219                            snapshot_id,
220                            sender,
221                        },
222                    )
223                })
224                .chain(snapshots),
225            |canister_id, sender, snapshot_id| {
226                pocket_ic.delete_canister_snapshot(canister_id, sender, snapshot_id)
227            },
228        );
229        return Err(match failure {
230            transfer::Failure::Call(failure) => ControllerSnapshotError::CaptureFailed {
231                canister_id,
232                sender,
233                source: Box::new(failure),
234                cleanup_failures,
235            },
236            transfer::Failure::Invalid(message) => ControllerSnapshotError::InvalidSnapshot {
237                canister_id,
238                message,
239                cleanup_failures,
240            },
241            transfer::Failure::Codec(source) => ControllerSnapshotError::SnapshotCodec {
242                canister_id,
243                source,
244                cleanup_failures,
245            },
246        });
247    }
248    Ok(ControllerSnapshots(snapshots))
249}
250
251impl ControllerSnapshots {
252    /// Return the number of captured canisters.
253    #[must_use]
254    pub fn len(&self) -> usize {
255        self.0.len()
256    }
257
258    /// Report whether the set contains no snapshots.
259    #[must_use]
260    pub fn is_empty(&self) -> bool {
261        self.0.is_empty()
262    }
263
264    /// Iterate over captured canister ids in deterministic principal order.
265    pub fn canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
266        self.0.keys().copied()
267    }
268
269    pub(super) fn iter(&self) -> impl Iterator<Item = (Principal, &[u8], Option<Principal>)> + '_ {
270        self.0.iter().map(|(canister_id, snapshot)| {
271            (
272                *canister_id,
273                snapshot.snapshot_id.as_slice(),
274                snapshot.sender,
275            )
276        })
277    }
278}
279
280fn try_take_snapshot(
281    pocket_ic: &PocketIc,
282    canister_id: Principal,
283    sender: Option<Principal>,
284    pending: &mut Vec<Vec<u8>>,
285) -> Result<ControllerSnapshot, transfer::Failure> {
286    let snapshot =
287        transport::capture_call(|| pocket_ic.take_canister_snapshot(canister_id, sender, None))?;
288    pending.push(snapshot.id.clone());
289    let uploaded = transfer::upload(pocket_ic, canister_id, sender, &snapshot.id, pending)?;
290    transport::capture_call(|| {
291        pocket_ic.delete_canister_snapshot(canister_id, sender, snapshot.id)
292    })?;
293    pending.clear();
294    Ok(ControllerSnapshot {
295        snapshot_id: uploaded,
296        sender,
297    })
298}
299
300fn cleanup_captured_snapshots(
301    snapshots: impl IntoIterator<Item = (Principal, ControllerSnapshot)>,
302    mut remove: impl FnMut(Principal, Option<Principal>, Vec<u8>) -> Result<(), RejectResponse>,
303) -> Vec<SnapshotCleanupFailure> {
304    let mut failures = Vec::new();
305    for (canister_id, snapshot) in snapshots {
306        let cleanup = transport::capture_call(|| {
307            remove(canister_id, snapshot.sender, snapshot.snapshot_id.clone())
308        });
309        if let Err(source) = cleanup {
310            failures.push(SnapshotCleanupFailure {
311                canister_id,
312                sender: snapshot.sender,
313                snapshot_id: snapshot.snapshot_id,
314                source: Box::new(source),
315            });
316        }
317    }
318    failures
319}
320
321fn restore_controller_snapshot(
322    pocket_ic: &PocketIc,
323    canister_id: Principal,
324    snapshot_id: &[u8],
325    funding: SnapshotRestoreFunding,
326    sender: Option<Principal>,
327) -> Result<(), ControllerSnapshotError> {
328    transport::capture_call(|| {
329        apply_snapshot_restore_funding(pocket_ic, canister_id, funding);
330        pocket_ic.load_canister_snapshot(canister_id, sender, snapshot_id.to_vec())
331    })
332    .map_err(|source| ControllerSnapshotError::RestoreFailed {
333        canister_id,
334        sender,
335        source: Box::new(source),
336    })
337}
338
339fn apply_snapshot_restore_funding(
340    pocket_ic: &PocketIc,
341    canister_id: Principal,
342    funding: SnapshotRestoreFunding,
343) {
344    match funding {
345        SnapshotRestoreFunding::Preserve => {}
346        SnapshotRestoreFunding::TopUpTo { minimum_cycles } => {
347            let top_up = minimum_cycles.saturating_sub(pocket_ic.cycle_balance(canister_id));
348            if top_up > 0 {
349                let _ = pocket_ic.add_cycles(canister_id, top_up);
350            }
351        }
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    #[test]
358    fn cleanup_preserves_causes_and_senders_and_continues_after_each_failure() {
359        use super::{
360            ControllerSnapshot, PocketIcCallFailure, SnapshotCleanupFailure,
361            cleanup_captured_snapshots,
362        };
363        use ic_platform::pocket_ic::{ErrorCode, RejectCode, RejectResponse};
364        use ic_platform::principal::ExternalPrincipal as Principal;
365
366        let ids = [1, 2, 3, 4].map(|byte| Principal::from_slice(&[byte]));
367        let sender = Some(Principal::anonymous());
368        let snapshots = ids
369            .map(|id| {
370                (
371                    id,
372                    ControllerSnapshot {
373                        snapshot_id: vec![7],
374                        sender,
375                    },
376                )
377            })
378            .into_iter()
379            .collect::<std::collections::BTreeMap<_, _>>();
380        let rejection = RejectResponse {
381            reject_code: RejectCode::CanisterReject,
382            reject_message: "cleanup rejected".into(),
383            error_code: ErrorCode::CanisterCalledTrap,
384            certified: false,
385        };
386        let transport = "reqwest::Error { kind: Request, url: \"http://127.0.0.1:1234/instances/0/update/delete_canister_snapshot\", source: hyper::Error(IncompleteMessage) }";
387        let mut calls = Vec::new();
388        let failures = cleanup_captured_snapshots(snapshots, |id, selected_sender, snapshot_id| {
389            assert_eq!(snapshot_id, [7]);
390            calls.push((id, selected_sender));
391            if id == ids[1] {
392                return Err(rejection.clone());
393            }
394            if id == ids[2] {
395                std::panic::panic_any(transport.to_owned());
396            }
397            Ok(())
398        });
399        assert_eq!(calls, ids.map(|id| (id, sender)));
400        let [
401            SnapshotCleanupFailure {
402                canister_id: rejected_id,
403                sender: rejected_sender,
404                source: rejected,
405                snapshot_id: rejected_snapshot,
406            },
407            SnapshotCleanupFailure {
408                canister_id: panicked_id,
409                sender: panicked_sender,
410                source: panicked,
411                snapshot_id: panicked_snapshot,
412            },
413        ] = failures.as_slice()
414        else {
415            panic!("expected both cleanup failures: {failures:?}");
416        };
417        assert_eq!(*rejected_id, ids[1]);
418        assert_eq!(rejected_snapshot, &[7]);
419        assert_eq!(panicked_snapshot, &[7]);
420        assert_eq!(*rejected_sender, sender);
421        let PocketIcCallFailure::Rejected(response) = rejected.as_ref() else {
422            panic!("expected rejection: {rejected:?}");
423        };
424        assert_eq!(response, &rejection);
425        assert_eq!(*panicked_id, ids[2]);
426        assert_eq!(*panicked_sender, sender);
427        let PocketIcCallFailure::Operation(source) = panicked.as_ref() else {
428            panic!("expected operation failure: {panicked:?}");
429        };
430        assert!(source.is_transport());
431        assert_eq!(source.message(), transport);
432    }
433}
ic_testkit/pic/snapshot.rs

ic_testkit/pic/
snapshot.rs
