feat(cli): Use IndexMap instead of BTreeMap where possible, to keep t…

…he ordering (#904) Co-authored-by: pietrodimarco-dfinity <124565147+pietrodimarco-dfinity@users.noreply.github.com>
dfinity · Sep 10, 2024 · 48da0a6 · 48da0a6
1 parent 8194404
commit 48da0a6
Show file tree

Hide file tree

Showing 26 changed files with 268 additions and 219 deletions.
diff --git a/Cargo.Bazel.lock b/Cargo.Bazel.lock
@@ -1,5 +1,5 @@
 {
-  "checksum": "8f276b46cbd4dab6df35c316a7bc7ed6716d52e987e3b60aa1a127416a00915a",
+  "checksum": "ac1c7f35e7cd235485ebc4e9085ef127ba2881abba9deee4940ac3a40ac97b9a",
   "crates": {
     "actix-codec 0.5.2": {
       "name": "actix-codec",
@@ -5739,13 +5739,13 @@
       },
       "license": "MIT OR Apache-2.0"
     },
-    "clap_complete 4.5.24": {
+    "clap_complete 4.5.26": {
       "name": "clap_complete",
-      "version": "4.5.24",
+      "version": "4.5.26",
       "repository": {
         "Http": {
-          "url": "https://static.crates.io/crates/clap_complete/4.5.24/download",
-          "sha256": "6d7db6eca8c205649e8d3ccd05aa5042b1800a784e56bc7c43524fde8abbfa9b"
+          "url": "https://static.crates.io/crates/clap_complete/4.5.26/download",
+          "sha256": "205d5ef6d485fa47606b98b0ddc4ead26eb850aaa86abfb562a94fb3280ecba0"
         }
       },
       "targets": [
@@ -5780,7 +5780,7 @@
           "selects": {}
         },
         "edition": "2021",
-        "version": "4.5.24"
+        "version": "4.5.26"
       },
       "license": "MIT OR Apache-2.0"
     },
@@ -8838,6 +8838,10 @@
               "id": "ic-base-types 0.9.0",
               "target": "ic_base_types"
             },
+            {
+              "id": "indexmap 2.5.0",
+              "target": "indexmap"
+            },
             {
               "id": "itertools 0.13.0",
               "target": "itertools"
@@ -10082,7 +10086,7 @@
               "target": "clap_num"
             },
             {
-              "id": "clap_complete 4.5.24",
+              "id": "clap_complete 4.5.26",
               "target": "clap_complete"
             },
             {
@@ -10213,6 +10217,10 @@
               "id": "ic-types 0.9.0",
               "target": "ic_types"
             },
+            {
+              "id": "indexmap 2.5.0",
+              "target": "indexmap"
+            },
             {
               "id": "itertools 0.13.0",
               "target": "itertools"
@@ -19828,6 +19836,10 @@
               "id": "ic-types 0.9.0",
               "target": "ic_types"
             },
+            {
+              "id": "indexmap 2.5.0",
+              "target": "indexmap"
+            },
             {
               "id": "itertools 0.13.0",
               "target": "itertools"
@@ -20071,6 +20083,10 @@
               "id": "ic-types 0.9.0",
               "target": "ic_types"
             },
+            {
+              "id": "indexmap 2.5.0",
+              "target": "indexmap"
+            },
             {
               "id": "registry-canister 0.9.0",
               "target": "registry_canister"
@@ -29689,6 +29705,10 @@
               "id": "ic-types 0.9.0",
               "target": "ic_types"
             },
+            {
+              "id": "indexmap 2.5.0",
+              "target": "indexmap"
+            },
             {
               "id": "rand 0.8.5",
               "target": "rand"

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -64,7 +64,7 @@ clap = { version = "4.5", features = [
     "string",
     "cargo",
 ] }
-clap_complete = "4.5.24"
+clap_complete = "4.5.26"
 clio = { version = "0.3.5", features = ["clap", "clap-parse"] }
 colored = "2.1.0"
 comfy-table = "7.1.1"

diff --git a/rs/cli/Cargo.toml b/rs/cli/Cargo.toml
@@ -49,6 +49,7 @@ ic-registry-subnet-type = { workspace = true }
 ic-sns-wasm = { workspace = true }
 ic-sys = { workspace = true }
 ic-types = { workspace = true }
+indexmap = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 pretty_env_logger = { workspace = true }

diff --git a/rs/cli/src/commands/hostos/rollout_from_node_group.rs b/rs/cli/src/commands/hostos/rollout_from_node_group.rs
@@ -7,7 +7,7 @@ use crate::{
     operations::hostos_rollout::{NodeGroupUpdate, NumberOfNodes},
 };
 
-#[derive(ValueEnum, Copy, Clone, Debug, Ord, Eq, PartialEq, PartialOrd, Default)]
+#[derive(ValueEnum, Copy, Clone, Debug, Ord, Eq, PartialEq, PartialOrd, Default, Hash)]
 pub enum NodeOwner {
     Dfinity,
     Others,
@@ -25,7 +25,7 @@ impl std::fmt::Display for NodeOwner {
     }
 }
 
-#[derive(ValueEnum, Copy, Clone, Debug, Ord, Eq, PartialEq, PartialOrd, Default)]
+#[derive(ValueEnum, Copy, Clone, Debug, Ord, Eq, PartialEq, PartialOrd, Default, Hash)]
 pub enum NodeAssignment {
     Unassigned,
     Assigned,

diff --git a/rs/cli/src/commands/registry.rs b/rs/cli/src/commands/registry.rs
@@ -16,6 +16,7 @@ use ic_protobuf::registry::{
 };
 use ic_registry_subnet_type::SubnetType;
 use ic_types::PrincipalId;
+use indexmap::IndexMap;
 use itertools::Itertools;
 use log::{info, warn};
 use serde::Serialize;
@@ -83,7 +84,7 @@ impl Registry {
 
         // Calculate number of rewardable nodes for node operators
         for node_operator in node_operators.values_mut() {
-            let mut nodes_by_health = BTreeMap::new();
+            let mut nodes_by_health = IndexMap::new();
             for node_details in nodes.iter().filter(|n| n.node_operator_id == node_operator.node_operator_principal_id) {
                 let node_id = node_details.node_id;
                 let health = node_details.status.to_string();
@@ -129,7 +130,7 @@ fn get_elected_host_os_versions(local_registry: &Arc<dyn LazyRegistry>) -> anyho
     local_registry.elected_hostos_records()
 }
 
-async fn get_node_operators(local_registry: &Arc<dyn LazyRegistry>, network: &Network) -> anyhow::Result<BTreeMap<PrincipalId, NodeOperator>> {
+async fn get_node_operators(local_registry: &Arc<dyn LazyRegistry>, network: &Network) -> anyhow::Result<IndexMap<PrincipalId, NodeOperator>> {
     let all_nodes = local_registry.nodes().await?;
     let operators = local_registry
         .operators()
@@ -163,7 +164,7 @@ async fn get_node_operators(local_registry: &Arc<dyn LazyRegistry>, network: &Ne
                 },
             )
         })
-        .collect::<BTreeMap<_, _>>();
+        .collect::<IndexMap<_, _>>();
     Ok(node_operators)
 }
 
@@ -203,7 +204,7 @@ fn get_unassigned_nodes(local_registry: &Arc<dyn LazyRegistry>) -> anyhow::Resul
 
 async fn get_nodes(
     local_registry: &Arc<dyn LazyRegistry>,
-    node_operators: &BTreeMap<PrincipalId, NodeOperator>,
+    node_operators: &IndexMap<PrincipalId, NodeOperator>,
     subnets: &[SubnetRecord],
     network: &Network,
 ) -> anyhow::Result<Vec<NodeDetails>> {
@@ -260,7 +261,7 @@ fn get_node_rewards_table(local_registry: &Arc<dyn LazyRegistry>, network: &Netw
                 panic!("Failed to get Node Rewards Table for mainnet")
             } else {
                 warn!("Failed to get Node Rewards Table for {}", network.name);
-                BTreeMap::new()
+                IndexMap::new()
             }
         }
     };
@@ -361,7 +362,7 @@ struct NodeDetails {
 struct SubnetRecord {
     subnet_id: PrincipalId,
     membership: Vec<String>,
-    nodes: BTreeMap<PrincipalId, NodeDetails>,
+    nodes: IndexMap<PrincipalId, NodeDetails>,
     max_ingress_bytes_per_message: u64,
     max_ingress_messages_per_block: u64,
     max_block_payload_size: u64,
@@ -393,7 +394,7 @@ struct NodeOperator {
     rewardable_nodes: BTreeMap<String, u32>,
     ipv6: Option<String>,
     total_up_nodes: u32,
-    nodes_health: BTreeMap<String, Vec<PrincipalId>>,
+    nodes_health: IndexMap<String, Vec<PrincipalId>>,
     rewards_correct: bool,
 }
 

diff --git a/rs/cli/src/commands/update_authorized_subnets.rs b/rs/cli/src/commands/update_authorized_subnets.rs
@@ -1,4 +1,5 @@
-use std::{collections::BTreeMap, path::PathBuf, sync::Arc};
+use indexmap::IndexMap;
+use std::{path::PathBuf, sync::Arc};
 
 use clap::{error::ErrorKind, Args};
 use ic_management_types::Subnet;
@@ -54,7 +55,7 @@ impl ExecutableCommand for UpdateAuthorizedSubnets {
 
         let registry = ctx.registry().await;
         let subnets = registry.subnets().await?;
-        let mut excluded_subnets = BTreeMap::new();
+        let mut excluded_subnets = IndexMap::new();
 
         let human_bytes = human_bytes::human_bytes(self.state_size_limit as f64);
         let agent = ctx.create_ic_agent_canister_client(None)?;
@@ -87,11 +88,11 @@ impl ExecutableCommand for UpdateAuthorizedSubnets {
             }
         }
 
-        let summary = construct_summary(&subnets, &excluded_subnets)?;
+        let summary = construct_summary(&subnets, &excluded_subnets, ctx.forum_post_link())?;
 
         let authorized = subnets
             .keys()
-            .filter(|subnet_id| !excluded_subnets.contains_key(subnet_id))
+            .filter(|subnet_id| !excluded_subnets.contains_key(*subnet_id))
             .cloned()
             .collect();
 
@@ -136,13 +137,19 @@ impl UpdateAuthorizedSubnets {
     }
 }
 
-fn construct_summary(subnets: &Arc<BTreeMap<PrincipalId, Subnet>>, excluded_subnets: &BTreeMap<PrincipalId, String>) -> anyhow::Result<String> {
+fn construct_summary(
+    subnets: &Arc<IndexMap<PrincipalId, Subnet>>,
+    excluded_subnets: &IndexMap<PrincipalId, String>,
+    forum_post_link: Option<String>,
+) -> anyhow::Result<String> {
     Ok(format!(
         "Updating the list of authorized subnets to:
 
 | Subnet id | Public | Description |
 | --------- | ------ | ----------- |
-{}",
+{}
+{}
+",
         subnets
             .values()
             .map(|s| {
@@ -154,6 +161,10 @@ fn construct_summary(subnets: &Arc<BTreeMap<PrincipalId, Subnet>>, excluded_subn
                     excluded_desc.map(|s| s.to_string()).unwrap_or_default()
                 )
             })
-            .join("\n")
+            .join("\n"),
+        match forum_post_link {
+            Some(link) => format!("\nForum post link: {}", link),
+            None => "".to_string(),
+        }
     ))
 }