diff --git a/Cargo.Bazel.lock b/Cargo.Bazel.lock index 1f15a2331..6d7158858 100644 --- 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 index 65f17edae..dc573fea0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1139,9 +1139,9 @@ dependencies = [ [[package]] name = "clap_complete" -version = "4.5.24" +version = "4.5.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d7db6eca8c205649e8d3ccd05aa5042b1800a784e56bc7c43524fde8abbfa9b" +checksum = "205d5ef6d485fa47606b98b0ddc4ead26eb850aaa86abfb562a94fb3280ecba0" dependencies = [ "clap 4.5.16", ] @@ -1749,6 +1749,7 @@ dependencies = [ "futures", "ic-base-types", "ic-management-types", + "indexmap 2.5.0", "itertools 0.13.0", "log", "rand", @@ -2029,6 +2030,7 @@ dependencies = [ "ic-sns-wasm", "ic-sys", "ic-types", + "indexmap 2.5.0", "itertools 0.13.0", "keyring", "log", @@ -3987,6 +3989,7 @@ dependencies = [ "ic-registry-nns-data-provider", "ic-registry-subnet-type", "ic-types", + "indexmap 2.5.0", "itertools 0.13.0", "lazy_static", "log", @@ -4039,6 +4042,7 @@ dependencies = [ "ic-protobuf", "ic-registry-subnet-type", "ic-types", + "indexmap 2.5.0", "registry-canister", "reqwest", "serde", @@ -5989,6 +5993,7 @@ dependencies = [ "ic-management-backend", "ic-management-types", "ic-types", + "indexmap 2.5.0", "pretty_assertions", "rand", "reqwest", diff --git a/Cargo.toml b/Cargo.toml index f962ee2bc..021bc34b4 100644 --- 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 index bef90da28..8bad18b74 100644 --- 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 index 3d33e358e..6c9e410f6 100644 --- 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 index 04c337b8f..1eb97a588 100644 --- 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) -> anyho local_registry.elected_hostos_records() } -async fn get_node_operators(local_registry: &Arc, network: &Network) -> anyhow::Result> { +async fn get_node_operators(local_registry: &Arc, network: &Network) -> anyhow::Result> { let all_nodes = local_registry.nodes().await?; let operators = local_registry .operators() @@ -163,7 +164,7 @@ async fn get_node_operators(local_registry: &Arc, network: &Ne }, ) }) - .collect::>(); + .collect::>(); Ok(node_operators) } @@ -203,7 +204,7 @@ fn get_unassigned_nodes(local_registry: &Arc) -> anyhow::Resul async fn get_nodes( local_registry: &Arc, - node_operators: &BTreeMap, + node_operators: &IndexMap, subnets: &[SubnetRecord], network: &Network, ) -> anyhow::Result> { @@ -260,7 +261,7 @@ fn get_node_rewards_table(local_registry: &Arc, 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, - nodes: BTreeMap, + nodes: IndexMap, 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, ipv6: Option, total_up_nodes: u32, - nodes_health: BTreeMap>, + nodes_health: IndexMap>, rewards_correct: bool, } diff --git a/rs/cli/src/commands/update_authorized_subnets.rs b/rs/cli/src/commands/update_authorized_subnets.rs index af1b1fb40..18ca9c2b3 100644 --- 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>, excluded_subnets: &BTreeMap) -> anyhow::Result { +fn construct_summary( + subnets: &Arc>, + excluded_subnets: &IndexMap, + forum_post_link: Option, +) -> anyhow::Result { 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>, 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(), + } )) } diff --git a/rs/cli/src/operations/hostos_rollout.rs b/rs/cli/src/operations/hostos_rollout.rs index 16b86c08e..88001f4a4 100644 --- a/rs/cli/src/operations/hostos_rollout.rs +++ b/rs/cli/src/operations/hostos_rollout.rs @@ -5,9 +5,10 @@ use ic_base_types::{NodeId, PrincipalId}; use ic_management_backend::health::{self, HealthStatusQuerier}; use ic_management_backend::proposal::ProposalAgent; use ic_management_types::{HealthStatus, Network, Node, Subnet, UpdateNodesHostosVersionsProposal}; +use indexmap::IndexMap; use log::{debug, info, warn}; use std::sync::Arc; -use std::{collections::BTreeMap, fmt::Display, str::FromStr}; +use std::{fmt::Display, str::FromStr}; use crate::commands::hostos::rollout_from_node_group::{NodeAssignment, NodeOwner}; @@ -39,7 +40,7 @@ impl Display for HostosRolloutReason { } } -#[derive(Copy, Clone, Debug, Ord, Eq, PartialEq, PartialOrd)] +#[derive(Copy, Clone, Debug, Ord, Eq, PartialEq, PartialOrd, Hash)] pub struct NodeGroup { pub assignment: NodeAssignment, pub owner: NodeOwner, @@ -129,8 +130,8 @@ enum CandidatesSelection { #[derive(Clone)] pub struct HostosRollout { nodes_all: Vec, - pub grouped_nodes: BTreeMap>, - pub subnets: Arc>, + pub grouped_nodes: IndexMap>, + pub subnets: Arc>, pub network: Network, pub proposal_agent: Arc, pub only_filter: Vec, @@ -139,15 +140,15 @@ pub struct HostosRollout { } impl HostosRollout { pub fn new( - nodes: Arc>, - subnets: Arc>, + nodes: Arc>, + subnets: Arc>, network: &Network, proposal_agent: Arc, rollout_version: &str, only_filter: &[String], exclude_filter: &[String], ) -> Self { - let grouped_nodes: BTreeMap> = nodes + let grouped_nodes: IndexMap> = nodes .values() .cloned() .map(|node| { @@ -169,7 +170,7 @@ impl HostosRollout { ); (NodeGroup::new(assignment, owner), node) }) - .fold(BTreeMap::new(), |mut acc, (node_group, node)| { + .fold(IndexMap::new(), |mut acc, (node_group, node)| { acc.entry(node_group).or_default().push(node); acc }); @@ -219,12 +220,12 @@ impl HostosRollout { .collect() } - async fn nodes_by_status(&self, nodes: Vec, nodes_health: BTreeMap) -> BTreeMap> { + async fn nodes_by_status(&self, nodes: Vec, nodes_health: IndexMap) -> IndexMap> { let nodes_by_status = nodes .iter() .cloned() .map(|node| (nodes_health.get(&node.principal).cloned().unwrap_or(HealthStatus::Unknown), node)) - .fold(BTreeMap::new(), |mut acc: BTreeMap>, (status, node)| { + .fold(IndexMap::new(), |mut acc: IndexMap>, (status, node)| { acc.entry(status).or_default().push(node); acc }); @@ -325,7 +326,7 @@ impl HostosRollout { async fn candidates_selection( &self, - nodes_health: BTreeMap, + nodes_health: IndexMap, nodes_with_open_proposals: Vec, nodes_in_group: Vec, ) -> anyhow::Result { @@ -363,7 +364,7 @@ impl HostosRollout { #[async_recursion] async fn with_nodes_health_and_open_proposals( &self, - nodes_health: BTreeMap, + nodes_health: IndexMap, nodes_with_open_proposals: Vec, update_group: NodeGroupUpdate, ) -> anyhow::Result { @@ -535,6 +536,7 @@ pub mod test { use crate::operations::hostos_rollout::NodeOwner::{Dfinity, Others}; use ic_management_backend::proposal::ProposalAgentImpl; use ic_management_types::{Network, Node, Operator, Provider, Subnet}; + use std::collections::BTreeMap; use std::net::Ipv6Addr; use super::*; @@ -548,7 +550,7 @@ pub mod test { let assigned_dfinity = gen_test_nodes(Some(subnet_id), 10, 0, version_one.clone(), true, false); let unassigned_dfinity_nodes = gen_test_nodes(None, 10, 10, version_one.clone(), true, false); let assigned_others_nodes = gen_test_nodes(Some(subnet_id), 10, 20, version_two.clone(), false, false); - let union: BTreeMap = { + let union: IndexMap = { assigned_dfinity .clone() .into_iter() @@ -559,7 +561,7 @@ pub mod test { let union_nodes = union.values().cloned().collect::>(); let subnet = { - let mut sub = BTreeMap::new(); + let mut sub = IndexMap::new(); sub.insert( subnet_id, Subnet { @@ -575,7 +577,7 @@ pub mod test { .keys() .cloned() .map(|principal| (principal, HealthStatus::Healthy)) - .collect::>(); + .collect::>(); let open_proposals: Vec = vec![]; @@ -685,8 +687,8 @@ pub mod test { hostos_version: String, dfinity_owned: bool, is_api_boundary_node: bool, - ) -> BTreeMap { - let mut n = BTreeMap::new(); + ) -> IndexMap { + let mut n = IndexMap::new(); for i in start_at_number..start_at_number + num_nodes { let node = Node { principal: PrincipalId::new_node_test_id(i), diff --git a/rs/cli/src/runner.rs b/rs/cli/src/runner.rs index 8a4d4d235..9f487fcab 100644 --- a/rs/cli/src/runner.rs +++ b/rs/cli/src/runner.rs @@ -1,6 +1,4 @@ use std::cell::RefCell; -use std::collections::BTreeMap; -use std::collections::BTreeSet; use std::sync::Arc; use decentralization::network::DecentralizedSubnet; @@ -28,6 +26,7 @@ use ic_management_types::NodeFeature; use ic_management_types::Release; use ic_management_types::TopologyChangePayload; use ic_types::PrincipalId; +use indexmap::{IndexMap, IndexSet}; use itertools::Itertools; use log::info; use log::warn; @@ -120,7 +119,7 @@ impl Runner { Ok(()) } - pub async fn health_of_nodes(&self) -> anyhow::Result> { + pub async fn health_of_nodes(&self) -> anyhow::Result> { let health_client = health::HealthClient::new(self.network.clone()); health_client.nodes().await } @@ -130,7 +129,7 @@ impl Runner { request: ic_management_types::requests::SubnetResizeRequest, motivation: String, forum_post_link: Option, - health_of_nodes: &BTreeMap, + health_of_nodes: &IndexMap, ) -> anyhow::Result<()> { let change = self .registry @@ -571,7 +570,7 @@ impl Runner { None => true, }) .map(|(id, subnet)| (*id, subnet.clone())) - .collect::>(); + .collect::>(); let (available_nodes, health_of_nodes) = try_join(self.registry.available_nodes().map_err(anyhow::Error::from), health_client.nodes()).await?; @@ -677,8 +676,8 @@ impl Runner { let ic_repo = self.ic_repo().await; let hosts = ic_repo.hostos_releases().await?; let active_releases = hosts.get_active_branches(); - let hostos_versions: BTreeSet = self.registry.nodes().await?.values().map(|s| s.hostos_version.clone()).collect(); - let versions_in_proposals: BTreeSet = self + let hostos_versions: IndexSet = self.registry.nodes().await?.values().map(|s| s.hostos_version.clone()).collect(); + let versions_in_proposals: IndexSet = self .proposal_agent .list_open_elect_hostos_proposals() .await? @@ -704,9 +703,9 @@ impl Runner { let ic_repo = self.ic_repo().await; let guests = ic_repo.guestos_releases().await?; let active_releases = guests.get_active_branches(); - let subnet_versions: BTreeSet = self.registry.subnets().await?.values().map(|s| s.replica_version.clone()).collect(); + let subnet_versions: IndexSet = self.registry.subnets().await?.values().map(|s| s.replica_version.clone()).collect(); let version_on_unassigned_nodes = self.registry.unassigned_nodes_replica_version().await?; - let versions_in_proposals: BTreeSet = self + let versions_in_proposals: IndexSet = self .proposal_agent .list_open_elect_replica_proposals() .await? @@ -816,7 +815,7 @@ pub fn replace_proposal_options(change: &SubnetChangeResponse, forum_post_link: }) } -fn nodes_by_dc(nodes: Vec) -> BTreeMap> { +fn nodes_by_dc(nodes: Vec) -> IndexMap> { nodes .iter() .cloned() @@ -832,7 +831,7 @@ fn nodes_by_dc(nodes: Vec) -> BTreeMap> { n.operator.datacenter, ) }) - .fold(BTreeMap::new(), |mut acc, (node_id, subnet, dc)| { + .fold(IndexMap::new(), |mut acc, (node_id, subnet, dc)| { acc.entry(dc.unwrap_or_default().name).or_default().push((node_id, subnet)); acc }) diff --git a/rs/cli/src/unit_tests/version.rs b/rs/cli/src/unit_tests/version.rs index 02f98defe..07f767961 100644 --- a/rs/cli/src/unit_tests/version.rs +++ b/rs/cli/src/unit_tests/version.rs @@ -1,7 +1,5 @@ -use std::{ - collections::BTreeMap, - sync::{Arc, RwLock}, -}; +use indexmap::IndexMap; +use std::sync::{Arc, RwLock}; use futures::future::ok; use ic_management_backend::{lazy_git::MockLazyGit, lazy_registry::MockLazyRegistry, proposal::MockProposalAgent}; @@ -35,7 +33,7 @@ async fn guest_os_elect_version_tests() { .returning(|| Box::pin(ok(Arc::new(ArtifactReleases::new(ic_management_types::Artifact::GuestOs))))); let mut registry = MockLazyRegistry::new(); - registry.expect_subnets().returning(|| Box::pin(ok(Arc::new(BTreeMap::new())))); + registry.expect_subnets().returning(|| Box::pin(ok(Arc::new(IndexMap::new())))); registry .expect_unassigned_nodes_replica_version() .returning(|| Box::pin(ok(Arc::new("some_ver".to_string())))); diff --git a/rs/decentralization/Cargo.toml b/rs/decentralization/Cargo.toml index 19732d935..05d7815fb 100644 --- a/rs/decentralization/Cargo.toml +++ b/rs/decentralization/Cargo.toml @@ -14,6 +14,7 @@ ahash = { workspace = true } anyhow = { workspace = true } ic-base-types = { workspace = true } ic-management-types = { workspace = true } +indexmap = { workspace = true } itertools = { workspace = true } log = { workspace = true } rand = { workspace = true } diff --git a/rs/decentralization/src/lib.rs b/rs/decentralization/src/lib.rs index c74f47ecc..e351ceeb2 100644 --- a/rs/decentralization/src/lib.rs +++ b/rs/decentralization/src/lib.rs @@ -1,8 +1,8 @@ pub mod nakamoto; pub mod network; pub mod subnets; +use indexmap::IndexMap; use itertools::Itertools; -use std::collections::BTreeMap; use std::fmt::{Display, Formatter}; use ic_base_types::PrincipalId; @@ -15,17 +15,17 @@ pub struct SubnetChangeResponse { pub removed_with_desc: Vec<(PrincipalId, String)>, #[serde(skip_serializing_if = "Option::is_none")] pub subnet_id: Option, - pub health_of_nodes: BTreeMap, + pub health_of_nodes: IndexMap, pub score_before: nakamoto::NakamotoScore, pub score_after: nakamoto::NakamotoScore, pub motivation: Option, pub comment: Option, pub run_log: Option>, - pub feature_diff: BTreeMap, + pub feature_diff: IndexMap, pub proposal_id: Option, } -pub type FeatureDiff = BTreeMap; +pub type FeatureDiff = IndexMap; impl SubnetChangeResponse { pub fn with_motivation(self, motivation: String) -> Self { @@ -34,7 +34,7 @@ impl SubnetChangeResponse { ..self } } - pub fn with_health_of_nodes(self, health_of_nodes: BTreeMap) -> Self { + pub fn with_health_of_nodes(self, health_of_nodes: IndexMap) -> Self { SubnetChangeResponse { health_of_nodes, ..self } } } @@ -45,7 +45,7 @@ impl From<&network::SubnetChange> for SubnetChangeResponse { added_with_desc: change.added().iter().map(|n| (n.0.id, n.1.clone())).collect(), removed_with_desc: change.removed().iter().map(|n| (n.0.id, n.1.clone())).collect(), subnet_id: if change.id == Default::default() { None } else { Some(change.id) }, - health_of_nodes: BTreeMap::new(), + health_of_nodes: IndexMap::new(), score_before: nakamoto::NakamotoScore::new_from_nodes(&change.old_nodes), score_after: nakamoto::NakamotoScore::new_from_nodes(&change.new_nodes), motivation: None, @@ -56,7 +56,7 @@ impl From<&network::SubnetChange> for SubnetChangeResponse { NodeFeature::variants() .into_iter() .map(|f| (f, FeatureDiff::new())) - .collect::>(), + .collect::>(), |mut acc, n| { for f in NodeFeature::variants() { acc.get_mut(&f).unwrap().entry(n.get_feature(&f)).or_insert((0, 0)).0 += 1; diff --git a/rs/decentralization/src/nakamoto/mod.rs b/rs/decentralization/src/nakamoto/mod.rs index 2b9fb455a..3a520d1b7 100644 --- a/rs/decentralization/src/nakamoto/mod.rs +++ b/rs/decentralization/src/nakamoto/mod.rs @@ -1,10 +1,11 @@ use crate::network::Node; use ahash::{AHashMap, AHasher}; +use indexmap::IndexMap; use itertools::Itertools; use serde::{Deserialize, Serialize}; use std::cell::RefCell; use std::cmp::Ordering; -use std::collections::{BTreeMap, VecDeque}; +use std::collections::VecDeque; use std::fmt::{self, Debug, Display, Formatter}; use std::hash::Hasher; use std::iter::{FromIterator, IntoIterator}; @@ -13,7 +14,7 @@ use ic_management_types::NodeFeature; #[derive(Eq, PartialEq, Clone, Serialize, Deserialize, Debug)] pub struct NodeFeatures { - pub feature_map: BTreeMap, + pub feature_map: IndexMap, } impl fmt::Display for NodeFeatures { @@ -32,7 +33,7 @@ impl NodeFeatures { #[cfg(test)] fn new_test_feature_set(value: &str) -> Self { - let mut result = BTreeMap::new(); + let mut result = IndexMap::new(); for feature in NodeFeature::variants() { result.insert(feature, value.to_string()); } @@ -50,7 +51,7 @@ impl NodeFeatures { impl FromIterator<(NodeFeature, &'static str)> for NodeFeatures { fn from_iter>(iter: I) -> Self { Self { - feature_map: BTreeMap::from_iter(iter.into_iter().map(|x| (x.0, String::from(x.1)))), + feature_map: IndexMap::from_iter(iter.into_iter().map(|x| (x.0, String::from(x.1)))), } } } @@ -58,7 +59,7 @@ impl FromIterator<(NodeFeature, &'static str)> for NodeFeatures { impl FromIterator<(NodeFeature, std::string::String)> for NodeFeatures { fn from_iter>(iter: I) -> Self { Self { - feature_map: BTreeMap::from_iter(iter), + feature_map: IndexMap::from_iter(iter), } } } @@ -77,9 +78,9 @@ thread_local! { /// For instance: [NodeFeature::NodeProvider], [NodeFeature::DataCenter], etc... /// For a complete reference check [NodeFeature] pub struct NakamotoScore { - coefficients: BTreeMap, - value_counts: BTreeMap>, - controlled_nodes: BTreeMap, + coefficients: IndexMap, + value_counts: IndexMap>, + controlled_nodes: IndexMap, avg_linear: f64, /// This field needs to be optional in case we get a -inf result because of @@ -101,13 +102,13 @@ pub struct NakamotoScore { impl NakamotoScore { /// Build a new NakamotoScore object from a slice of [NodeFeatures]. pub fn new_from_slice_node_features(slice_node_features: &[NodeFeatures]) -> Self { - let mut features_to_nodes_map = BTreeMap::new(); + let mut features_to_nodes_map = IndexMap::new(); for feature in NodeFeature::variants() { features_to_nodes_map.insert(feature, Vec::new()); } - // Convert a Vec> into a Vec> into a Vec> for node_features in slice_node_features.iter() { for feature in NodeFeature::variants() { @@ -158,9 +159,9 @@ impl NakamotoScore { let scores = nakamoto_calc .clone() .map(|(f, n, _)| (f, n.0 as f64)) - .collect::>(); + .collect::>(); - let controlled_nodes = nakamoto_calc.clone().map(|(f, n, _)| (f, n.1)).collect::>(); + let controlled_nodes = nakamoto_calc.clone().map(|(f, n, _)| (f, n.1)).collect::>(); let value_counts = nakamoto_calc.map(|(f, _, value_counts)| (f, value_counts)).collect(); @@ -290,7 +291,7 @@ impl NakamotoScore { } /// Get a Map with all the features and the corresponding Nakamoto score - pub fn scores_individual(&self) -> BTreeMap { + pub fn scores_individual(&self) -> IndexMap { self.coefficients.clone() } @@ -568,6 +569,7 @@ mod tests { use ahash::HashSet; use ic_base_types::PrincipalId; use ic_management_types::HealthStatus; + use indexmap::IndexMap; use itertools::Itertools; use regex::Regex; @@ -604,7 +606,7 @@ mod tests { let score = NakamotoScore::new_from_slice_node_features(&features); let score_expected = NakamotoScore { - coefficients: BTreeMap::from([ + coefficients: IndexMap::from([ (NodeFeature::City, 1.), (NodeFeature::Country, 1.), (NodeFeature::Continent, 1.), @@ -612,8 +614,8 @@ mod tests { (NodeFeature::NodeProvider, 1.), (NodeFeature::DataCenter, 1.), ]), - value_counts: BTreeMap::new(), - controlled_nodes: BTreeMap::from([ + value_counts: IndexMap::new(), + controlled_nodes: IndexMap::from([ (NodeFeature::City, 1), (NodeFeature::Country, 1), (NodeFeature::Continent, 1), @@ -771,7 +773,7 @@ mod tests { .iter() .chain(subnet_initial.nodes.iter()) .map(|n| (n.id, HealthStatus::Healthy)) - .collect::>(); + .collect::>(); println!( "initial {} Countries {:?}", @@ -828,7 +830,7 @@ mod tests { ) ); let nodes_available = new_test_nodes_with_overrides("spare", 7, 2, 0, (&NodeFeature::NodeProvider, &["NP6", "NP7"])); - let health_of_nodes = nodes_available.iter().map(|n| (n.id, HealthStatus::Healthy)).collect::>(); + let health_of_nodes = nodes_available.iter().map(|n| (n.id, HealthStatus::Healthy)).collect::>(); println!( "initial {} NPs {:?}", @@ -886,7 +888,7 @@ mod tests { .iter() .chain(subnet_initial.nodes.iter()) .map(|n| (n.id, HealthStatus::Healthy)) - .collect::>(); + .collect::>(); println!( "initial {} NPs {:?}", @@ -1099,8 +1101,8 @@ mod tests { (n, HealthStatus::Healthy) } }) - .collect::>(); - let mut important = BTreeMap::new(); + .collect::>(); + let mut important = IndexMap::new(); important.insert(subnet.principal, subnet); @@ -1124,7 +1126,7 @@ mod tests { .iter() .chain(subnet_initial.nodes.iter()) .map(|n| (n.id, HealthStatus::Healthy)) - .collect::>(); + .collect::>(); let change_initial = SubnetChangeRequest::new(subnet_initial.clone(), nodes_available, Vec::new(), Vec::new(), Vec::new(), None); diff --git a/rs/decentralization/src/network.rs b/rs/decentralization/src/network.rs index 339f57fea..a58f38d26 100644 --- a/rs/decentralization/src/network.rs +++ b/rs/decentralization/src/network.rs @@ -8,12 +8,12 @@ use anyhow::anyhow; use futures::future::BoxFuture; use ic_base_types::PrincipalId; use ic_management_types::{HealthStatus, MinNakamotoCoefficients, NetworkError, NodeFeature}; +use indexmap::IndexMap; use itertools::Itertools; use log::{debug, info, warn}; use rand::{seq::SliceRandom, SeedableRng}; use serde::{Deserialize, Serialize}; use std::cmp::Ordering; -use std::collections::BTreeMap; use std::fmt::{Debug, Display, Formatter}; use std::hash::Hash; @@ -880,7 +880,7 @@ pub trait TopologyManager: SubnetQuerier + AvailableNodesQuerier + Sync { include_nodes: Vec, exclude_nodes: Vec, only_nodes: Vec, - health_of_nodes: &'a BTreeMap, + health_of_nodes: &'a IndexMap, ) -> BoxFuture<'a, Result> { Box::pin(async move { SubnetChangeRequest { @@ -1029,7 +1029,7 @@ impl SubnetChangeRequest { mut self, optimize_count: usize, replacements_unhealthy_with_desc: &[(Node, String)], - health_of_nodes: &BTreeMap, + health_of_nodes: &IndexMap, ) -> Result { let old_nodes = self.subnet.nodes.clone(); self.subnet = self.subnet.without_nodes(replacements_unhealthy_with_desc.to_owned())?; @@ -1042,7 +1042,7 @@ impl SubnetChangeRequest { Ok(SubnetChange { old_nodes, ..result }) } - pub fn rescue(mut self, health_of_nodes: &BTreeMap) -> Result { + pub fn rescue(mut self, health_of_nodes: &IndexMap) -> Result { let old_nodes = self.subnet.nodes.clone(); let nodes_to_remove = self .subnet @@ -1074,7 +1074,7 @@ impl SubnetChangeRequest { how_many_nodes_to_add: usize, how_many_nodes_to_remove: usize, how_many_nodes_unhealthy: usize, - health_of_nodes: &BTreeMap, + health_of_nodes: &IndexMap, ) -> Result { let old_nodes = self.subnet.nodes.clone(); @@ -1137,7 +1137,7 @@ impl SubnetChangeRequest { /// Evaluates the subnet change request to simulate the requested topology /// change. Command returns all the information about the subnet before /// and after the change. - pub fn evaluate(self, health_of_nodes: &BTreeMap) -> Result { + pub fn evaluate(self, health_of_nodes: &IndexMap) -> Result { self.resize(0, 0, 0, health_of_nodes) } } @@ -1235,18 +1235,18 @@ impl Ord for NetworkHealSubnets { } pub struct NetworkHealRequest { - pub subnets: BTreeMap, + pub subnets: IndexMap, } impl NetworkHealRequest { - pub fn new(subnets: BTreeMap) -> Self { + pub fn new(subnets: IndexMap) -> Self { Self { subnets } } pub async fn heal_and_optimize( &self, mut available_nodes: Vec, - health_of_nodes: &BTreeMap, + health_of_nodes: &IndexMap, ) -> Result, NetworkError> { let mut subnets_changed = Vec::new(); let subnets_to_heal = unhealthy_with_nodes(&self.subnets, health_of_nodes) diff --git a/rs/decentralization/src/subnets.rs b/rs/decentralization/src/subnets.rs index 83d90ed7c..0400b748b 100644 --- a/rs/decentralization/src/subnets.rs +++ b/rs/decentralization/src/subnets.rs @@ -1,17 +1,17 @@ -use std::{collections::BTreeMap, sync::Arc}; - use crate::network::Node; use ic_base_types::PrincipalId; use ic_management_types::{ requests::{NodeRemoval, NodeRemovalReason}, HealthStatus, Subnet, }; +use indexmap::IndexMap; use itertools::Itertools; +use std::sync::Arc; pub async fn unhealthy_with_nodes( - subnets: &BTreeMap, - nodes_health: &BTreeMap, -) -> BTreeMap> { + subnets: &IndexMap, + nodes_health: &IndexMap, +) -> IndexMap> { subnets .clone() .into_iter() @@ -31,7 +31,7 @@ pub async fn unhealthy_with_nodes( None } }) - .collect::>() + .collect::>() } pub struct NodesRemover { @@ -45,14 +45,14 @@ pub struct NodesRemover { impl NodesRemover { pub fn remove_nodes( &self, - mut healths: std::collections::BTreeMap, - nodes_with_proposals: Arc>, + mut healths: IndexMap, + nodes_with_proposals: Arc>, ) -> (Vec, String) { let nodes_to_rm = nodes_with_proposals .values() .cloned() .map(|n| { - let status = healths.remove(&n.principal).unwrap_or(ic_management_types::HealthStatus::Unknown); + let status = healths.shift_remove(&n.principal).unwrap_or(ic_management_types::HealthStatus::Unknown); (n, status) }) .filter(|(n, _)| n.proposal.is_none()) diff --git a/rs/ic-management-backend/Cargo.toml b/rs/ic-management-backend/Cargo.toml index baa5fcea8..f3ab77a48 100644 --- a/rs/ic-management-backend/Cargo.toml +++ b/rs/ic-management-backend/Cargo.toml @@ -41,6 +41,7 @@ ic-registry-local-store = { workspace = true } ic-registry-nns-data-provider = { workspace = true } ic-registry-subnet-type = { workspace = true } ic-types = { workspace = true } +indexmap = { workspace = true } itertools = { workspace = true } lazy_static = { workspace = true } log = { workspace = true } diff --git a/rs/ic-management-backend/src/endpoints/mod.rs b/rs/ic-management-backend/src/endpoints/mod.rs index d614b0b69..18fd8dac0 100644 --- a/rs/ic-management-backend/src/endpoints/mod.rs +++ b/rs/ic-management-backend/src/endpoints/mod.rs @@ -12,9 +12,9 @@ use decentralization::network::AvailableNodesQuerier; use ic_management_types::Network; use ic_registry_nns_data_provider::registry::RegistryCanister; use ic_types::PrincipalId; +use indexmap::IndexMap; use log::{debug, info}; use serde::{Deserialize, Serialize}; -use std::collections::BTreeMap; use std::ops::Deref; use std::str::FromStr; use std::sync::Arc; @@ -224,10 +224,10 @@ async fn nodes_healths(registry: web::Data>> .map(|n| { ( n.principal, - healths.remove(&n.principal).unwrap_or(ic_management_types::HealthStatus::Unknown), + healths.shift_remove(&n.principal).unwrap_or(ic_management_types::HealthStatus::Unknown), ) }) - .collect::>() + .collect::>() })) } diff --git a/rs/ic-management-backend/src/endpoints/subnet.rs b/rs/ic-management-backend/src/endpoints/subnet.rs index ab1b4ec93..80f3d0fef 100644 --- a/rs/ic-management-backend/src/endpoints/subnet.rs +++ b/rs/ic-management-backend/src/endpoints/subnet.rs @@ -2,8 +2,8 @@ use super::*; use crate::subnets::get_proposed_subnet_changes; use ic_base_types::PrincipalId; use ic_management_types::Node; +use indexmap::IndexMap; use serde::Deserialize; -use std::collections::BTreeMap; #[derive(Deserialize)] struct SubnetRequest { @@ -20,7 +20,7 @@ pub(crate) async fn change_preview( let subnet = subnets .get(&request.subnet) .ok_or_else(|| actix_web::error::ErrorNotFound(anyhow::format_err!("subnet {} not found", request.subnet)))?; - let registry_nodes: BTreeMap = registry.read().await.nodes(); + let registry_nodes: IndexMap = registry.read().await.nodes(); get_proposed_subnet_changes(®istry_nodes, subnet) .map_err(actix_web::error::ErrorBadRequest) diff --git a/rs/ic-management-backend/src/health.rs b/rs/ic-management-backend/src/health.rs index a2ad9740d..c3532d4d6 100644 --- a/rs/ic-management-backend/src/health.rs +++ b/rs/ic-management-backend/src/health.rs @@ -1,7 +1,5 @@ -use std::{ - collections::{BTreeMap, HashSet}, - str::FromStr, -}; +use indexmap::IndexMap; +use std::{collections::HashSet, str::FromStr}; use ic_base_types::PrincipalId; use ic_management_types::{HealthStatus, Network}; @@ -26,14 +24,14 @@ impl HealthClient { } impl HealthStatusQuerier for HealthClient { - async fn subnet(&self, subnet: PrincipalId) -> anyhow::Result> { + async fn subnet(&self, subnet: PrincipalId) -> anyhow::Result> { match &self.implementation { HealthStatusQuerierImplementations::Dashboard(c) => c.subnet(subnet).await, HealthStatusQuerierImplementations::Prometheus(c) => c.subnet(subnet).await, } } - async fn nodes(&self) -> anyhow::Result> { + async fn nodes(&self) -> anyhow::Result> { match &self.implementation { HealthStatusQuerierImplementations::Dashboard(c) => c.nodes().await, HealthStatusQuerierImplementations::Prometheus(c) => c.nodes().await, @@ -57,8 +55,8 @@ impl From for HealthStatusQuerierImplementations { } pub trait HealthStatusQuerier { - fn subnet(&self, subnet: PrincipalId) -> impl std::future::Future>> + Send; - fn nodes(&self) -> impl std::future::Future>> + Send; + fn subnet(&self, subnet: PrincipalId) -> impl std::future::Future>> + Send; + fn nodes(&self) -> impl std::future::Future>> + Send; } pub struct PublicDashboardHealthClient { @@ -186,7 +184,7 @@ struct ShortNodeInfo { } impl HealthStatusQuerier for PublicDashboardHealthClient { - async fn subnet(&self, subnet: PrincipalId) -> anyhow::Result> { + async fn subnet(&self, subnet: PrincipalId) -> anyhow::Result> { Ok(self .get_all_nodes() .await? @@ -199,7 +197,7 @@ impl HealthStatusQuerier for PublicDashboardHealthClient { .collect()) } - async fn nodes(&self) -> anyhow::Result> { + async fn nodes(&self) -> anyhow::Result> { Ok(self.get_all_nodes().await?.into_iter().map(|n| (n.node_id, n.status)).collect()) } } @@ -219,7 +217,7 @@ impl PrometheusHealthClient { } impl HealthStatusQuerier for PrometheusHealthClient { - async fn subnet(&self, subnet: PrincipalId) -> anyhow::Result> { + async fn subnet(&self, subnet: PrincipalId) -> anyhow::Result> { let ic_name = self.network.legacy_name(); let subnet_name = subnet.to_string(); let query_up = Selector::new() @@ -264,7 +262,7 @@ impl HealthStatusQuerier for PrometheusHealthClient { .collect()) } - async fn nodes(&self) -> anyhow::Result> { + async fn nodes(&self) -> anyhow::Result> { let query = format!( r#"ic_replica_orchestrator:health_state:bottomk_1{{ic="{network}"}}"#, network = self.network.legacy_name(), diff --git a/rs/ic-management-backend/src/lazy_registry.rs b/rs/ic-management-backend/src/lazy_registry.rs index c46509820..b53fbe23c 100644 --- a/rs/ic-management-backend/src/lazy_registry.rs +++ b/rs/ic-management-backend/src/lazy_registry.rs @@ -1,5 +1,4 @@ -use std::collections::BTreeMap; -use std::collections::{BTreeSet, HashSet}; +use indexmap::{IndexMap, IndexSet}; use std::net::Ipv6Addr; use std::str::FromStr; use std::sync::Arc; @@ -63,15 +62,15 @@ pub trait LazyRegistry: fn sync_with_nns(&self) -> BoxFuture<'_, anyhow::Result<()>>; - fn operators(&self) -> BoxFuture<'_, anyhow::Result>>>; + fn operators(&self) -> BoxFuture<'_, anyhow::Result>>>; - fn nodes(&self) -> BoxFuture<'_, anyhow::Result>>>; + fn nodes(&self) -> BoxFuture<'_, anyhow::Result>>>; fn firewall_rule_set(&self, firewall_rule_scope: FirewallRulesScope) -> BoxFuture<'_, anyhow::Result>; - fn subnets(&self) -> BoxFuture<'_, anyhow::Result>>>; + fn subnets(&self) -> BoxFuture<'_, anyhow::Result>>>; - fn nodes_with_proposals(&self) -> BoxFuture<'_, anyhow::Result>>>; + fn nodes_with_proposals(&self) -> BoxFuture<'_, anyhow::Result>>>; fn nns_replica_version(&self) -> BoxFuture<'_, anyhow::Result>> { Box::pin(async { @@ -118,7 +117,7 @@ pub trait LazyRegistry: fn get_api_boundary_nodes(&self) -> anyhow::Result>; - fn get_node_rewards_table(&self) -> anyhow::Result>; + fn get_node_rewards_table(&self) -> anyhow::Result>; fn get_unassigned_nodes(&self) -> anyhow::Result>; @@ -130,7 +129,7 @@ pub trait LazyRegistry: fn update_proposal_data(&self) -> BoxFuture<'_, anyhow::Result<()>>; - fn subnets_and_proposals(&self) -> BoxFuture<'_, anyhow::Result>>> { + fn subnets_and_proposals(&self) -> BoxFuture<'_, anyhow::Result>>> { Box::pin(async { let subnets = self.subnets().await?; @@ -174,14 +173,14 @@ where local_registry: LocalRegistry, network: Network, - subnets: RwLock>>>, - nodes: RwLock>>>, - operators: RwLock>>>, + subnets: RwLock>>>, + nodes: RwLock>>>, + operators: RwLock>>>, node_labels_guests: RwLock>>>, elected_guestos: RwLock>>>, elected_hostos: RwLock>>>, unassigned_nodes_replica_version: RwLock>>, - firewall_rule_set: RwLock>>>, + firewall_rule_set: RwLock>>>, no_sync: bool, proposal_agent: Arc, } @@ -236,17 +235,17 @@ pub trait LazyRegistryFamilyEntries { fn get_latest_version(&self) -> RegistryVersion; } -fn get_family_entries(reg: &impl LazyRegistryFamilyEntries) -> anyhow::Result> { +fn get_family_entries(reg: &impl LazyRegistryFamilyEntries) -> anyhow::Result> { let family = get_family_entries_versioned::(reg)?; Ok(family.into_iter().map(|(k, (_, v))| (k, v)).collect()) } -fn get_family_entries_versioned(reg: &impl LazyRegistryFamilyEntries) -> anyhow::Result> { +fn get_family_entries_versioned(reg: &impl LazyRegistryFamilyEntries) -> anyhow::Result> { get_family_entries_of_version(reg, reg.get_latest_version()) } fn get_family_entries_of_version( reg: &impl LazyRegistryFamilyEntries, version: RegistryVersion, -) -> anyhow::Result> { +) -> anyhow::Result> { let prefix_length = T::KEY_PREFIX.len(); Ok(reg .get_key_family(T::KEY_PREFIX, version)? @@ -406,7 +405,7 @@ impl LazyRegistry for LazyRegistryImpl { }) } - fn operators(&self) -> BoxFuture<'_, anyhow::Result>>> { + fn operators(&self) -> BoxFuture<'_, anyhow::Result>>> { Box::pin(async { if let Some(operators) = self.operators.read().await.as_ref() { return Ok(operators.to_owned()); @@ -417,11 +416,11 @@ impl LazyRegistry for LazyRegistryImpl { false => query_ic_dashboard_list::(&self.network, "v3/node-providers").await?, true => NodeProvidersResponse { node_providers: vec![] }, }; - let node_providers: BTreeMap<_, _> = node_providers.node_providers.iter().map(|p| (p.principal_id, p)).collect(); + let node_providers: IndexMap<_, _> = node_providers.node_providers.iter().map(|p| (p.principal_id, p)).collect(); let data_centers = get_family_entries::(self)?; let operators = get_family_entries::(self)?; - let records: BTreeMap<_, _> = operators + let records: IndexMap<_, _> = operators .iter() .map(|(p, or)| { let principal = PrincipalId::from_str(p).expect("Invalid operator principal id"); @@ -462,7 +461,7 @@ impl LazyRegistry for LazyRegistryImpl { longitude: dc.gps.clone().map(|l| l.longitude as f64), } }), - rewardable_nodes: or.rewardable_nodes.clone(), + rewardable_nodes: or.rewardable_nodes.iter().map(|(k, v)| (k.clone(), *v)).collect(), ipv6: or.ipv6().to_string(), }, ) @@ -475,7 +474,7 @@ impl LazyRegistry for LazyRegistryImpl { }) } - fn nodes(&self) -> BoxFuture<'_, anyhow::Result>>> { + fn nodes(&self) -> BoxFuture<'_, anyhow::Result>>> { Box::pin(async { if let Some(nodes) = self.nodes.read().await.as_ref() { return Ok(nodes.to_owned()); @@ -483,11 +482,11 @@ impl LazyRegistry for LazyRegistryImpl { let node_entries = get_family_entries::(self)?; let versioned_node_entries = get_family_entries_versioned::(self)?; - let dfinity_dcs = DFINITY_DCS.split(' ').map(|dc| dc.to_string().to_lowercase()).collect::>(); + let dfinity_dcs = DFINITY_DCS.split(' ').map(|dc| dc.to_string().to_lowercase()).collect::>(); let api_boundary_nodes = get_family_entries::(self)?; let guests = self.node_labels().await?; let operators = self.operators().await?; - let nodes: BTreeMap<_, _> = node_entries + let nodes: IndexMap<_, _> = node_entries .iter() .map(|(p, nr)| { let guest = Self::node_record_guest(guests.clone(), nr); @@ -587,7 +586,7 @@ impl LazyRegistry for LazyRegistryImpl { let bag = Arc::make_mut(arc_map); bag.insert(key.to_owned(), value.clone()); } else { - let mut all = BTreeMap::new(); + let mut all = IndexMap::new(); all.insert(key.to_owned(), value.clone()); *opt_arc_map = Some(Arc::new(all)); } @@ -596,7 +595,7 @@ impl LazyRegistry for LazyRegistryImpl { }) } - fn subnets(&self) -> BoxFuture<'_, anyhow::Result>>> { + fn subnets(&self) -> BoxFuture<'_, anyhow::Result>>> { Box::pin(async { if let Some(subnets) = self.subnets.read().await.as_ref() { return Ok(subnets.to_owned()); @@ -604,7 +603,7 @@ impl LazyRegistry for LazyRegistryImpl { let all_nodes = self.nodes().await?; - let subnets: BTreeMap<_, _> = get_family_entries::(self)? + let subnets: IndexMap<_, _> = get_family_entries::(self)? .iter() .enumerate() .map(|(i, (p, sr))| { @@ -672,7 +671,7 @@ impl LazyRegistry for LazyRegistryImpl { }) } - fn nodes_with_proposals(&self) -> BoxFuture<'_, anyhow::Result>>> { + fn nodes_with_proposals(&self) -> BoxFuture<'_, anyhow::Result>>> { Box::pin(async { let nodes = self.nodes().await?; if nodes.iter().any(|(_, n)| n.proposal.is_some()) { @@ -694,7 +693,7 @@ impl LazyRegistry for LazyRegistryImpl { let subnets = self.subnets().await?; let topology_proposals = self.proposal_agent.list_open_topology_proposals().await?; - let nodes: BTreeMap<_, _> = nodes + let nodes: IndexMap<_, _> = nodes .iter() .map(|(p, n)| { let proposal = topology_proposals @@ -706,7 +705,7 @@ impl LazyRegistry for LazyRegistryImpl { }) .collect(); - let subnets: BTreeMap<_, _> = subnets + let subnets: IndexMap<_, _> = subnets .iter() .map(|(p, s)| { let proposal = topology_proposals @@ -785,7 +784,7 @@ impl LazyRegistry for LazyRegistryImpl { .collect_vec()) } - fn get_node_rewards_table(&self) -> anyhow::Result> { + fn get_node_rewards_table(&self) -> anyhow::Result> { get_family_entries::(self).map_err(|e| anyhow::anyhow!("Couldn't get node rewards table: {:?}", e)) } @@ -852,7 +851,7 @@ impl SubnetQuerier for LazyRegistryImpl { let subnets = nodes .iter() .map(|n| reg_nodes.get(&n.id).and_then(|n| n.subnet_id)) - .collect::>(); + .collect::>(); if subnets.len() > 1 { return Err(NetworkError::IllegalRequest("Nodes don't belong to the same subnet".to_owned())); } @@ -924,21 +923,21 @@ mock! { fn sync_with_nns(&self) -> BoxFuture<'_, anyhow::Result<()>>; - fn operators(&self) -> BoxFuture<'_, anyhow::Result>>>; + fn operators(&self) -> BoxFuture<'_, anyhow::Result>>>; - fn nodes(&self) -> BoxFuture<'_, anyhow::Result>>>; + fn nodes(&self) -> BoxFuture<'_, anyhow::Result>>>; fn firewall_rule_set(&self, firewall_rule_scope: FirewallRulesScope) -> BoxFuture<'_, anyhow::Result>; - fn subnets(&self) -> BoxFuture<'_, anyhow::Result>>>; + fn subnets(&self) -> BoxFuture<'_, anyhow::Result>>>; - fn nodes_with_proposals(&self) -> BoxFuture<'_, anyhow::Result>>>; + fn nodes_with_proposals(&self) -> BoxFuture<'_, anyhow::Result>>>; fn unassigned_nodes_replica_version(&self) -> BoxFuture<'_, anyhow::Result>>; fn get_api_boundary_nodes(&self) -> anyhow::Result>; - fn get_node_rewards_table(&self) -> anyhow::Result>; + fn get_node_rewards_table(&self) -> anyhow::Result>; fn get_unassigned_nodes(&self) -> anyhow::Result>; diff --git a/rs/ic-management-backend/src/registry.rs b/rs/ic-management-backend/src/registry.rs index 1049510ea..19568ffd6 100644 --- a/rs/ic-management-backend/src/registry.rs +++ b/rs/ic-management-backend/src/registry.rs @@ -36,6 +36,7 @@ use ic_registry_local_store::{Changelog, ChangelogEntry, KeyMutation, LocalStore use ic_registry_nns_data_provider::registry::RegistryCanister; use ic_registry_subnet_type::SubnetType; use ic_types::PrincipalId; +use indexmap::{IndexMap, IndexSet}; use itertools::Itertools; use lazy_static::lazy_static; use log::{debug, error, info, warn}; @@ -43,14 +44,11 @@ use regex::Regex; use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; use std::convert::TryFrom; +use std::net::Ipv6Addr; use std::ops::Add; use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::Arc; -use std::{ - collections::{BTreeMap, BTreeSet}, - net::Ipv6Addr, -}; use tokio::sync::RwLock; use tokio::time::{sleep, Duration}; use url::Url; @@ -68,11 +66,11 @@ pub struct RegistryState { pub local_registry: Arc, version: u64, - subnets: BTreeMap, - nodes: BTreeMap, - operators: BTreeMap, + subnets: IndexMap, + nodes: IndexMap, + operators: IndexMap, node_labels_guests: Vec, - known_subnets: BTreeMap, + known_subnets: IndexMap, guestos_releases: ArtifactReleases, hostos_releases: ArtifactReleases, @@ -115,13 +113,13 @@ impl RegistryEntry for ApiBoundaryNodeRecord { } pub trait RegistryFamilyEntries { - fn get_family_entries(&self) -> Result>; - fn get_family_entries_versioned(&self) -> Result>; - fn get_family_entries_of_version(&self, version: RegistryVersion) -> Result>; + fn get_family_entries(&self) -> Result>; + fn get_family_entries_versioned(&self) -> Result>; + fn get_family_entries_of_version(&self, version: RegistryVersion) -> Result>; } impl RegistryFamilyEntries for LocalRegistry { - fn get_family_entries(&self) -> Result> { + fn get_family_entries(&self) -> Result> { let prefix_length = T::KEY_PREFIX.len(); Ok(self .get_key_family(T::KEY_PREFIX, self.get_latest_version())? @@ -131,14 +129,14 @@ impl RegistryFamilyEntries for LocalRegistry { .unwrap_or_else(|_| panic!("failed to get entry {} for type {}", key, std::any::type_name::())) .map(|v| (key[prefix_length..].to_string(), T::decode(v.as_slice()).expect("invalid registry value"))) }) - .collect::>()) + .collect::>()) } - fn get_family_entries_versioned(&self) -> Result> { + fn get_family_entries_versioned(&self) -> Result> { self.get_family_entries_of_version(self.get_latest_version()) } - fn get_family_entries_of_version(&self, version: RegistryVersion) -> Result> { + fn get_family_entries_of_version(&self, version: RegistryVersion) -> Result> { let prefix_length = T::KEY_PREFIX.len(); Ok(self .get_key_family(T::KEY_PREFIX, version)? @@ -155,7 +153,7 @@ impl RegistryFamilyEntries for LocalRegistry { }) .unwrap_or_else(|_| panic!("failed to get entry {} for type {}", key, std::any::type_name::())) }) - .collect::>()) + .collect::>()) } } @@ -218,9 +216,9 @@ impl RegistryState { network: network.clone(), local_registry, version: 0, - subnets: BTreeMap::::new(), - nodes: BTreeMap::new(), - operators: BTreeMap::new(), + subnets: IndexMap::::new(), + nodes: IndexMap::new(), + operators: IndexMap::new(), node_labels_guests: Vec::new(), guestos_releases: ArtifactReleases::new(Artifact::GuestOs), hostos_releases: ArtifactReleases::new(Artifact::HostOs), @@ -386,9 +384,9 @@ impl RegistryState { } fn update_operators(&mut self, providers: &[NodeProviderDetails]) -> Result<()> { - let providers = providers.iter().map(|p| (p.principal_id, p)).collect::>(); - let data_center_records: BTreeMap = self.local_registry.get_family_entries()?; - let operator_records: BTreeMap = self.local_registry.get_family_entries()?; + let providers = providers.iter().map(|p| (p.principal_id, p)).collect::>(); + let data_center_records: IndexMap = self.local_registry.get_family_entries()?; + let operator_records: IndexMap = self.local_registry.get_family_entries()?; self.operators = operator_records .iter() @@ -423,7 +421,7 @@ impl RegistryState { longitude: dc.gps.clone().map(|l| l.longitude as f64), } }), - rewardable_nodes: or.rewardable_nodes.clone(), + rewardable_nodes: or.rewardable_nodes.iter().map(|(k, v)| (k.clone(), *v)).collect(), ipv6: or.ipv6().to_string(), }, ) @@ -443,7 +441,7 @@ impl RegistryState { fn update_nodes(&mut self) -> Result<()> { let node_entries = self.local_registry.get_family_entries_versioned::()?; let dfinity_dcs = DFINITY_DCS.split(' ').map(|dc| dc.to_string().to_lowercase()).collect::>(); - let api_boundary_nodes: BTreeMap = self.local_registry.get_family_entries()?; + let api_boundary_nodes: IndexMap = self.local_registry.get_family_entries()?; self.nodes = node_entries .iter() @@ -607,15 +605,15 @@ impl RegistryState { self.version } - pub fn subnets(&self) -> BTreeMap { + pub fn subnets(&self) -> IndexMap { self.subnets.clone() } - pub fn nodes(&self) -> BTreeMap { + pub fn nodes(&self) -> IndexMap { self.nodes.clone() } - pub async fn nodes_with_proposals(&self) -> Result> { + pub async fn nodes_with_proposals(&self) -> Result> { let nodes = self.nodes.clone(); let proposal_agent = proposal::ProposalAgentImpl::new(self.network.get_nns_urls()); @@ -644,7 +642,7 @@ impl RegistryState { proposal_agent.list_open_elect_hostos_proposals().await } - pub async fn subnets_with_proposals(&self) -> Result> { + pub async fn subnets_with_proposals(&self) -> Result> { let subnets = self.subnets.clone(); let proposal_agent = proposal::ProposalAgentImpl::new(self.network.get_nns_urls()); @@ -697,8 +695,8 @@ impl RegistryState { async fn retireable_hostos_versions(&self) -> Result> { let active_releases = self.hostos_releases.get_active_branches(); - let hostos_versions: BTreeSet = self.nodes.values().map(|s| s.hostos_version.clone()).collect(); - let versions_in_proposals: BTreeSet = self + let hostos_versions: IndexSet = self.nodes.values().map(|s| s.hostos_version.clone()).collect(); + let versions_in_proposals: IndexSet = self .open_elect_hostos_proposals() .await? .iter() @@ -721,9 +719,9 @@ impl RegistryState { async fn retireable_guestos_versions(&self) -> Result> { let active_releases = self.guestos_releases.get_active_branches(); - let subnet_versions: BTreeSet = self.subnets.values().map(|s| s.replica_version.clone()).collect(); + let subnet_versions: IndexSet = self.subnets.values().map(|s| s.replica_version.clone()).collect(); let version_on_unassigned_nodes = self.get_unassigned_nodes_replica_version().await?; - let versions_in_proposals: BTreeSet = self + let versions_in_proposals: IndexSet = self .open_elect_replica_proposals() .await? .iter() @@ -754,7 +752,7 @@ impl RegistryState { ) } - pub fn operators(&self) -> BTreeMap { + pub fn operators(&self) -> IndexMap { self.operators.clone() } @@ -858,7 +856,7 @@ impl SubnetQuerier for RegistryState { .to_vec() .iter() .map(|n| self.nodes.get(&n.id).and_then(|n| n.subnet_id)) - .collect::>(); + .collect::>(); if subnets.len() > 1 { return Err(NetworkError::IllegalRequest("nodes don't belong to the same subnet".to_string())); } diff --git a/rs/ic-management-backend/src/release.rs b/rs/ic-management-backend/src/release.rs index 41b29185a..e1cda437a 100644 --- a/rs/ic-management-backend/src/release.rs +++ b/rs/ic-management-backend/src/release.rs @@ -6,10 +6,10 @@ use chrono::{Datelike, NaiveDate, NaiveTime, TimeZone, Utc, Weekday}; use futures_util::future::try_join_all; use ic_management_types::{Network, Release, Subnet}; use ic_types::PrincipalId; +use indexmap::{IndexMap, IndexSet}; use itertools::Itertools; use log::info; use serde::{Deserialize, Serialize}; -use std::collections::{BTreeMap, BTreeSet}; use std::str::FromStr; use strum_macros::{Display, EnumString}; @@ -144,7 +144,7 @@ impl RolloutConfig { pub struct RolloutBuilder { pub proposal_agent: ProposalAgentImpl, pub prometheus_client: prometheus_http_query::Client, - pub subnets: BTreeMap, + pub subnets: IndexMap, pub releases: Vec, pub network: Network, } @@ -323,8 +323,8 @@ impl RolloutBuilder { }); if let Some(last_stage) = stages.last_mut() { - let releases = last_stage.updates.iter().map(|u| u.replica_release.clone()).collect::>(); - let mut update_states = BTreeMap::new(); + let releases = last_stage.updates.iter().map(|u| u.replica_release.clone()).collect::>(); + let mut update_states = IndexMap::new(); for release in releases { let release_update_states = get_update_states(&self.network, &self.prometheus_client, &release, last_stage.start_date_time).await?; update_states.insert(release, release_update_states); @@ -526,7 +526,7 @@ async fn get_update_states( prometheus_client: &prometheus_http_query::Client, release: &Release, since: chrono::DateTime, -) -> Result> { +) -> Result> { const STATE_FIELD: &str = "state"; let query = format!( r#" @@ -624,7 +624,7 @@ pub async fn list_subnets_release_statuses( proposal_agent: &ProposalAgentImpl, prometheus_client: &prometheus_http_query::Client, network: Network, - subnets: BTreeMap, + subnets: IndexMap, releases: Vec, ) -> anyhow::Result> { let mut subnet_update_proposals = proposal_agent.list_update_subnet_version_proposals().await?; @@ -633,7 +633,7 @@ pub async fn list_subnets_release_statuses( let latest_updates = subnets .keys() .map(|p| (p, subnet_update_proposals.iter().find(|sup| sup.payload.subnet_id == *p))) - .collect::>(); + .collect::>(); let active_releases = latest_updates .values() @@ -658,7 +658,7 @@ pub async fn list_subnets_release_statuses( .unwrap_or_else(Utc::now), ) }) - .collect::>(); + .collect::>(); let updates_statuses_for_revision = try_join_all(active_releases.clone().into_iter().cloned().map(|r| async { let retryable = || async { @@ -677,7 +677,7 @@ pub async fn list_subnets_release_statuses( })) .await? .into_iter() - .collect::>(); + .collect::>(); let latest_releases = releases.iter().rev().unique_by(|r| r.branch.clone()).collect::>(); diff --git a/rs/ic-management-backend/src/subnets.rs b/rs/ic-management-backend/src/subnets.rs index b94bf88a1..9cf68a5ea 100644 --- a/rs/ic-management-backend/src/subnets.rs +++ b/rs/ic-management-backend/src/subnets.rs @@ -1,11 +1,10 @@ -use std::collections::BTreeMap; - use decentralization::{network::SubnetChange, SubnetChangeResponse}; use ic_base_types::PrincipalId; use ic_management_types::{Node, TopologyChangeProposal}; +use indexmap::IndexMap; pub fn get_proposed_subnet_changes( - all_nodes: &BTreeMap, + all_nodes: &IndexMap, subnet: &ic_management_types::Subnet, ) -> Result { if let Some(proposal) = &subnet.proposal { @@ -125,8 +124,8 @@ mod tests { assert_eq!(change.removed_with_desc.iter().map(|x| x.0).collect::>(), node_ids_removed); } - fn gen_test_nodes(subnet_id: PrincipalId, num_nodes: u64, start_at_number: u64) -> BTreeMap { - let mut nodes = BTreeMap::new(); + fn gen_test_nodes(subnet_id: PrincipalId, num_nodes: u64, start_at_number: u64) -> IndexMap { + let mut nodes = IndexMap::new(); for i in start_at_number..start_at_number + num_nodes { let node = Node { principal: PrincipalId::new_node_test_id(i), diff --git a/rs/ic-management-types/Cargo.toml b/rs/ic-management-types/Cargo.toml index 7f0889e71..799a8f05d 100644 --- a/rs/ic-management-types/Cargo.toml +++ b/rs/ic-management-types/Cargo.toml @@ -15,6 +15,7 @@ ic-base-types = { workspace = true } ic-nns-governance = { workspace = true } ic-registry-subnet-type = { workspace = true } ic-types = { workspace = true } +indexmap = { workspace = true } registry-canister = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } diff --git a/rs/np-notifications/Cargo.toml b/rs/np-notifications/Cargo.toml index 45aff305c..2f4d01ab3 100644 --- a/rs/np-notifications/Cargo.toml +++ b/rs/np-notifications/Cargo.toml @@ -16,6 +16,7 @@ clap = { workspace = true } ic-management-backend = { path = "../ic-management-backend" } ic-management-types = { path = "../ic-management-types" } ic-types = { workspace = true } +indexmap = { workspace = true } rand = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } diff --git a/rs/np-notifications/src/nodes_status.rs b/rs/np-notifications/src/nodes_status.rs index 6c327c5c2..5b3bb013b 100644 --- a/rs/np-notifications/src/nodes_status.rs +++ b/rs/np-notifications/src/nodes_status.rs @@ -1,27 +1,39 @@ -use std::collections::{BTreeMap, BTreeSet}; +use indexmap::{IndexMap, IndexSet}; use ic_management_types::HealthStatus; use ic_types::PrincipalId; use crate::notification::Notification; -#[derive(Debug, PartialOrd, Ord, PartialEq, Eq, Clone)] +#[derive(Debug, PartialEq, Eq, Clone)] pub struct NodesStatus { - nodes: BTreeMap, + nodes: IndexMap, } -impl From> for NodesStatus { - fn from(tree: BTreeMap) -> Self { +impl Ord for NodesStatus { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.nodes.iter().cmp(&other.nodes) + } +} + +impl PartialOrd for NodesStatus { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl From> for NodesStatus { + fn from(tree: IndexMap) -> Self { Self { nodes: tree } } } impl NodesStatus { pub fn _new() -> Self { - Self { nodes: BTreeMap::new() } + Self { nodes: IndexMap::new() } } - pub fn get_set_of_node_ids(&self) -> BTreeSet { + pub fn get_set_of_node_ids(&self) -> IndexSet { self.nodes.keys().copied().collect() } @@ -29,7 +41,7 @@ impl NodesStatus { self.nodes.get(&id) } - pub fn updated_from_map(&self, map: BTreeMap) -> (NodesStatus, Vec) { + pub fn updated_from_map(&self, map: IndexMap) -> (NodesStatus, Vec) { self.updated(Self::from(map)) } @@ -43,7 +55,7 @@ impl NodesStatus { let current_status_node_ids = self.get_set_of_node_ids(); let new_status_node_ids = new_statuses.get_set_of_node_ids(); - let added_nodes: BTreeSet = new_status_node_ids.difference(¤t_status_node_ids).cloned().collect(); + let added_nodes: IndexSet = new_status_node_ids.difference(¤t_status_node_ids).cloned().collect(); for node_id in added_nodes { notifications.push(Notification { @@ -59,7 +71,7 @@ impl NodesStatus { }) } - let removed_nodes: BTreeSet = current_status_node_ids.difference(&new_status_node_ids).cloned().collect(); + let removed_nodes: IndexSet = current_status_node_ids.difference(&new_status_node_ids).cloned().collect(); for node_id in removed_nodes { notifications.push(Notification { @@ -74,7 +86,7 @@ impl NodesStatus { }) } - let kept_nodes: BTreeSet = current_status_node_ids.intersection(&new_status_node_ids).cloned().collect(); + let kept_nodes: IndexSet = current_status_node_ids.intersection(&new_status_node_ids).cloned().collect(); for node_id in kept_nodes { if self.get(node_id) != new_statuses.get(node_id) { @@ -120,14 +132,14 @@ mod tests { ]; let statuses = NodesStatus { - nodes: BTreeMap::from([ + nodes: IndexMap::from([ (ids[0], HealthStatus::Healthy), (ids[1], HealthStatus::Healthy), (ids[2], HealthStatus::Healthy), ]), }; let new_statuses = NodesStatus { - nodes: BTreeMap::from([ + nodes: IndexMap::from([ (ids[0], HealthStatus::Healthy), (ids[1], HealthStatus::Degraded), (ids[3], HealthStatus::Healthy),