canic_testkit/pic/
snapshot.rs

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

canic_testkit/pic/
snapshot.rs
