diff --git a/rs/cli/src/runner.rs b/rs/cli/src/runner.rs index f340f7dc1..2e2eb7ae6 100644 --- a/rs/cli/src/runner.rs +++ b/rs/cli/src/runner.rs @@ -20,14 +20,15 @@ use log::{info, warn}; use registry_canister::mutations::do_change_subnet_membership::ChangeSubnetMembershipPayload; use std::cell::RefCell; use std::collections::BTreeMap; -use std::sync::mpsc; +use std::sync::{mpsc, Arc}; use std::thread; use tabled::builder::Builder; use tabled::settings::Style; pub struct Runner { pub ic_admin: ic_admin::IcAdminWrapper, - registry: RegistryState, + registry: RefCell>>, + network: Network, dashboard_backend_client: RefCell>, backend_srv: RefCell>, } @@ -43,7 +44,7 @@ impl Runner { let backend_url = format!("http://localhost:{}/", backend_port); let (tx, rx) = mpsc::channel(); - let target_network_backend = self.registry.network(); + let target_network_backend = self.registry().await.network(); thread::spawn(move || { let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async move { @@ -71,16 +72,41 @@ impl Runner { Ok(()) } - pub async fn new(ic_admin: ic_admin::IcAdminWrapper, network: &Network) -> anyhow::Result { - let mut registry = registry::RegistryState::new(network, true).await; - let node_providers = query_ic_dashboard_list::(network, "v3/node-providers") - .await? + pub async fn registry(&self) -> Arc { + { + if let Some(ref registry) = *self.registry.borrow() { + return Arc::clone(registry); + } + } + + // Create a new registry state + let new_registry = Arc::new(registry::RegistryState::new(&self.network, true).await); + + // Fetch node providers + let node_providers = query_ic_dashboard_list::(&self.network, "v3/node-providers") + .await + .expect("Failed to get node providers") .node_providers; - registry.update_node_details(&node_providers).await?; + // Update node details + Arc::get_mut(&mut Arc::clone(&new_registry)) + .expect("Failed to get mutable reference") + .update_node_details(&node_providers) + .await + .expect("Failed to update node details"); + + // Replace the registry in self with the new registry state + self.registry.replace(Some(Arc::clone(&new_registry))); + + // Return the Arc to the new registry state + new_registry + } + + pub async fn new(ic_admin: ic_admin::IcAdminWrapper, network: &Network) -> anyhow::Result { Ok(Self { ic_admin, - registry, + registry: RefCell::new(None), + network: network.clone(), dashboard_backend_client: RefCell::new(None), backend_srv: RefCell::new(None), }) @@ -317,7 +343,7 @@ impl Runner { version: &String, exclude: &Option>, ) -> anyhow::Result, String)>> { - let elected_versions = self.registry.blessed_versions(&Artifact::HostOs).await.unwrap(); + let elected_versions = self.registry().await.blessed_versions(&Artifact::HostOs).await.unwrap(); if !elected_versions.contains(&version.to_string()) { return Err(anyhow::anyhow!(format!( "The version {} has not being elected.\nVersions elected are: {:?}", @@ -325,10 +351,10 @@ impl Runner { ))); } let hostos_rollout = HostosRollout::new( - self.registry.nodes(), - self.registry.subnets(), - &self.registry.network(), - ProposalAgent::new(self.registry.get_nns_urls()), + self.registry().await.nodes(), + self.registry().await.subnets(), + &self.registry().await.network(), + ProposalAgent::new(self.registry().await.get_nns_urls()), version, exclude, ); @@ -521,11 +547,16 @@ impl Runner { pub async fn decentralization_change(&self, change: &ChangeSubnetMembershipPayload) -> Result<(), anyhow::Error> { if let Some(id) = change.get_subnet() { - let subnet_before = self.registry.subnet(SubnetQueryBy::SubnetId(id)).await.map_err(|e| anyhow::anyhow!(e))?; + let subnet_before = self + .registry() + .await + .subnet(SubnetQueryBy::SubnetId(id)) + .await + .map_err(|e| anyhow::anyhow!(e))?; let nodes_before = subnet_before.nodes.clone(); - let added_nodes = self.registry.get_decentralized_nodes(&change.get_added_node_ids()); - let removed_nodes = self.registry.get_decentralized_nodes(&change.get_added_node_ids()); + let added_nodes = self.registry().await.get_decentralized_nodes(&change.get_added_node_ids()); + let removed_nodes = self.registry().await.get_decentralized_nodes(&change.get_added_node_ids()); let subnet_after = subnet_before .with_nodes(added_nodes) @@ -545,7 +576,8 @@ impl Runner { pub async fn subnet_rescue(&self, subnet: &PrincipalId, keep_nodes: Option>, dry_run: bool) -> anyhow::Result<()> { let change_request = self - .registry + .registry() + .await .modify_subnet_nodes(SubnetQueryBy::SubnetId(*subnet)) .await .map_err(|e| anyhow::anyhow!(e))?;