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

ic_testkit/pic/
snapshot.rs
