ic_testkit/pic/
snapshot.rs

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

ic_testkit/pic/
snapshot.rs
