perf(nns): Read heap neurons for cardinalities

rs/nns/governance/canbench/canbench_results.yml

@@ -1,73 +1,73 @@
 benches:
   add_neuron_active_maximum:
     total:
-      instructions: 36179013
+      instructions: 36164619
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_active_typical:
     total:
-      instructions: 1835446
+      instructions: 1834736
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_maximum:
     total:
-      instructions: 96157759
+      instructions: 96124805
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_typical:
     total:
-      instructions: 7371051
+      instructions: 7369666
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_all_heap:
     total:
-      instructions: 34547039
+      instructions: 34451954
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_heap_neurons_stable_index:
     total:
-      instructions: 56681936
+      instructions: 56587595
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_stable_everything:
     total:
-      instructions: 371913075
+      instructions: 371815016
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_stable_neurons_with_heap_index:
     total:
-      instructions: 349727939
+      instructions: 349629136
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   centralized_following_all_stable:
     total:
-      instructions: 173821479
+      instructions: 174310560
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   compute_ballots_for_new_proposal_with_stable_neurons:
     total:
-      instructions: 1815411
+      instructions: 1815415
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   draw_maturity_from_neurons_fund_heap:
     total:
-      instructions: 7336319
+      instructions: 7233119
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   draw_maturity_from_neurons_fund_stable:
     total:
-      instructions: 10301600
+      instructions: 10300850
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -79,13 +79,13 @@ benches:
     scopes: {}
   list_active_neurons_fund_neurons_stable:
     total:
-      instructions: 2450275
+      instructions: 2450279
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   list_neurons_heap:
     total:
-      instructions: 3900688
+      instructions: 3900362
       heap_increase: 9
       stable_memory_increase: 0
     scopes: {}
@@ -97,13 +97,13 @@ benches:
     scopes: {}
   list_neurons_ready_to_unstake_maturity_stable:
     total:
-      instructions: 36937653
+      instructions: 36937657
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   list_neurons_stable:
     total:
-      instructions: 97962451
+      instructions: 97878953
       heap_increase: 5
       stable_memory_increase: 0
     scopes: {}
@@ -115,19 +115,19 @@ benches:
     scopes: {}
   list_ready_to_spawn_neuron_ids_stable:
     total:
-      instructions: 36926354
+      instructions: 36926358
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   neuron_data_validation_heap:
     total:
-      instructions: 531679423
+      instructions: 531768288
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   neuron_data_validation_stable:
     total:
-      instructions: 883798615
+      instructions: 667794511
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -139,25 +139,25 @@ benches:
     scopes: {}
   neuron_metrics_calculation_stable:
     total:
-      instructions: 2476865
+      instructions: 2476869
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   range_neurons_performance:
     total:
-      instructions: 47700442
+      instructions: 47700446
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   single_vote_all_stable:
     total:
-      instructions: 2474986
+      instructions: 2477911
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   update_recent_ballots_stable_memory:
     total:
-      instructions: 236963
+      instructions: 236980
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}

rs/nns/governance/src/neuron_data_validation.rs

