1use super::{
2 BackupOperation, BackupOperationKind, BackupPlan, BackupPlanError, BackupScopeKind,
3 BackupTarget, ControlAuthority, ControlAuthoritySource, QuiescencePolicy,
4 SnapshotReadAuthority,
5};
6use crate::{
7 discovery::{RegistryEntry, SnapshotTarget, targets_from_registry},
8 manifest::IdentityMode,
9};
10use candid::Principal;
11use std::{
12 collections::{BTreeMap, BTreeSet},
13 str::FromStr,
14};
15
16pub struct BackupPlanBuildInput<'a> {
21 pub plan_id: String,
22 pub run_id: String,
23 pub fleet: String,
24 pub network: String,
25 pub root_canister_id: String,
26 pub selected_canister_id: Option<String>,
27 pub selected_scope_kind: BackupScopeKind,
28 pub include_descendants: bool,
29 pub topology_hash_before_quiesce: String,
30 pub registry: &'a [RegistryEntry],
31 pub control_authority: ControlAuthority,
32 pub snapshot_read_authority: SnapshotReadAuthority,
33 pub quiescence_policy: QuiescencePolicy,
34 pub identity_mode: IdentityMode,
35}
36
37pub fn build_backup_plan(input: BackupPlanBuildInput<'_>) -> Result<BackupPlan, BackupPlanError> {
39 let snapshot_read_authority = input.snapshot_read_authority.clone();
40 let quiescence_policy = input.quiescence_policy.clone();
41 let root_included = input.selected_scope_kind == BackupScopeKind::MaintenanceRoot;
42 let selected_subtree_root = selected_subtree_root(&input)?;
43 let snapshot_targets = snapshot_targets(&input)?;
44 let target_depths = target_depths(input.registry);
45 let targets = snapshot_targets
46 .into_iter()
47 .map(|target| {
48 backup_target(
49 target,
50 &target_depths,
51 input.control_authority.clone(),
52 input.snapshot_read_authority.clone(),
53 input.identity_mode.clone(),
54 )
55 })
56 .collect::<Vec<_>>();
57 let phases = build_backup_phases(&targets);
58
59 let plan = BackupPlan {
60 plan_id: input.plan_id,
61 run_id: input.run_id,
62 fleet: input.fleet,
63 network: input.network,
64 root_canister_id: input.root_canister_id,
65 selected_subtree_root,
66 selected_scope_kind: input.selected_scope_kind,
67 include_descendants: input.include_descendants,
68 root_included,
69 requires_root_controller: input.control_authority.source
70 == ControlAuthoritySource::RootController,
71 snapshot_read_authority,
72 quiescence_policy,
73 topology_hash_before_quiesce: input.topology_hash_before_quiesce,
74 targets,
75 phases,
76 };
77 plan.validate()?;
78 Ok(plan)
79}
80
81pub fn resolve_backup_selector(
83 registry: &[RegistryEntry],
84 selector: &str,
85) -> Result<String, BackupPlanError> {
86 validate_nonempty("selector", selector)?;
87 if Principal::from_str(selector).is_ok() {
88 return registry
89 .iter()
90 .find(|entry| entry.pid == selector)
91 .map(|entry| entry.pid.clone())
92 .ok_or_else(|| BackupPlanError::UnknownSelector(selector.to_string()));
93 }
94
95 let matches = registry
96 .iter()
97 .filter(|entry| entry.role.as_deref() == Some(selector))
98 .map(|entry| entry.pid.clone())
99 .collect::<Vec<_>>();
100 match matches.as_slice() {
101 [canister] => Ok(canister.clone()),
102 [] => Err(BackupPlanError::UnknownSelector(selector.to_string())),
103 _ => Err(BackupPlanError::AmbiguousSelector {
104 selector: selector.to_string(),
105 matches,
106 }),
107 }
108}
109
110fn selected_subtree_root(
111 input: &BackupPlanBuildInput<'_>,
112) -> Result<Option<String>, BackupPlanError> {
113 match input.selected_scope_kind {
114 BackupScopeKind::NonRootFleet => Ok(None),
115 BackupScopeKind::Member | BackupScopeKind::Subtree | BackupScopeKind::MaintenanceRoot => {
116 input
117 .selected_canister_id
118 .clone()
119 .ok_or(BackupPlanError::EmptyField("selected_canister_id"))
120 .map(Some)
121 }
122 }
123}
124
125fn snapshot_targets(
126 input: &BackupPlanBuildInput<'_>,
127) -> Result<Vec<SnapshotTarget>, BackupPlanError> {
128 match input.selected_scope_kind {
129 BackupScopeKind::Member | BackupScopeKind::Subtree | BackupScopeKind::MaintenanceRoot => {
130 let selected = input
131 .selected_canister_id
132 .as_deref()
133 .ok_or(BackupPlanError::EmptyField("selected_canister_id"))?;
134 let recursive = input.selected_scope_kind == BackupScopeKind::Subtree
135 || input.include_descendants
136 || input.selected_scope_kind == BackupScopeKind::MaintenanceRoot;
137 targets_from_registry(input.registry, selected, recursive)
138 .map_err(BackupPlanError::from)
139 }
140 BackupScopeKind::NonRootFleet => Ok(input
141 .registry
142 .iter()
143 .filter(|entry| entry.pid != input.root_canister_id)
144 .map(|entry| SnapshotTarget {
145 canister_id: entry.pid.clone(),
146 role: entry.role.clone(),
147 parent_canister_id: entry.parent_pid.clone(),
148 })
149 .collect()),
150 }
151}
152
153fn backup_target(
154 target: SnapshotTarget,
155 target_depths: &BTreeMap<String, u32>,
156 control_authority: ControlAuthority,
157 snapshot_read_authority: SnapshotReadAuthority,
158 identity_mode: IdentityMode,
159) -> BackupTarget {
160 BackupTarget {
161 depth: target_depths
162 .get(&target.canister_id)
163 .copied()
164 .unwrap_or_default(),
165 canister_id: target.canister_id,
166 role: target.role,
167 parent_canister_id: target.parent_canister_id,
168 control_authority,
169 snapshot_read_authority,
170 identity_mode,
171 expected_module_hash: None,
172 }
173}
174
175fn target_depths(registry: &[RegistryEntry]) -> BTreeMap<String, u32> {
176 let parents = registry
177 .iter()
178 .map(|entry| (entry.pid.as_str(), entry.parent_pid.as_deref()))
179 .collect::<BTreeMap<_, _>>();
180 registry
181 .iter()
182 .map(|entry| {
183 (
184 entry.pid.clone(),
185 target_depth(entry.pid.as_str(), &parents),
186 )
187 })
188 .collect()
189}
190
191fn target_depth(canister_id: &str, parents: &BTreeMap<&str, Option<&str>>) -> u32 {
192 let mut depth = 0;
193 let mut current = canister_id;
194 let mut seen = BTreeSet::new();
195
196 while let Some(Some(parent)) = parents.get(current) {
197 if !seen.insert(current) {
198 break;
199 }
200 depth += 1;
201 current = parent;
202 }
203
204 depth
205}
206
207fn build_backup_phases(targets: &[BackupTarget]) -> Vec<BackupOperation> {
208 let mut phases = vec![
209 operation(
210 "validate-topology",
211 BackupOperationKind::ValidateTopology,
212 None,
213 ),
214 operation(
215 "validate-control-authority",
216 BackupOperationKind::ValidateControlAuthority,
217 None,
218 ),
219 operation(
220 "validate-snapshot-read-authority",
221 BackupOperationKind::ValidateSnapshotReadAuthority,
222 None,
223 ),
224 operation(
225 "validate-quiescence-policy",
226 BackupOperationKind::ValidateQuiescencePolicy,
227 None,
228 ),
229 ];
230
231 let mut top_down = targets.iter().collect::<Vec<_>>();
232 top_down.sort_by(|left, right| {
233 left.depth
234 .cmp(&right.depth)
235 .then_with(|| left.canister_id.cmp(&right.canister_id))
236 });
237 for target in &top_down {
238 phases.push(operation(
239 format!("stop-{}", target.canister_id),
240 BackupOperationKind::Stop,
241 Some(target.canister_id.clone()),
242 ));
243 }
244 for target in &top_down {
245 phases.push(operation(
246 format!("snapshot-{}", target.canister_id),
247 BackupOperationKind::CreateSnapshot,
248 Some(target.canister_id.clone()),
249 ));
250 }
251
252 let mut bottom_up = top_down;
253 bottom_up.reverse();
254 for target in &bottom_up {
255 phases.push(operation(
256 format!("start-{}", target.canister_id),
257 BackupOperationKind::Start,
258 Some(target.canister_id.clone()),
259 ));
260 }
261
262 for target in targets {
263 phases.push(operation(
264 format!("download-{}", target.canister_id),
265 BackupOperationKind::DownloadSnapshot,
266 Some(target.canister_id.clone()),
267 ));
268 phases.push(operation(
269 format!("verify-{}", target.canister_id),
270 BackupOperationKind::VerifyArtifact,
271 Some(target.canister_id.clone()),
272 ));
273 }
274 phases.push(operation(
275 "finalize-manifest",
276 BackupOperationKind::FinalizeManifest,
277 None,
278 ));
279
280 phases
281 .into_iter()
282 .enumerate()
283 .map(|(index, mut phase)| {
284 phase.order = u32::try_from(index).unwrap_or(u32::MAX);
285 phase
286 })
287 .collect()
288}
289
290fn operation(
291 operation_id: impl Into<String>,
292 kind: BackupOperationKind,
293 target_canister_id: Option<String>,
294) -> BackupOperation {
295 BackupOperation {
296 operation_id: operation_id.into(),
297 order: 0,
298 kind,
299 target_canister_id,
300 }
301}
302
303fn validate_nonempty(field: &'static str, value: &str) -> Result<(), BackupPlanError> {
304 if value.trim().is_empty() {
305 Err(BackupPlanError::EmptyField(field))
306 } else {
307 Ok(())
308 }
309}