1use crate::{
2 icp::{IcpCli, IcpCommandError},
3 install_root::{
4 InstallState, read_named_fleet_install_state, read_named_fleet_install_state_from_root,
5 },
6 registry::{RegistryEntry, RegistryParseError, parse_registry_entries},
7 replica_query,
8};
9use std::{collections::BTreeMap, path::Path};
10use thiserror::Error as ThisError;
11
12#[derive(Clone, Debug, Eq, PartialEq)]
17pub struct InstalledFleetRequest {
18 pub fleet: String,
19 pub network: String,
20 pub icp: String,
21 pub detect_lost_local_root: bool,
22}
23
24#[derive(Clone, Debug, Eq, PartialEq)]
29pub struct InstalledFleetResolution {
30 pub source: InstalledFleetSource,
31 pub state: InstallState,
32 pub registry: InstalledFleetRegistry,
33 pub topology: ResolvedFleetTopology,
34}
35
36#[derive(Clone, Copy, Debug, Eq, PartialEq)]
41pub enum InstalledFleetSource {
42 LocalReplica,
43 IcpCli,
44}
45
46#[derive(Clone, Debug, Eq, PartialEq)]
51pub struct InstalledFleetRegistry {
52 pub root_canister_id: String,
53 pub entries: Vec<RegistryEntry>,
54}
55
56#[derive(Clone, Debug, Eq, PartialEq)]
61pub struct ResolvedFleetTopology {
62 pub root_canister_id: String,
63 pub children_by_parent: BTreeMap<Option<String>, Vec<String>>,
64 pub roles_by_canister: BTreeMap<String, String>,
65}
66
67#[derive(Debug, ThisError)]
72pub enum InstalledFleetError {
73 #[error("fleet {fleet} is not installed on network {network}")]
74 NoInstalledFleet { network: String, fleet: String },
75
76 #[error("failed to read canic fleet state: {0}")]
77 InstallState(String),
78
79 #[error("local replica query failed: {0}")]
80 ReplicaQuery(String),
81
82 #[error("icp command failed: {command}\n{stderr}")]
83 IcpFailed { command: String, stderr: String },
84
85 #[error(
86 "fleet {fleet} points to root {root}, but that canister is not present on network {network}"
87 )]
88 LostLocalFleet {
89 fleet: String,
90 network: String,
91 root: String,
92 },
93
94 #[error(transparent)]
95 Registry(#[from] RegistryParseError),
96
97 #[error(transparent)]
98 Io(#[from] std::io::Error),
99}
100
101pub fn resolve_installed_fleet(
102 request: &InstalledFleetRequest,
103) -> Result<InstalledFleetResolution, InstalledFleetError> {
104 let state = read_installed_fleet_state(&request.network, &request.fleet)?;
105 let (source, registry_json) = query_registry(request, &state.root_canister_id)?;
106 installed_fleet_resolution(request, state, source, registry_json)
107}
108
109pub fn resolve_installed_fleet_from_root(
110 request: &InstalledFleetRequest,
111 icp_root: &Path,
112) -> Result<InstalledFleetResolution, InstalledFleetError> {
113 let state = read_installed_fleet_state_from_root(&request.network, &request.fleet, icp_root)?;
114 let (source, registry_json) =
115 query_registry_from_root(request, &state.root_canister_id, icp_root)?;
116 installed_fleet_resolution(request, state, source, registry_json)
117}
118
119fn installed_fleet_resolution(
120 _request: &InstalledFleetRequest,
121 state: InstallState,
122 source: InstalledFleetSource,
123 registry_json: String,
124) -> Result<InstalledFleetResolution, InstalledFleetError> {
125 let entries = parse_registry_entries(®istry_json)?;
126 let registry = InstalledFleetRegistry {
127 root_canister_id: state.root_canister_id.clone(),
128 entries,
129 };
130 let topology = ResolvedFleetTopology::from_registry(®istry);
131 Ok(InstalledFleetResolution {
132 source,
133 state,
134 registry,
135 topology,
136 })
137}
138
139pub fn read_installed_fleet_state(
140 network: &str,
141 fleet: &str,
142) -> Result<InstallState, InstalledFleetError> {
143 read_named_fleet_install_state(network, fleet)
144 .map_err(|err| InstalledFleetError::InstallState(err.to_string()))?
145 .ok_or_else(|| InstalledFleetError::NoInstalledFleet {
146 network: network.to_string(),
147 fleet: fleet.to_string(),
148 })
149}
150
151pub fn read_installed_fleet_state_from_root(
152 network: &str,
153 fleet: &str,
154 icp_root: &Path,
155) -> Result<InstallState, InstalledFleetError> {
156 read_named_fleet_install_state_from_root(icp_root, network, fleet)
157 .map_err(|err| InstalledFleetError::InstallState(err.to_string()))?
158 .ok_or_else(|| InstalledFleetError::NoInstalledFleet {
159 network: network.to_string(),
160 fleet: fleet.to_string(),
161 })
162}
163
164impl ResolvedFleetTopology {
165 fn from_registry(registry: &InstalledFleetRegistry) -> Self {
166 let mut children_by_parent = BTreeMap::<Option<String>, Vec<String>>::new();
167 let mut roles_by_canister = BTreeMap::new();
168 for entry in ®istry.entries {
169 children_by_parent
170 .entry(entry.parent_pid.clone())
171 .or_default()
172 .push(entry.pid.clone());
173 if let Some(role) = &entry.role {
174 roles_by_canister.insert(entry.pid.clone(), role.clone());
175 }
176 }
177 for children in children_by_parent.values_mut() {
178 children.sort();
179 }
180 Self {
181 root_canister_id: registry.root_canister_id.clone(),
182 children_by_parent,
183 roles_by_canister,
184 }
185 }
186}
187
188fn query_registry(
189 request: &InstalledFleetRequest,
190 root: &str,
191) -> Result<(InstalledFleetSource, String), InstalledFleetError> {
192 if replica_query::should_use_local_replica_query(Some(&request.network)) {
193 return replica_query::query_subnet_registry_json(Some(&request.network), root)
194 .map(|registry| (InstalledFleetSource::LocalReplica, registry))
195 .map_err(|err| local_registry_error(request, root, err.to_string()));
196 }
197
198 IcpCli::new(&request.icp, None, Some(request.network.clone()))
199 .canister_query_output(root, "canic_subnet_registry", Some("json"))
200 .map(|registry| (InstalledFleetSource::IcpCli, registry))
201 .map_err(installed_fleet_icp_error)
202}
203
204fn query_registry_from_root(
205 request: &InstalledFleetRequest,
206 root: &str,
207 icp_root: &Path,
208) -> Result<(InstalledFleetSource, String), InstalledFleetError> {
209 if replica_query::should_use_local_replica_query(Some(&request.network)) {
210 return replica_query::query_subnet_registry_json_from_root(
211 Some(&request.network),
212 root,
213 icp_root,
214 )
215 .map(|registry| (InstalledFleetSource::LocalReplica, registry))
216 .map_err(|err| local_registry_error(request, root, err.to_string()));
217 }
218
219 IcpCli::new(&request.icp, None, Some(request.network.clone()))
220 .with_cwd(icp_root)
221 .canister_query_output(root, "canic_subnet_registry", Some("json"))
222 .map(|registry| (InstalledFleetSource::IcpCli, registry))
223 .map_err(installed_fleet_icp_error)
224}
225
226fn local_registry_error(
227 request: &InstalledFleetRequest,
228 root: &str,
229 error: String,
230) -> InstalledFleetError {
231 if request.detect_lost_local_root && is_canister_not_found_error(&error) {
232 return InstalledFleetError::LostLocalFleet {
233 fleet: request.fleet.clone(),
234 network: request.network.clone(),
235 root: root.to_string(),
236 };
237 }
238 InstalledFleetError::ReplicaQuery(error)
239}
240
241fn is_canister_not_found_error(error: &str) -> bool {
242 error.contains("Canister ") && error.contains(" not found")
243}
244
245fn installed_fleet_icp_error(error: IcpCommandError) -> InstalledFleetError {
246 match error {
247 IcpCommandError::Io(err) => InstalledFleetError::Io(err),
248 IcpCommandError::Failed { command, stderr } => {
249 InstalledFleetError::IcpFailed { command, stderr }
250 }
251 IcpCommandError::Json {
252 command, output, ..
253 } => InstalledFleetError::IcpFailed {
254 command,
255 stderr: output,
256 },
257 IcpCommandError::SnapshotIdUnavailable { output } => InstalledFleetError::IcpFailed {
258 command: "icp canister snapshot create".to_string(),
259 stderr: output,
260 },
261 }
262}
263
264#[cfg(test)]
265mod tests {
266 use super::*;
267
268 #[test]
270 fn topology_indexes_registry_entries() {
271 let registry = InstalledFleetRegistry {
272 root_canister_id: "root-id".to_string(),
273 entries: vec![
274 RegistryEntry {
275 pid: "child-b".to_string(),
276 role: Some("worker".to_string()),
277 kind: None,
278 parent_pid: Some("root-id".to_string()),
279 module_hash: None,
280 },
281 RegistryEntry {
282 pid: "root-id".to_string(),
283 role: Some("root".to_string()),
284 kind: None,
285 parent_pid: None,
286 module_hash: None,
287 },
288 RegistryEntry {
289 pid: "child-a".to_string(),
290 role: Some("app".to_string()),
291 kind: None,
292 parent_pid: Some("root-id".to_string()),
293 module_hash: None,
294 },
295 ],
296 };
297
298 let topology = ResolvedFleetTopology::from_registry(®istry);
299
300 assert_eq!(
301 topology
302 .children_by_parent
303 .get(&Some("root-id".to_string())),
304 Some(&vec!["child-a".to_string(), "child-b".to_string()])
305 );
306 assert_eq!(topology.roles_by_canister["child-a"], "app");
307 assert_eq!(topology.root_canister_id, "root-id");
308 }
309
310 #[test]
312 fn detects_local_canister_not_found_error() {
313 assert!(is_canister_not_found_error(
314 "local replica rejected query: code=3 message=Canister uxrrr-q7777-77774-qaaaq-cai not found"
315 ));
316 assert!(!is_canister_not_found_error(
317 "local replica rejected query: code=5 message=some other failure"
318 ));
319 }
320}