Skip to content

Commit

Permalink
feat(cli): improve node provider queries with info from governance ca…
Browse files Browse the repository at this point in the history
…nister (#1175)
  • Loading branch information
sasa-tomic authored Jan 7, 2025
1 parent 8c31aa8 commit 8608c99
Show file tree
Hide file tree
Showing 7 changed files with 67 additions and 6 deletions.
10 changes: 9 additions & 1 deletion Cargo.Bazel.lock
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
{
"checksum": "9b114dd1f58e0ade31049e0ecd6bbc1aed942d2d1d326f2eb31404064ca5d87a",
"checksum": "f057248fe3333221bba8683e81adaa371ced055af99b77f49deffbb1ddd79d2e",
"crates": {
"actix-codec 0.5.2": {
"name": "actix-codec",
Expand Down Expand Up @@ -9761,6 +9761,10 @@
"id": "ic-types 0.9.0",
"target": "ic_types"
},
{
"id": "icp-ledger 0.9.0",
"target": "icp_ledger"
},
{
"id": "indexmap 2.7.0",
"target": "indexmap"
Expand Down Expand Up @@ -15298,6 +15302,10 @@
"id": "ic-nns-governance 0.9.0",
"target": "ic_nns_governance"
},
{
"id": "ic-nns-governance-api 0.9.0",
"target": "ic_nns_governance_api"
},
{
"id": "ic-protobuf 0.9.0",
"target": "ic_protobuf"
Expand Down
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,7 @@ ic-sns-wasm = { git = "https://github.com/dfinity/ic.git", rev = "a700085f0575a4
cycles-minting-canister = { git = "https://github.com/dfinity/ic.git", rev = "a700085f0575a440cb44941da69cfa2efeea3ab9" }
ic-icrc1-test-utils = { git = "https://github.com/dfinity/ic.git", rev = "a700085f0575a440cb44941da69cfa2efeea3ab9" }
rosetta-core = { git = "https://github.com/dfinity/ic.git", rev = "a700085f0575a440cb44941da69cfa2efeea3ab9" }
icp-ledger = { git = "https://github.com/dfinity/ic.git", rev = "a700085f0575a440cb44941da69cfa2efeea3ab9" }
icrc-ledger-types = { git = "https://github.com/dfinity/ic.git", rev = "a700085f0575a440cb44941da69cfa2efeea3ab9" }
ic-transport-types = "0.39.1"
ic-utils = "0.39.0"
Expand Down
1 change: 1 addition & 0 deletions rs/cli/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ ic-registry-subnet-type = { workspace = true }
ic-sns-wasm = { workspace = true }
ic-sys = { workspace = true }
ic-types = { workspace = true }
icp-ledger = { workspace = true }
indexmap = { workspace = true }
itertools = { workspace = true }
keyring = { workspace = true }
Expand Down
48 changes: 43 additions & 5 deletions rs/cli/src/commands/registry.rs
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
use std::{
collections::{BTreeMap, HashMap},
collections::{BTreeMap, HashMap, HashSet},
net::Ipv6Addr,
path::PathBuf,
str::FromStr,
sync::Arc,
};

use clap::Args;
use ic_canisters::governance::GovernanceCanisterWrapper;
use ic_canisters::IcAgentCanisterClient;
use ic_management_backend::{health::HealthStatusQuerier, lazy_registry::LazyRegistry};
use ic_management_types::{HealthStatus, Network};
use ic_protobuf::registry::{
Expand All @@ -19,6 +21,7 @@ use ic_protobuf::registry::{
};
use ic_registry_subnet_type::SubnetType;
use ic_types::PrincipalId;
use icp_ledger::AccountIdentifier;
use indexmap::IndexMap;
use itertools::Itertools;
use log::{info, warn};
Expand Down Expand Up @@ -160,7 +163,7 @@ impl Registry {
node_operators: node_operators.values().cloned().collect_vec(),
node_rewards_table,
api_bns,
node_providers: get_node_providers(&local_registry).await?,
node_providers: get_node_providers(&local_registry, ctx.network()).await?,
})
}
}
Expand Down Expand Up @@ -417,24 +420,58 @@ fn get_api_boundary_nodes(local_registry: &Arc<dyn LazyRegistry>) -> anyhow::Res
Ok(api_bns)
}

async fn get_node_providers(local_registry: &Arc<dyn LazyRegistry>) -> anyhow::Result<Vec<NodeProvider>> {
async fn get_node_providers(local_registry: &Arc<dyn LazyRegistry>, network: &Network) -> anyhow::Result<Vec<NodeProvider>> {
let all_nodes = local_registry.nodes().await?;
let node_providers = local_registry

// Get the node providers from the node operator records, and from the governance canister, and merge them
let nns_urls = network.get_nns_urls();
let url = nns_urls.first().ok_or(anyhow::anyhow!("No NNS URLs provided"))?.to_owned();
let canister_client = IcAgentCanisterClient::from_anonymous(url)?;
let gov = GovernanceCanisterWrapper::from(canister_client);
let gov_node_providers: HashMap<PrincipalId, String> = gov
.get_node_providers()
.await?
.iter()
.map(|p| {
(
p.id.unwrap_or(PrincipalId::new_anonymous()),
match &p.reward_account {
Some(account) => AccountIdentifier::from_slice(&account.hash).unwrap().to_string(),
None => "".to_string(),
},
)
})
.collect();
let mut reg_node_providers = local_registry
.operators()
.await?
.values()
.map(|operator| operator.provider.clone())
.collect_vec();
let reg_provider_ids = reg_node_providers.iter().map(|provider| provider.principal).collect::<HashSet<_>>();
for principal in gov_node_providers.keys() {
if !reg_provider_ids.contains(principal) {
reg_node_providers.push(ic_management_types::Provider {
principal: *principal,
name: None,
website: None,
});
}
}
let reg_node_providers = reg_node_providers
.into_iter()
.sorted_by_key(|provider| provider.principal)
.dedup_by(|x, y| x.principal == y.principal)
.collect_vec();

Ok(node_providers
Ok(reg_node_providers
.iter()
.map(|provider| {
let provider_nodes = all_nodes.values().filter(|node| node.operator.provider.principal == provider.principal);

NodeProvider {
principal: provider.principal,
reward_account: gov_node_providers.get(&provider.principal).cloned().unwrap_or_default(),
total_nodes: provider_nodes.clone().count(),
nodes_in_subnet: provider_nodes.clone().filter(|node| node.subnet_id.is_some()).count(),
nodes_per_dc: provider_nodes
Expand Down Expand Up @@ -562,6 +599,7 @@ pub struct NodeRewardsTableFlattened {
struct NodeProvider {
name: String,
principal: PrincipalId,
reward_account: String,
total_nodes: usize,
nodes_in_subnet: usize,
nodes_per_dc: HashMap<String, usize>,
Expand Down
1 change: 1 addition & 0 deletions rs/ic-canisters/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ ic-management-canister-types = { workspace = true }
ic-nns-common = { workspace = true }
ic-nns-constants = { workspace = true }
ic-nns-governance = { workspace = true }
ic-nns-governance-api = { workspace = true }
ic-protobuf = { workspace = true }
ic-registry-transport = { workspace = true }
ic-utils = { workspace = true }
Expand Down
10 changes: 10 additions & 0 deletions rs/ic-canisters/src/governance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@ use ic_nns_governance::pb::v1::ManageNeuron;
use ic_nns_governance::pb::v1::ManageNeuronResponse;
use ic_nns_governance::pb::v1::Neuron;
use ic_nns_governance::pb::v1::NeuronInfo;
use ic_nns_governance::pb::v1::NodeProvider as PbNodeProvider;
use ic_nns_governance::pb::v1::ProposalInfo;
use ic_nns_governance_api::pb::v1::ListNodeProvidersResponse;
use serde::{self, Serialize};
use std::str::FromStr;
use std::time::Duration;
Expand Down Expand Up @@ -207,6 +209,14 @@ impl GovernanceCanisterWrapper {
.await
}

pub async fn get_node_providers(&self) -> anyhow::Result<Vec<PbNodeProvider>> {
let response = self
.query::<ListNodeProvidersResponse>("list_node_providers", candid::encode_one(())?)
.await?;
let node_providers = response.node_providers.into_iter().map(PbNodeProvider::from).collect();
Ok(node_providers)
}

async fn query<T>(&self, method_name: &str, args: Vec<u8>) -> anyhow::Result<T>
where
T: candid::CandidType + for<'a> candid::Deserialize<'a>,
Expand Down

0 comments on commit 8608c99

Please sign in to comment.