ic_testkit/pic/
snapshot.rs

1use std::collections::HashMap;
2
3use candid::Principal;
4
5use super::{ControllerSnapshots, Pic};
6
7const SNAPSHOT_RESTORE_MIN_CYCLES: u128 = 200_000_000_000_000;
8
9impl Pic {
10    /// Capture one restorable snapshot per canister using a shared controller.
11    pub fn capture_controller_snapshots<I>(
12        &self,
13        controller_id: Principal,
14        canister_ids: I,
15    ) -> Option<ControllerSnapshots>
16    where
17        I: IntoIterator<Item = Principal>,
18    {
19        let mut snapshots = HashMap::new();
20
21        for canister_id in canister_ids {
22            let Some(snapshot) = self.try_take_controller_snapshot(controller_id, canister_id)
23            else {
24                eprintln!(
25                    "capture_controller_snapshots: snapshot capture unavailable for {canister_id}"
26                );
27                return None;
28            };
29            snapshots.insert(canister_id, snapshot);
30        }
31
32        Some(ControllerSnapshots::new(snapshots))
33    }
34
35    /// Restore a previously captured snapshot set using the same controller.
36    pub fn restore_controller_snapshots(
37        &self,
38        controller_id: Principal,
39        snapshots: &ControllerSnapshots,
40    ) {
41        for (canister_id, snapshot_id, sender) in snapshots.iter() {
42            self.restore_controller_snapshot(controller_id, canister_id, sender, snapshot_id);
43        }
44    }
45
46    // Capture one snapshot with sender fallbacks that match controller ownership.
47    fn try_take_controller_snapshot(
48        &self,
49        controller_id: Principal,
50        canister_id: Principal,
51    ) -> Option<(Vec<u8>, Option<Principal>)> {
52        let candidates = controller_sender_candidates(controller_id, canister_id);
53        let mut last_err = None;
54
55        for sender in candidates {
56            match self.inner.take_canister_snapshot(canister_id, sender, None) {
57                Ok(snapshot) => return Some((snapshot.id, sender)),
58                Err(err) => last_err = Some((sender, err)),
59            }
60        }
61
62        if let Some((sender, err)) = last_err {
63            eprintln!(
64                "failed to capture canister snapshot for {canister_id} using sender {sender:?}: {err}"
65            );
66        }
67        None
68    }
69
70    // Restore one snapshot with sender fallbacks that match controller ownership.
71    fn restore_controller_snapshot(
72        &self,
73        controller_id: Principal,
74        canister_id: Principal,
75        snapshot_sender: Option<Principal>,
76        snapshot_id: &[u8],
77    ) {
78        let fallback_sender = if snapshot_sender.is_some() {
79            None
80        } else {
81            Some(controller_id)
82        };
83        let candidates = [snapshot_sender, fallback_sender];
84        let mut last_err = None;
85
86        for sender in candidates {
87            self.ensure_snapshot_restore_cycles(canister_id);
88            match self
89                .inner
90                .load_canister_snapshot(canister_id, sender, snapshot_id.to_vec())
91            {
92                Ok(()) => return,
93                Err(err) => last_err = Some((sender, err)),
94            }
95        }
96
97        let (sender, err) =
98            last_err.expect("snapshot restore must have at least one sender attempt");
99        panic!(
100            "failed to restore canister snapshot for {canister_id} using sender {sender:?}: {err}"
101        );
102    }
103
104    fn ensure_snapshot_restore_cycles(&self, canister_id: Principal) {
105        let balance = self.inner.cycle_balance(canister_id);
106        if balance < SNAPSHOT_RESTORE_MIN_CYCLES {
107            let top_up = SNAPSHOT_RESTORE_MIN_CYCLES - balance;
108            let _ = self.inner.add_cycles(canister_id, top_up);
109        }
110    }
111}
112
113// Prefer the likely controller sender first to reduce noisy management-call failures.
114fn controller_sender_candidates(
115    controller_id: Principal,
116    canister_id: Principal,
117) -> [Option<Principal>; 2] {
118    if canister_id == controller_id {
119        [None, Some(controller_id)]
120    } else {
121        [Some(controller_id), None]
122    }
123}
ic_testkit/pic/snapshot.rs

ic_testkit/pic/
snapshot.rs
