diff --git a/rs/registry/canister/src/mutations/node_management/do_remove_nodes.rs b/rs/registry/canister/src/mutations/node_management/do_remove_nodes.rs index 8312293aac1f..7a8e75aed91f 100644 --- a/rs/registry/canister/src/mutations/node_management/do_remove_nodes.rs +++ b/rs/registry/canister/src/mutations/node_management/do_remove_nodes.rs @@ -1,14 +1,15 @@ use crate::mutations::node_management::common::{ - find_subnet_for_node, get_node_operator_id_for_node, get_node_operator_record, - get_subnet_list_record, make_remove_node_registry_mutations, - make_update_node_operator_mutation, + find_subnet_for_node, get_node_operator_record, get_subnet_list_record, + make_remove_node_registry_mutations, make_update_node_operator_mutation, }; use crate::{common::LOG_PREFIX, registry::Registry}; -use candid::{CandidType, Deserialize}; +use candid::{CandidType, Deserialize, Principal}; #[cfg(target_arch = "wasm32")] use dfn_core::println; use ic_base_types::NodeId; +use ic_protobuf::registry::node_operator::v1::NodeOperatorRecord; use ic_registry_keys::make_subnet_record_key; +use ic_types::PrincipalId; use serde::Serialize; use std::collections::HashMap; @@ -19,7 +20,7 @@ impl Registry { // This hashmap tracks node operators for which mutations have already been // determined; increments to node allowance should not be idempotent - let mut node_operator_hmap = HashMap::::new(); + let mut node_operator_hmap = HashMap::::new(); // 1. De-duplicate the node list let mut nodes_to_be_removed = payload.node_ids.clone(); @@ -31,23 +32,22 @@ impl Registry { let subnet_list_record = get_subnet_list_record(self); // 3. Loop through each node - let mutations = nodes_to_be_removed + let mut mutations : Vec<_> = nodes_to_be_removed .into_iter().flat_map(|node_to_remove| { // 4. Skip nodes that are not in the registry. // This tackles the race condition where a node is removed from the registry // by another transaction before this transaction is processed. - if self.get_node(node_to_remove).is_none() { - println!("{LOG_PREFIX}do_remove_nodes: node {node_to_remove} not found in registry, skipping"); - return vec![]; + let node_record = match self.get_node(node_to_remove) { + Some(n) => n, + None => { + println!("{LOG_PREFIX}do_remove_nodes: node {node_to_remove} not found in registry, skipping"); + return vec![]; + } }; - // 5. Find the node operator id for this record - // and abort if the node record is not found - let node_operator_id = get_node_operator_id_for_node(self, node_to_remove) - .map_err(|e| format!("{LOG_PREFIX}do_remove_nodes: Aborting node removal: {e}")) - .unwrap(); + let node_operator_id = PrincipalId(Principal::from_slice(&node_record.node_operator_id)); - // 6. Ensure node is not in a subnet + // 5. Ensure node is not in a subnet let is_node_in_subnet = find_subnet_for_node(self, node_to_remove, &subnet_list_record); if let Some(subnet_id) = is_node_in_subnet { panic!("{}do_remove_nodes: Cannot remove a node that is a member of a subnet. This node is a member of Subnet: {}", @@ -56,41 +56,32 @@ impl Registry { ); } - // 7. Retrieve the NO record and increment its node allowance by 1 - let mut new_node_operator_record = get_node_operator_record(self, node_operator_id) + // 6. Retrieve the NO record, cache it and increment its node allowance by 1 + let new_node_operator_record = node_operator_hmap.entry(node_operator_id).or_insert(get_node_operator_record(self, node_operator_id) .map_err(|err| { format!( "{LOG_PREFIX}do_remove_nodes: Aborting node removal: {err}" ) }) - .unwrap(); + .unwrap()); + new_node_operator_record.node_allowance += 1; - // Use the hashmap to track whether the same NO has already been mutated in the same call - new_node_operator_record.node_allowance = match node_operator_hmap.get(&node_operator_id.to_string()) { - Some(n) => { - *n + 1 - } - None => { - new_node_operator_record.node_allowance + 1 - } - }; - node_operator_hmap.insert(node_operator_id.to_string(), new_node_operator_record.node_allowance); - // 8. Finally, generate the following mutations: + // 7. Finally, generate the following mutations: // * Delete the node // * Delete entries for node encryption keys - // * Increment NO's allowance by 1 - let mut mutations = make_remove_node_registry_mutations(self, node_to_remove); - // mutation to update node operator value - mutations.push(make_update_node_operator_mutation( - node_operator_id, - &new_node_operator_record, - )); - - mutations + make_remove_node_registry_mutations(self, node_to_remove) }).collect(); - // 8. Apply mutations after checking invariants + // 8. Create node operator update mutations + for (node_operator_id, updated_node_operator) in node_operator_hmap { + mutations.push(make_update_node_operator_mutation( + node_operator_id, + &updated_node_operator, + )); + } + + // 9. Apply mutations after checking invariants self.maybe_apply_mutation_internal(mutations); println!("{LOG_PREFIX}do_remove_nodes finished: {payload:?}"); @@ -221,6 +212,22 @@ mod tests { // Verify node operator allowance was incremented by 2 let updated_operator = get_node_operator_record(®istry, node_operator_id).unwrap(); assert_eq!(updated_operator.node_allowance, initial_allowance + 2); + + // Verify that there was only a single mutation in the latest version + let latest_version = registry.latest_version(); + let key = make_node_operator_record_key(node_operator_id); + let number_of_mutations_for_operator = registry + .store + .get(key.as_bytes()) + .map(|all_mutations| { + all_mutations + .iter() + .filter(|mutation| mutation.version == latest_version) + .count() + }) + .unwrap(); + + assert_eq!(number_of_mutations_for_operator, 1); } #[test] diff --git a/rs/registry/canister/unreleased_changelog.md b/rs/registry/canister/unreleased_changelog.md index 8d02ace0dd2c..b55d07a996d6 100644 --- a/rs/registry/canister/unreleased_changelog.md +++ b/rs/registry/canister/unreleased_changelog.md @@ -19,4 +19,6 @@ on the process that this file is part of, see ## Fixed +* When performing `RemoveNodes`, generate 1 update mutation per node operator key. Before this change, a single node operator record would be changed multiple times in a single version if the remove nodes proposal removed multiple nodes from the same node operator, which caused confusion. + ## Security