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 module_hash: entry.module_hash.clone(),
149 })
150 .collect()),
151 }
152}
153
154fn backup_target(
155 target: SnapshotTarget,
156 target_depths: &BTreeMap<String, u32>,
157 control_authority: ControlAuthority,
158 snapshot_read_authority: SnapshotReadAuthority,
159 identity_mode: IdentityMode,
160) -> BackupTarget {
161 BackupTarget {
162 depth: target_depths
163 .get(&target.canister_id)
164 .copied()
165 .unwrap_or_default(),
166 canister_id: target.canister_id,
167 role: target.role,
168 parent_canister_id: target.parent_canister_id,
169 control_authority,
170 snapshot_read_authority,
171 identity_mode,
172 expected_module_hash: target.module_hash,
173 }
174}
175
176fn target_depths(registry: &[RegistryEntry]) -> BTreeMap<String, u32> {
177 let parents = registry
178 .iter()
179 .map(|entry| (entry.pid.as_str(), entry.parent_pid.as_deref()))
180 .collect::<BTreeMap<_, _>>();
181 registry
182 .iter()
183 .map(|entry| {
184 (
185 entry.pid.clone(),
186 target_depth(entry.pid.as_str(), &parents),
187 )
188 })
189 .collect()
190}
191
192fn target_depth(canister_id: &str, parents: &BTreeMap<&str, Option<&str>>) -> u32 {
193 let mut depth = 0;
194 let mut current = canister_id;
195 let mut seen = BTreeSet::new();
196
197 while let Some(Some(parent)) = parents.get(current) {
198 if !seen.insert(current) {
199 break;
200 }
201 depth += 1;
202 current = parent;
203 }
204
205 depth
206}
207
208fn build_backup_phases(targets: &[BackupTarget]) -> Vec<BackupOperation> {
209 let mut phases = vec![
210 operation(
211 "validate-topology",
212 BackupOperationKind::ValidateTopology,
213 None,
214 ),
215 operation(
216 "validate-control-authority",
217 BackupOperationKind::ValidateControlAuthority,
218 None,
219 ),
220 operation(
221 "validate-snapshot-read-authority",
222 BackupOperationKind::ValidateSnapshotReadAuthority,
223 None,
224 ),
225 operation(
226 "validate-quiescence-policy",
227 BackupOperationKind::ValidateQuiescencePolicy,
228 None,
229 ),
230 ];
231
232 let mut top_down = targets.iter().collect::<Vec<_>>();
233 top_down.sort_by(|left, right| {
234 left.depth
235 .cmp(&right.depth)
236 .then_with(|| left.canister_id.cmp(&right.canister_id))
237 });
238 for target in &top_down {
239 phases.push(operation(
240 format!("stop-{}", target.canister_id),
241 BackupOperationKind::Stop,
242 Some(target.canister_id.clone()),
243 ));
244 }
245 for target in &top_down {
246 phases.push(operation(
247 format!("snapshot-{}", target.canister_id),
248 BackupOperationKind::CreateSnapshot,
249 Some(target.canister_id.clone()),
250 ));
251 }
252
253 let mut bottom_up = top_down;
254 bottom_up.reverse();
255 for target in &bottom_up {
256 phases.push(operation(
257 format!("start-{}", target.canister_id),
258 BackupOperationKind::Start,
259 Some(target.canister_id.clone()),
260 ));
261 }
262
263 for target in targets {
264 phases.push(operation(
265 format!("download-{}", target.canister_id),
266 BackupOperationKind::DownloadSnapshot,
267 Some(target.canister_id.clone()),
268 ));
269 phases.push(operation(
270 format!("verify-{}", target.canister_id),
271 BackupOperationKind::VerifyArtifact,
272 Some(target.canister_id.clone()),
273 ));
274 }
275 phases.push(operation(
276 "finalize-manifest",
277 BackupOperationKind::FinalizeManifest,
278 None,
279 ));
280
281 phases
282 .into_iter()
283 .enumerate()
284 .map(|(index, mut phase)| {
285 phase.order = u32::try_from(index).unwrap_or(u32::MAX);
286 phase
287 })
288 .collect()
289}
290
291fn operation(
292 operation_id: impl Into<String>,
293 kind: BackupOperationKind,
294 target_canister_id: Option<String>,
295) -> BackupOperation {
296 BackupOperation {
297 operation_id: operation_id.into(),
298 order: 0,
299 kind,
300 target_canister_id,
301 }
302}
303
304fn validate_nonempty(field: &'static str, value: &str) -> Result<(), BackupPlanError> {
305 if value.trim().is_empty() {
306 Err(BackupPlanError::EmptyField(field))
307 } else {
308 Ok(())
309 }
310}