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 mut cleanup_failures = cleanup_snapshot_ids(pocket_ic, canister_id, sender, pending);
213        cleanup_failures.extend(cleanup_captured_snapshots(
214            snapshots,
215            |canister_id, sender, snapshot_id| {
216                pocket_ic.delete_canister_snapshot(canister_id, sender, snapshot_id)
217            },
218        ));
219        return Err(match failure {
220            transfer::Failure::Call(failure) => ControllerSnapshotError::CaptureFailed {
221                canister_id,
222                sender,
223                source: Box::new(failure),
224                cleanup_failures,
225            },
226            transfer::Failure::Invalid(message) => ControllerSnapshotError::InvalidSnapshot {
227                canister_id,
228                message,
229                cleanup_failures,
230            },
231            transfer::Failure::Codec(source) => ControllerSnapshotError::SnapshotCodec {
232                canister_id,
233                source,
234                cleanup_failures,
235            },
236        });
237    }
238    Ok(ControllerSnapshots(snapshots))
239}
240
241impl ControllerSnapshots {
242    /// Return the number of captured canisters.
243    #[must_use]
244    pub fn len(&self) -> usize {
245        self.0.len()
246    }
247
248    /// Report whether the set contains no snapshots.
249    #[must_use]
250    pub fn is_empty(&self) -> bool {
251        self.0.is_empty()
252    }
253
254    /// Iterate over captured canister ids in deterministic principal order.
255    pub fn canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
256        self.0.keys().copied()
257    }
258
259    pub(super) fn iter(&self) -> impl Iterator<Item = (Principal, &[u8], Option<Principal>)> + '_ {
260        self.0.iter().map(|(canister_id, snapshot)| {
261            (
262                *canister_id,
263                snapshot.snapshot_id.as_slice(),
264                snapshot.sender,
265            )
266        })
267    }
268}
269
270fn try_take_snapshot(
271    pocket_ic: &PocketIc,
272    canister_id: Principal,
273    sender: Option<Principal>,
274    pending: &mut Vec<Vec<u8>>,
275) -> Result<ControllerSnapshot, transfer::Failure> {
276    let snapshot =
277        transport::capture_call(|| pocket_ic.take_canister_snapshot(canister_id, sender, None))?;
278    pending.push(snapshot.id.clone());
279    let uploaded = transfer::upload(pocket_ic, canister_id, sender, &snapshot.id, pending)?;
280    transport::capture_call(|| {
281        pocket_ic.delete_canister_snapshot(canister_id, sender, snapshot.id)
282    })?;
283    pending.clear();
284    Ok(ControllerSnapshot {
285        snapshot_id: uploaded,
286        sender,
287    })
288}
289
290fn cleanup_snapshot_ids(
291    pic: &PocketIc,
292    canister_id: Principal,
293    sender: Option<Principal>,
294    ids: Vec<Vec<u8>>,
295) -> Vec<SnapshotCleanupFailure> {
296    ids.into_iter()
297        .filter_map(|id| {
298            transport::capture_call(|| {
299                pic.delete_canister_snapshot(canister_id, sender, id.clone())
300            })
301            .err()
302            .map(|source| SnapshotCleanupFailure {
303                canister_id,
304                sender,
305                snapshot_id: id,
306                source: Box::new(source),
307            })
308        })
309        .collect()
310}
311
312fn cleanup_captured_snapshots(
313    snapshots: BTreeMap<Principal, ControllerSnapshot>,
314    mut remove: impl FnMut(Principal, Option<Principal>, Vec<u8>) -> Result<(), RejectResponse>,
315) -> Vec<SnapshotCleanupFailure> {
316    let mut failures = Vec::new();
317    for (canister_id, snapshot) in snapshots {
318        let cleanup = transport::capture_call(|| {
319            remove(canister_id, snapshot.sender, snapshot.snapshot_id.clone())
320        });
321        if let Err(source) = cleanup {
322            failures.push(SnapshotCleanupFailure {
323                canister_id,
324                sender: snapshot.sender,
325                snapshot_id: snapshot.snapshot_id,
326                source: Box::new(source),
327            });
328        }
329    }
330    failures
331}
332
333fn restore_controller_snapshot(
334    pocket_ic: &PocketIc,
335    canister_id: Principal,
336    snapshot_id: &[u8],
337    funding: SnapshotRestoreFunding,
338    sender: Option<Principal>,
339) -> Result<(), ControllerSnapshotError> {
340    transport::capture_call(|| {
341        apply_snapshot_restore_funding(pocket_ic, canister_id, funding);
342        pocket_ic.load_canister_snapshot(canister_id, sender, snapshot_id.to_vec())
343    })
344    .map_err(|source| ControllerSnapshotError::RestoreFailed {
345        canister_id,
346        sender,
347        source: Box::new(source),
348    })
349}
350
351fn apply_snapshot_restore_funding(
352    pocket_ic: &PocketIc,
353    canister_id: Principal,
354    funding: SnapshotRestoreFunding,
355) {
356    match funding {
357        SnapshotRestoreFunding::Preserve => {}
358        SnapshotRestoreFunding::TopUpTo { minimum_cycles } => {
359            let top_up = minimum_cycles.saturating_sub(pocket_ic.cycle_balance(canister_id));
360            if top_up > 0 {
361                let _ = pocket_ic.add_cycles(canister_id, top_up);
362            }
363        }
364    }
365}
366
367#[cfg(test)]
368mod tests {
369    #[test]
370    fn cleanup_preserves_causes_and_senders_and_continues_after_each_failure() {
371        use super::{
372            ControllerSnapshot, PocketIcCallFailure, SnapshotCleanupFailure,
373            cleanup_captured_snapshots,
374        };
375        use ic_platform::pocket_ic::{ErrorCode, RejectCode, RejectResponse};
376        use ic_platform::principal::ExternalPrincipal as Principal;
377
378        let ids = [1, 2, 3, 4].map(|byte| Principal::from_slice(&[byte]));
379        let sender = Some(Principal::anonymous());
380        let snapshots = ids
381            .map(|id| {
382                (
383                    id,
384                    ControllerSnapshot {
385                        snapshot_id: vec![7],
386                        sender,
387                    },
388                )
389            })
390            .into_iter()
391            .collect();
392        let rejection = RejectResponse {
393            reject_code: RejectCode::CanisterReject,
394            reject_message: "cleanup rejected".into(),
395            error_code: ErrorCode::CanisterCalledTrap,
396            certified: false,
397        };
398        let transport = "reqwest::Error { kind: Request, url: \"http://127.0.0.1:1234/instances/0/update/delete_canister_snapshot\", source: hyper::Error(IncompleteMessage) }";
399        let mut calls = Vec::new();
400        let failures = cleanup_captured_snapshots(snapshots, |id, selected_sender, snapshot_id| {
401            assert_eq!(snapshot_id, [7]);
402            calls.push((id, selected_sender));
403            if id == ids[1] {
404                return Err(rejection.clone());
405            }
406            if id == ids[2] {
407                std::panic::panic_any(transport.to_owned());
408            }
409            Ok(())
410        });
411        assert_eq!(calls, ids.map(|id| (id, sender)));
412        let [
413            SnapshotCleanupFailure {
414                canister_id: rejected_id,
415                sender: rejected_sender,
416                source: rejected,
417                snapshot_id: rejected_snapshot,
418            },
419            SnapshotCleanupFailure {
420                canister_id: panicked_id,
421                sender: panicked_sender,
422                source: panicked,
423                snapshot_id: panicked_snapshot,
424            },
425        ] = failures.as_slice()
426        else {
427            panic!("expected both cleanup failures: {failures:?}");
428        };
429        assert_eq!(*rejected_id, ids[1]);
430        assert_eq!(rejected_snapshot, &[7]);
431        assert_eq!(panicked_snapshot, &[7]);
432        assert_eq!(*rejected_sender, sender);
433        let PocketIcCallFailure::Rejected(response) = rejected.as_ref() else {
434            panic!("expected rejection: {rejected:?}");
435        };
436        assert_eq!(response, &rejection);
437        assert_eq!(*panicked_id, ids[2]);
438        assert_eq!(*panicked_sender, sender);
439        let PocketIcCallFailure::Operation(source) = panicked.as_ref() else {
440            panic!("expected operation failure: {panicked:?}");
441        };
442        assert!(source.is_transport());
443        assert_eq!(source.message(), transport);
444    }
445}
ic_testkit/pic/snapshot.rs

ic_testkit/pic/
snapshot.rs
