ic_testkit/pic/
snapshot.rs

1use std::collections::BTreeMap;
2
3use candid::Principal;
4use pocket_ic::{PocketIc, RejectResponse};
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    fn capture_snapshots<I>(
97        &self,
98        targets: I,
99    ) -> Result<ControllerSnapshots, ControllerSnapshotError>
100    where
101        I: IntoIterator<Item = CanisterSnapshotTarget>;
102
103    /// Restore each snapshot with its exact captured sender and selected funding.
104    ///
105    /// `Preserve` never adds cycles; `TopUpTo` is evaluated before each restore.
106    /// PocketIC may charge cycles for the restore itself. Rejection or panic stops
107    /// the set without retrying another sender or rewinding earlier restores.
108    fn restore_snapshots(
109        &self,
110        snapshots: &ControllerSnapshots,
111        funding: SnapshotRestoreFunding,
112    ) -> Result<(), ControllerSnapshotError>;
113}
114
115impl PocketIcSnapshotExt for PocketIc {
116    fn capture_snapshots<I>(
117        &self,
118        targets: I,
119    ) -> Result<ControllerSnapshots, ControllerSnapshotError>
120    where
121        I: IntoIterator<Item = CanisterSnapshotTarget>,
122    {
123        capture_snapshot_set(self, targets)
124    }
125
126    fn restore_snapshots(
127        &self,
128        snapshots: &ControllerSnapshots,
129        funding: SnapshotRestoreFunding,
130    ) -> Result<(), ControllerSnapshotError> {
131        for (canister_id, snapshot_id, sender) in snapshots.iter() {
132            restore_controller_snapshot(self, canister_id, snapshot_id, funding, sender)?;
133        }
134        Ok(())
135    }
136}
137
138fn capture_snapshot_set<I>(
139    pocket_ic: &PocketIc,
140    targets: I,
141) -> Result<ControllerSnapshots, ControllerSnapshotError>
142where
143    I: IntoIterator<Item = CanisterSnapshotTarget>,
144{
145    // Validate the entire input before issuing any management calls. The map
146    // also owns the deterministic canister order.
147    let mut ordered_targets = BTreeMap::new();
148    for target in targets {
149        if ordered_targets
150            .insert(target.canister_id, target.sender)
151            .is_some()
152        {
153            return Err(ControllerSnapshotError::DuplicateCanisterId {
154                canister_id: target.canister_id,
155            });
156        }
157    }
158    let mut snapshots = BTreeMap::new();
159    for (canister_id, sender) in ordered_targets {
160        let failure = match try_take_snapshot(pocket_ic, canister_id, sender) {
161            Ok(snapshot) => {
162                snapshots.insert(canister_id, snapshot);
163                continue;
164            }
165            Err(failure) => failure,
166        };
167        let cleanup_failures =
168            cleanup_captured_snapshots(snapshots, |canister_id, sender, snapshot_id| {
169                pocket_ic.delete_canister_snapshot(canister_id, sender, snapshot_id)
170            });
171        return Err(ControllerSnapshotError::CaptureFailed {
172            canister_id,
173            sender,
174            source: Box::new(failure),
175            cleanup_failures,
176        });
177    }
178    Ok(ControllerSnapshots(snapshots))
179}
180
181impl ControllerSnapshots {
182    /// Return the number of captured canisters.
183    #[must_use]
184    pub fn len(&self) -> usize {
185        self.0.len()
186    }
187
188    /// Report whether the set contains no snapshots.
189    #[must_use]
190    pub fn is_empty(&self) -> bool {
191        self.0.is_empty()
192    }
193
194    /// Iterate over captured canister ids in deterministic principal order.
195    pub fn canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
196        self.0.keys().copied()
197    }
198
199    pub(super) fn iter(&self) -> impl Iterator<Item = (Principal, &[u8], Option<Principal>)> + '_ {
200        self.0.iter().map(|(canister_id, snapshot)| {
201            (
202                *canister_id,
203                snapshot.snapshot_id.as_slice(),
204                snapshot.sender,
205            )
206        })
207    }
208}
209
210fn try_take_snapshot(
211    pocket_ic: &PocketIc,
212    canister_id: Principal,
213    sender: Option<Principal>,
214) -> Result<ControllerSnapshot, PocketIcCallFailure> {
215    let snapshot =
216        transport::capture_call(|| pocket_ic.take_canister_snapshot(canister_id, sender, None))?;
217    Ok(ControllerSnapshot {
218        snapshot_id: snapshot.id,
219        sender,
220    })
221}
222
223fn cleanup_captured_snapshots(
224    snapshots: BTreeMap<Principal, ControllerSnapshot>,
225    mut remove: impl FnMut(Principal, Option<Principal>, Vec<u8>) -> Result<(), RejectResponse>,
226) -> Vec<SnapshotCleanupFailure> {
227    let mut failures = Vec::new();
228    for (canister_id, snapshot) in snapshots {
229        let cleanup =
230            transport::capture_call(|| remove(canister_id, snapshot.sender, snapshot.snapshot_id));
231        if let Err(source) = cleanup {
232            failures.push(SnapshotCleanupFailure {
233                canister_id,
234                sender: snapshot.sender,
235                source: Box::new(source),
236            });
237        }
238    }
239    failures
240}
241
242fn restore_controller_snapshot(
243    pocket_ic: &PocketIc,
244    canister_id: Principal,
245    snapshot_id: &[u8],
246    funding: SnapshotRestoreFunding,
247    sender: Option<Principal>,
248) -> Result<(), ControllerSnapshotError> {
249    transport::capture_call(|| {
250        apply_snapshot_restore_funding(pocket_ic, canister_id, funding);
251        pocket_ic.load_canister_snapshot(canister_id, sender, snapshot_id.to_vec())
252    })
253    .map_err(|source| ControllerSnapshotError::RestoreFailed {
254        canister_id,
255        sender,
256        source: Box::new(source),
257    })
258}
259
260fn apply_snapshot_restore_funding(
261    pocket_ic: &PocketIc,
262    canister_id: Principal,
263    funding: SnapshotRestoreFunding,
264) {
265    match funding {
266        SnapshotRestoreFunding::Preserve => {}
267        SnapshotRestoreFunding::TopUpTo { minimum_cycles } => {
268            let top_up = minimum_cycles.saturating_sub(pocket_ic.cycle_balance(canister_id));
269            if top_up > 0 {
270                let _ = pocket_ic.add_cycles(canister_id, top_up);
271            }
272        }
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    #[test]
279    fn cleanup_preserves_causes_and_senders_and_continues_after_each_failure() {
280        use super::{
281            ControllerSnapshot, PocketIcCallFailure, SnapshotCleanupFailure,
282            cleanup_captured_snapshots,
283        };
284        use candid::Principal;
285        use pocket_ic::{ErrorCode, RejectCode, RejectResponse};
286
287        let ids = [1, 2, 3, 4].map(|byte| Principal::from_slice(&[byte]));
288        let sender = Some(Principal::anonymous());
289        let snapshots = ids
290            .map(|id| {
291                (
292                    id,
293                    ControllerSnapshot {
294                        snapshot_id: vec![7],
295                        sender,
296                    },
297                )
298            })
299            .into_iter()
300            .collect();
301        let rejection = RejectResponse {
302            reject_code: RejectCode::CanisterReject,
303            reject_message: "cleanup rejected".into(),
304            error_code: ErrorCode::CanisterCalledTrap,
305            certified: false,
306        };
307        let transport = "reqwest::Error { kind: Request, url: \"http://127.0.0.1:1234/instances/0/update/delete_canister_snapshot\", source: hyper::Error(IncompleteMessage) }";
308        let mut calls = Vec::new();
309        let failures = cleanup_captured_snapshots(snapshots, |id, selected_sender, snapshot_id| {
310            assert_eq!(snapshot_id, [7]);
311            calls.push((id, selected_sender));
312            if id == ids[1] {
313                return Err(rejection.clone());
314            }
315            if id == ids[2] {
316                std::panic::panic_any(transport.to_owned());
317            }
318            Ok(())
319        });
320        assert_eq!(calls, ids.map(|id| (id, sender)));
321        let [
322            SnapshotCleanupFailure {
323                canister_id: rejected_id,
324                sender: rejected_sender,
325                source: rejected,
326            },
327            SnapshotCleanupFailure {
328                canister_id: panicked_id,
329                sender: panicked_sender,
330                source: panicked,
331            },
332        ] = failures.as_slice()
333        else {
334            panic!("expected both cleanup failures: {failures:?}");
335        };
336        assert_eq!(*rejected_id, ids[1]);
337        assert_eq!(*rejected_sender, sender);
338        let PocketIcCallFailure::Rejected(response) = rejected.as_ref() else {
339            panic!("expected rejection: {rejected:?}");
340        };
341        assert_eq!(response, &rejection);
342        assert_eq!(*panicked_id, ids[2]);
343        assert_eq!(*panicked_sender, sender);
344        let PocketIcCallFailure::Operation(source) = panicked.as_ref() else {
345            panic!("expected operation failure: {panicked:?}");
346        };
347        assert!(source.is_transport());
348        assert_eq!(source.message(), transport);
349    }
350}
ic_testkit/pic/snapshot.rs

ic_testkit/pic/
snapshot.rs
