From 63cfb73b6aa98deacc57ada949c445bb4ef74736 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sa=C5=A1a=20Tomi=C4=87?= <sasa.tomic@dfinity.org>
Date: Wed, 13 Nov 2024 16:49:32 +0100
Subject: [PATCH] fix display of decentralization changes

---
 rs/cli/src/runner.rs                          |  6 +++---
 rs/ic-management-backend/src/lazy_registry.rs | 13 ++++--------
 rs/ic-management-backend/src/registry.rs      | 21 +++----------------
 rs/ic-management-types/src/lib.rs             |  3 ++-
 4 files changed, 12 insertions(+), 31 deletions(-)

diff --git a/rs/cli/src/runner.rs b/rs/cli/src/runner.rs
index deebf89bf..1ce4fad4e 100644
--- a/rs/cli/src/runner.rs
+++ b/rs/cli/src/runner.rs
@@ -822,7 +822,7 @@ impl Runner {
     ) -> anyhow::Result<()> {
         let subnet_before = match override_subnet_nodes {
             Some(nodes) => {
-                let nodes = self.registry.get_decentralized_nodes(&nodes).await?;
+                let nodes = self.registry.get_nodes_from_ids(&nodes).await?;
                 DecentralizedSubnet::new_with_subnet_id_and_nodes(change.subnet_id, nodes)
             }
             None => self
@@ -835,11 +835,11 @@ impl Runner {
         let health_of_nodes = self.health_of_nodes().await?;
 
         // Simulate node removal
-        let removed_nodes = self.registry.get_decentralized_nodes(&change.get_removed_node_ids()).await?;
+        let removed_nodes = self.registry.get_nodes_from_ids(&change.get_removed_node_ids()).await?;
         let subnet_mid = subnet_before.without_nodes(&removed_nodes).map_err(|e| anyhow::anyhow!(e))?;
 
         // Now simulate node addition
-        let added_nodes = self.registry.get_decentralized_nodes(&change.get_added_node_ids()).await?;
+        let added_nodes = self.registry.get_nodes_from_ids(&change.get_added_node_ids()).await?;
 
         let subnet_after = subnet_mid.with_nodes(&added_nodes);
 
diff --git a/rs/ic-management-backend/src/lazy_registry.rs b/rs/ic-management-backend/src/lazy_registry.rs
index c8eefbb6d..c776bb67d 100644
--- a/rs/ic-management-backend/src/lazy_registry.rs
+++ b/rs/ic-management-backend/src/lazy_registry.rs
@@ -102,15 +102,10 @@ pub trait LazyRegistry:
         })
     }
 
-    fn get_decentralized_nodes<'a>(&'a self, principals: &'a [PrincipalId]) -> BoxFuture<'a, anyhow::Result<Vec<Node>>> {
+    fn get_nodes_from_ids<'a>(&'a self, principals: &'a [PrincipalId]) -> BoxFuture<'a, anyhow::Result<Vec<Node>>> {
         Box::pin(async {
-            Ok(self
-                .nodes()
-                .await?
-                .values()
-                .filter(|n| principals.contains(&n.principal))
-                .cloned()
-                .collect_vec())
+            let all_nodes = self.nodes().await?;
+            Ok(principals.iter().filter_map(|p| all_nodes.get(p).cloned()).collect())
         })
     }
 
@@ -753,7 +748,7 @@ impl LazyRegistry for LazyRegistryImpl {
         })
     }
 
-    fn get_decentralized_nodes<'a>(&'a self, principals: &'a [PrincipalId]) -> BoxFuture<'a, anyhow::Result<Vec<Node>>> {
+    fn get_nodes_from_ids<'a>(&'a self, principals: &'a [PrincipalId]) -> BoxFuture<'a, anyhow::Result<Vec<Node>>> {
         Box::pin(async {
             Ok(self
                 .nodes()
diff --git a/rs/ic-management-backend/src/registry.rs b/rs/ic-management-backend/src/registry.rs
index 59971cbcb..f4f7d52dd 100644
--- a/rs/ic-management-backend/src/registry.rs
+++ b/rs/ic-management-backend/src/registry.rs
@@ -787,12 +787,9 @@ impl RegistryState {
         self.network.get_nns_urls()
     }
 
-    pub fn get_decentralized_nodes(&self, principals: &[PrincipalId]) -> Vec<Node> {
-        self.nodes()
-            .values()
-            .filter(|node| principals.contains(&node.principal))
-            .cloned()
-            .collect_vec()
+    pub fn get_nodes_from_ids(&self, principals: &[PrincipalId]) -> Vec<Node> {
+        let all_nodes = self.nodes();
+        principals.iter().filter_map(|p| all_nodes.get(p).cloned()).collect()
     }
 
     pub async fn get_unassigned_nodes_replica_version(&self) -> Result<String, anyhow::Error> {
@@ -810,18 +807,6 @@ impl RegistryState {
             _ => Err(anyhow::anyhow!("No GuestOS version for unassigned nodes found".to_string(),)),
         }
     }
-
-    #[allow(dead_code)]
-    pub async fn node(&self, node_id: PrincipalId) -> Node {
-        self.nodes
-            .iter()
-            .filter(|(&id, _)| id == node_id)
-            .collect::<Vec<_>>()
-            .first()
-            .unwrap()
-            .1
-            .clone()
-    }
 }
 
 impl decentralization::network::TopologyManager for RegistryState {}
diff --git a/rs/ic-management-types/src/lib.rs b/rs/ic-management-types/src/lib.rs
index 19bd1934c..639f6253b 100644
--- a/rs/ic-management-types/src/lib.rs
+++ b/rs/ic-management-types/src/lib.rs
@@ -330,6 +330,7 @@ impl Node {
     }
     pub fn get_features(&self) -> NodeFeatures {
         let features = if let Some(features) = &self.cached_features.get() {
+            // Return a clone of the cached value, if it exists
             (*features).clone()
         } else {
             let country = self
@@ -383,7 +384,7 @@ impl Node {
     }
 
     pub fn get_feature(&self, feature: &NodeFeature) -> Option<String> {
-        self.cached_features.get().and_then(|fts| fts.get(feature))
+        self.get_features().get(feature)
     }
 
     pub fn matches_feature_value(&self, value: &str) -> bool {