feat: Update registry helpers away from blessed versions

rs/registry/admin/bin/main.rs

@@ -81,23 +81,23 @@ use ic_protobuf::registry::{
     node_operator::v1::NodeOperatorRecord,
     node_rewards::v2::{NodeRewardRate, UpdateNodeRewardsTableProposalPayload},
     provisional_whitelist::v1::ProvisionalWhitelist as ProvisionalWhitelistProto,
-    replica_version::v1::{BlessedReplicaVersions, ReplicaVersionRecord},
+    replica_version::v1::ReplicaVersionRecord,
     routing_table::v1::CanisterMigrations,
     subnet::v1::{SubnetListRecord, SubnetRecord as SubnetRecordProto},
     unassigned_nodes_config::v1::UnassignedNodesConfigRecord,
 };
 use ic_registry_client::client::RegistryClientImpl;
 use ic_registry_client_helpers::{
     chain_keys::ChainKeysRegistry, crypto::CryptoRegistry, deserialize_registry_value,
-    ecdsa_keys::EcdsaKeysRegistry, hostos_version::HostosRegistry, subnet::SubnetRegistry,
+    ecdsa_keys::EcdsaKeysRegistry, hostos_version::HostosRegistry,
+    replica_version::ReplicaVersionRegistry, subnet::SubnetRegistry,
 };
 use ic_registry_keys::{
     API_BOUNDARY_NODE_RECORD_KEY_PREFIX, FirewallRulesScope, NODE_OPERATOR_RECORD_KEY_PREFIX,
     NODE_RECORD_KEY_PREFIX, NODE_REWARDS_TABLE_KEY, ROOT_SUBNET_ID_KEY,
     get_node_operator_id_from_record_key, get_node_record_node_id, is_node_operator_record_key,
-    is_node_record_key, make_api_boundary_node_record_key, make_blessed_replica_versions_key,
-    make_canister_migrations_record_key, make_crypto_node_key,
-    make_crypto_threshold_signing_pubkey_key, make_crypto_tls_cert_key,
+    is_node_record_key, make_api_boundary_node_record_key, make_canister_migrations_record_key,
+    make_crypto_node_key, make_crypto_threshold_signing_pubkey_key, make_crypto_tls_cert_key,
     make_data_center_record_key, make_firewall_config_record_key, make_firewall_rules_record_key,
     make_node_operator_record_key, make_node_record_key, make_provisional_whitelist_record_key,
     make_replica_version_key, make_subnet_list_record_key, make_subnet_record_key,
@@ -740,7 +740,7 @@ struct GetGuestOsVersionCmd {
 }
 
 /// Sub-command to submit a proposal to upgrade the replicas running a specific
-/// subnet to the given (blessed) version.
+/// subnet to the given (elected) version.
 #[derive_common_proposal_fields]
 #[derive(Parser, ProposalMetadata)]
 struct ProposeToDeployGuestosToAllSubnetNodesCmd {
@@ -4892,12 +4892,28 @@ async fn main() {
             .await;
         }
         SubCommand::GetElectedGuestosVersions => {
-            print_and_get_last_value::<BlessedReplicaVersions>(
-                make_blessed_replica_versions_key().as_bytes().to_vec(),
-                &registry_canister,
-                opts.json,
-            )
-            .await;
+            let registry_client = RegistryClientImpl::new(
+                Arc::new(NnsDataProvider::new(
+                    tokio::runtime::Handle::current(),
+                    reachable_nns_urls.clone(),
+                )),
+                None,
+            );
+
+            // maximum number of retries, let the user ctrl+c if necessary
+            registry_client
+                .try_polling_latest_version(usize::MAX)
+                .unwrap();
+
+            let guestos_versions = registry_client
+                .get_all_replica_version_records(registry_client.get_latest_version())
+                .unwrap();
+
+            if let Some(guestos_versions) = guestos_versions {
+                for (version, _) in guestos_versions {
+                    println!("{}", version);
+                }
+            }
         }
         SubCommand::GetRoutingTable(cmd) => {
             let registry_version = cmd.registry_version.map(RegistryVersion::from);

rs/registry/helpers/BUILD.bazel

@@ -40,6 +40,7 @@ rust_test(
         "//rs/registry/provisional_whitelist",
         "//rs/registry/routing_table",
         "//rs/registry/subnet_features",
+        "//rs/registry/transport",
         "//rs/types/base_types",
         "//rs/types/management_canister_types",
         "//rs/types/types",
@@ -65,6 +66,7 @@ rust_test_suite(
         "//rs/registry/provisional_whitelist",
         "//rs/registry/routing_table",
         "//rs/registry/subnet_features",
+        "//rs/registry/transport",
         "//rs/types/base_types",
         "//rs/types/management_canister_types",
         "//rs/types/types",

rs/registry/helpers/src/replica_version.rs

@@ -6,10 +6,10 @@ pub use ic_types::replica_version::ReplicaVersion;
 pub use ic_types::{NodeId, RegistryVersion, SubnetId};
 
 pub trait ReplicaVersionRegistry {
-    fn get_replica_versions(
+    fn get_all_replica_version_records(
         &self,
         version: RegistryVersion,
-    ) -> RegistryClientResult<Vec<ReplicaVersionRecord>>;
+    ) -> RegistryClientResult<Vec<(String, ReplicaVersionRecord)>>;
 
     fn get_replica_version_record(
         &self,
@@ -29,18 +29,22 @@ pub trait ReplicaVersionRegistry {
 }
 
 impl<T: RegistryClient + ?Sized> ReplicaVersionRegistry for T {
-    fn get_replica_versions(
+    fn get_all_replica_version_records(
         &self,
         version: RegistryVersion,
-    ) -> RegistryClientResult<Vec<ReplicaVersionRecord>> {
+    ) -> RegistryClientResult<Vec<(String, ReplicaVersionRecord)>> {
+        // Note this `get_key_family` impl does not strip the prefix from keys. The impl in the registry canister, does.
         let keys = self.get_key_family(REPLICA_VERSION_KEY_PREFIX, version)?;
 
         let mut records = Vec::new();
         for key in keys {
             let bytes = self.get_value(&key, version);
             let replica_version_proto =
                 deserialize_registry_value::<ReplicaVersionRecord>(bytes)?.unwrap_or_default();
-            records.push(replica_version_proto)
+            records.push((
+                key[REPLICA_VERSION_KEY_PREFIX.len()..].to_string(),
+                replica_version_proto,
+            ))
         }
 
         Ok(Some(records))
@@ -60,13 +64,13 @@ impl<T: RegistryClient + ?Sized> ReplicaVersionRegistry for T {
         version: RegistryVersion,
     ) -> Result<Vec<Vec<u8>>, String> {
         let replica_versions = self
-            .get_replica_versions(version)
+            .get_all_replica_version_records(version)
             .map_err(|err| format!("Failed to get replica versions: {err}"))?
-            .ok_or_else(|| "Blessed replica versions not found in registry".to_string())?;
+            .ok_or_else(|| "Elected replica versions not found in registry".to_string())?;
 
         let measurements = replica_versions
             .into_iter()
-            .flat_map(|record| {
+            .flat_map(|(_, record)| {
                 record
                     .guest_launch_measurements
                     .unwrap_or_default()

rs/tests/driver/BUILD.bazel

@@ -105,7 +105,6 @@ rust_library(
         "@crate_index//:macaddr",
         "@crate_index//:nix",
         "@crate_index//:num_cpus",
-        "@crate_index//:prost",
         "@crate_index//:rand",
         "@crate_index//:rand_chacha",
         "@crate_index//:regex",

rs/tests/driver/Cargo.toml

@@ -84,7 +84,6 @@ nix = { workspace = true }
 num_cpus = "1.13.1"
 on_wire = { path = "../../rust_canisters/on_wire" }
 phantom_newtype = { path = "../../phantom_newtype" }
-prost = { workspace = true }
 rand = { workspace = true }
 rand_chacha = { workspace = true }
 regex = { workspace = true }

rs/tests/driver/src/driver/test_env_api.rs

@@ -167,17 +167,15 @@ use ic_nns_test_utils::{
 };
 use ic_prep_lib::prep_state_directory::IcPrepStateDir;
 use ic_protobuf::registry::{
-    node::v1 as pb_node,
-    replica_version::v1::{BlessedReplicaVersions, ReplicaVersionRecord},
-    subnet::v1 as pb_subnet,
+    node::v1 as pb_node, replica_version::v1::ReplicaVersionRecord, subnet::v1 as pb_subnet,
 };
 use ic_registry_client_helpers::{
     api_boundary_node::ApiBoundaryNodeRegistry,
     node::NodeRegistry,
+    replica_version::ReplicaVersionRegistry,
     routing_table::RoutingTableRegistry,
     subnet::{SubnetListRegistry, SubnetRegistry},
 };
-use ic_registry_keys::REPLICA_VERSION_KEY_PREFIX;
 use ic_registry_local_registry::LocalRegistry;
 use ic_registry_routing_table::CanisterIdRange;
 use ic_registry_subnet_type::SubnetType;
@@ -188,8 +186,6 @@ use ic_types::{
 };
 use ic_utils::interfaces::ManagementCanister;
 use icp_ledger::{AccountIdentifier, LedgerCanisterInitPayload, Tokens};
-use itertools::Itertools;
-use prost::Message;
 use registry_canister::init::{RegistryCanisterInitPayload, RegistryCanisterInitPayloadBuilder};
 use serde::{Deserialize, Serialize};
 use slog::{Logger, debug, info, warn};
@@ -531,61 +527,14 @@ impl TopologySnapshot {
         )
     }
 
-    pub fn elected_replica_versions(&self) -> anyhow::Result<Vec<String>> {
-        Ok(self
-            .local_registry
-            .get_key_family(
-                "blessed_replica_versions",
-                self.local_registry.get_latest_version(),
-            )
-            .map_err(anyhow::Error::from)?
-            .iter()
-            .filter_map(|key| {
-                let r = self
-                    .local_registry
-                    .get_versioned_value(key, self.local_registry.get_latest_version())
-                    .unwrap_or_else(|_| {
-                        panic!("Failed to get entry {key} for blessed replica versions")
-                    });
-
-                r.as_ref().map(|v| {
-                    BlessedReplicaVersions::decode(v.as_slice()).expect("Invalid registry value")
-                })
-            })
-            .collect_vec()
-            .first()
-            .ok_or(anyhow::anyhow!(
-                "Failed to find any blessed replica versions"
-            ))?
-            .blessed_version_ids
-            .clone())
-    }
+    pub fn replica_version_records(&self) -> Result<Vec<(String, ReplicaVersionRecord)>> {
+        let registry_version = self.local_registry.get_latest_version();
 
-    pub fn replica_version_records(&self) -> anyhow::Result<Vec<(String, ReplicaVersionRecord)>> {
         Ok(self
             .local_registry
-            .get_key_family(
-                REPLICA_VERSION_KEY_PREFIX,
-                self.local_registry.get_latest_version(),
-            )
-            .map_err(anyhow::Error::from)?
-            .iter()
-            .map(|key| {
-                let r = self
-                    .local_registry
-                    .get_versioned_value(key, self.local_registry.get_latest_version())
-                    .unwrap_or_else(|_| panic!("Failed to get entry for replica version {key}"));
-                (
-                    key[REPLICA_VERSION_KEY_PREFIX.len()..].to_string(),
-                    r.as_ref()
-                        .map(|v| {
-                            ReplicaVersionRecord::decode(v.as_slice())
-                                .expect("Invalid registry value")
-                        })
-                        .unwrap(),
-                )
-            })
-            .collect_vec())
+            .get_all_replica_version_records(registry_version)
+            // get_all_replica_version_records always returns Some
+            .map(|v| v.unwrap())?)
     }
 
     /// The subnet id of the root subnet.
@@ -1398,27 +1347,21 @@ impl<T: HasTopologySnapshot> GetFirstHealthyNodeSnapshot for T {
         })
     }
     fn get_first_healthy_nns_node_snapshot(&self) -> IcNodeSnapshot {
-        let root_subnet_id = get_root_subnet_id_from_snapshot(self);
+        let root_subnet_id = self.topology_snapshot().root_subnet_id();
         self.get_first_healthy_node_snapshot_where(|s| s.subnet_id == root_subnet_id)
     }
     fn get_first_healthy_non_nns_node_snapshot(&self) -> IcNodeSnapshot {
-        let root_subnet_id = get_root_subnet_id_from_snapshot(self);
+        let root_subnet_id = self.topology_snapshot().root_subnet_id();
         self.get_first_healthy_node_snapshot_where(|s| s.subnet_id != root_subnet_id)
     }
     fn get_first_healthy_system_but_not_nns_node_snapshot(&self) -> IcNodeSnapshot {
-        let root_subnet_id = get_root_subnet_id_from_snapshot(self);
+        let root_subnet_id = self.topology_snapshot().root_subnet_id();
         self.get_first_healthy_node_snapshot_where(|s| {
             s.subnet_type() == SubnetType::System && s.subnet_id != root_subnet_id
         })
     }
 }
 
-fn get_root_subnet_id_from_snapshot<T: HasTopologySnapshot>(env: &T) -> SubnetId {
-    let ts = env.topology_snapshot();
-    ts.local_registry
-        .get_root_subnet_id(ts.registry_version)
-        .unwrap_result(ts.registry_version, "root_subnet_id")
-}
 pub trait HasRegistryLocalStore {
     fn registry_local_store_path(&self, name: &str) -> Option<PathBuf>;
 }