@@ -1,5 +1,4 @@
 use crate::{
-    is_active_neurons_in_stable_memory_enabled,
     neuron::Neuron,
     neuron_store::NeuronStore,
     pb::v1::Topic,
@@ -472,10 +471,10 @@ impl CardinalityAndRangeValidator for PrincipalIndexValidator {
     const HEAP_NEURON_RANGE_CHUNK_SIZE: usize = 200;
 
     fn validate_cardinalities(neuron_store: &NeuronStore) -> Option<ValidationIssue> {
-        let cardinality_primary_heap: u64 = neuron_store.with_active_neurons_iter(|iter| {
-            iter.map(|neuron| neuron.principal_ids_with_special_permissions().len() as u64)
-                .sum()
-        });
+        let cardinality_primary_heap: u64 = neuron_store
+            .heap_neurons_iter()
+            .map(|neuron| neuron.principal_ids_with_special_permissions().len() as u64)
+            .sum();
         let cardinality_primary_stable = with_stable_neuron_store(|stable_neuron_store|
                     // `stable_neuron_store.len()` is for the controllers.
                     stable_neuron_store.lens().hot_keys + stable_neuron_store.len() as u64);
@@ -530,10 +529,10 @@ impl CardinalityAndRangeValidator for FollowingIndexValidator {
     const HEAP_NEURON_RANGE_CHUNK_SIZE: usize = 40;
 
     fn validate_cardinalities(neuron_store: &NeuronStore) -> Option<ValidationIssue> {
-        let cardinality_primary_heap: u64 = neuron_store.with_active_neurons_iter(|iter| {
-            iter.map(|neuron| neuron.topic_followee_pairs().len() as u64)
-                .sum()
-        });
+        let cardinality_primary_heap: u64 = neuron_store
+            .heap_neurons_iter()
+            .map(|neuron| neuron.topic_followee_pairs().len() as u64)
+            .sum();
         let cardinality_primary_stable =
             with_stable_neuron_store(|stable_neuron_store| stable_neuron_store.lens().followees);
         let cardinality_primary = cardinality_primary_heap + cardinality_primary_stable;
@@ -585,20 +584,14 @@ impl CardinalityAndRangeValidator for KnownNeuronIndexValidator {
     // entry lookup takes ~130K instructions.
     const HEAP_NEURON_RANGE_CHUNK_SIZE: usize = 300000;
     fn validate_cardinalities(neuron_store: &NeuronStore) -> Option<ValidationIssue> {
-        let cardinality_active_neurons = neuron_store.with_active_neurons_iter(|iter| {
-            iter.filter(|neuron| neuron.known_neuron_data.is_some())
-                .count() as u64
-        });
-        let cardinality_stable_neurons = with_stable_neuron_store(|stable_neuron_store| {
+        let cardinality_primary_heap = neuron_store
+            .heap_neurons_iter()
+            .filter(|neuron| neuron.known_neuron_data.is_some())
+            .count() as u64;
+        let cardinality_primary_stable = with_stable_neuron_store(|stable_neuron_store| {
             stable_neuron_store.lens().known_neuron_data
         });
-        // NOTE - this will not be completely correct during the migration of active
-        // heap neurons to stable storage, but it will self-correct after the migration is finished.
-        let cardinality_primary = if is_active_neurons_in_stable_memory_enabled() {
-            cardinality_stable_neurons
-        } else {
-            cardinality_active_neurons + cardinality_stable_neurons
-        };
+        let cardinality_primary = cardinality_primary_heap + cardinality_primary_stable;
         let cardinality_index =
             with_stable_neuron_indexes(|indexes| indexes.known_neuron().num_entries()) as u64;
         if cardinality_primary != cardinality_index {
@@ -682,6 +675,7 @@ mod tests {
     use maplit::{btreemap, hashmap};
 
     use crate::{
+        is_active_neurons_in_stable_memory_enabled,
         neuron::{DissolveStateAndAge, NeuronBuilder},
         pb::v1::{neuron::Followees, KnownNeuronData},
         storage::{with_stable_neuron_indexes_mut, with_stable_neuron_store_mut},

rs/nns/governance/src/neuron_store.rs

@@ -934,6 +934,14 @@ impl NeuronStore {
         self.heap_neurons.range(range).map(|(_, neuron)| neuron)
     }
 
+    // TODO remove this after we no longer need to validate neurons in heap.
+    /// Returns an iterator over all neurons in the heap. There is no guarantee that the active
+    /// neurons are all in the heap as they are being migrated to stable memory, so the caller
+    /// should be aware of different storage locations.
+    pub fn heap_neurons_iter(&self) -> impl Iterator<Item = &Neuron> {
+        self.heap_neurons.values()
+    }
+
     fn is_active_neurons_fund_neuron(neuron: &Neuron, now: u64) -> bool {
         !neuron.is_inactive(now) && neuron.is_a_neurons_fund_member()
     }