fix(nns): Avoid cloning heap_neurons to avoid performance penalty (#2726

)

This PR fixes a regression in a previous change that causes performance
for heartbeats to drop while still using heap neurons. This fixes the
method to use Cow instead, so that heap neurons are not unnecessarily
cloned.

rs/nns/governance/src/governance/tla/mod.rs

@@ -1,6 +1,7 @@
 use itertools::Itertools;
 use std::{
     collections::{BTreeMap, BTreeSet},
+    ops::Deref,
     thread,
 };
 
@@ -33,7 +34,8 @@ pub use claim_neuron::claim_neuron_desc;
 fn neuron_global(gov: &Governance) -> TlaValue {
     let neuron_map: BTreeMap<u64, TlaValue> = with_stable_neuron_store(|store| {
         gov.neuron_store.with_active_neurons_iter(|iter| {
-            iter.chain(store.range_neurons(std::ops::RangeFull))
+            iter.map(|n| n.deref().clone())
+                .chain(store.range_neurons(std::ops::RangeFull))
                 .map(|neuron| {
                     (
                         neuron.id().id,

rs/nns/governance/src/neuron_store.rs

@@ -807,14 +807,14 @@ impl NeuronStore {
     }
     pub fn with_active_neurons_iter<R>(
         &self,
-        callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Neuron> + 'b>) -> R,
+        callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Cow<Neuron>> + 'b>) -> R,
     ) -> R {
         self.with_active_neurons_iter_sections(callback, NeuronSections::all())
     }
 
     fn with_active_neurons_iter_sections<R>(
         &self,
-        callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Neuron> + 'b>) -> R,
+        callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Cow<Neuron>> + 'b>) -> R,
         sections: NeuronSections,
     ) -> R {
         if self.use_stable_memory_for_all_neurons {
@@ -825,13 +825,13 @@ impl NeuronStore {
                     stable_store
                         .range_neurons_sections(.., sections)
                         .filter(|n| !n.is_inactive(now))
-                        .chain(self.heap_neurons.values().cloned()),
-                ) as Box<dyn Iterator<Item = Neuron>>;
+                        .map(Cow::Owned)
+                        .chain(self.heap_neurons.values().map(Cow::Borrowed)),
+                );
                 callback(iter)
             })
         } else {
-            let iter =
-                Box::new(self.heap_neurons.values().cloned()) as Box<dyn Iterator<Item = Neuron>>;
+            let iter = Box::new(self.heap_neurons.values().map(Cow::Borrowed));
             callback(iter)
         }
     }
@@ -870,9 +870,13 @@ impl NeuronStore {
     fn filter_map_active_neurons<R>(
         &self,
         filter: impl Fn(&Neuron) -> bool,
-        f: impl FnMut(Neuron) -> R,
+        f: impl Fn(&Neuron) -> R,
     ) -> Vec<R> {
-        self.with_active_neurons_iter(|iter| iter.filter(|n| filter(n)).map(f).collect())
+        self.with_active_neurons_iter(|iter| {
+            iter.filter(|n| filter(n.as_ref()))
+                .map(|n| f(n.as_ref()))
+                .collect()
+        })
     }
 
     fn is_active_neurons_fund_neuron(neuron: &Neuron, now: u64) -> bool {
@@ -932,7 +936,7 @@ impl NeuronStore {
         let mut deciding_voting_power: u128 = 0;
         let mut potential_voting_power: u128 = 0;
 
-        let mut process_neuron = |neuron: Neuron| {
+        let mut process_neuron = |neuron: &Neuron| {
             if neuron.is_inactive(now_seconds)
                 || neuron.dissolve_delay_seconds(now_seconds)
                     < MIN_DISSOLVE_DELAY_FOR_VOTE_ELIGIBILITY_SECONDS
@@ -956,7 +960,7 @@ impl NeuronStore {
         self.with_active_neurons_iter_sections(
             |iter| {
                 for neuron in iter {
-                    process_neuron(neuron);
+                    process_neuron(neuron.as_ref());
                 }
             },
             NeuronSections::default(),

rs/nns/governance/tests/governance.rs

@@ -126,7 +126,7 @@ use std::{
     collections::{BTreeMap, HashSet, VecDeque},
     convert::{TryFrom, TryInto},
     iter::{self, once},
-    ops::Div,
+    ops::{Deref, Div},
     path::PathBuf,
     sync::{Arc, Mutex},
 };
@@ -7200,9 +7200,11 @@ fn test_manage_and_reward_node_providers() {
         ProposalStatus::Executed
     );
     // Find the neuron...
-    let neuron = gov
-        .neuron_store
-        .with_active_neurons_iter(|mut iter| iter.find(|x| x.controller() == np_pid).unwrap());
+    let neuron = gov.neuron_store.with_active_neurons_iter(|mut iter| {
+        iter.find(|x| x.controller() == np_pid)
+            .map(|n| n.deref().clone())
+            .unwrap()
+    });
     assert_eq!(neuron.stake_e8s(), 99_999_999);
     // Find the transaction in the ledger...
     driver.assert_account_contains(
@@ -7508,9 +7510,11 @@ fn test_manage_and_reward_multiple_node_providers() {
 
     // Check third reward
     // Find the neuron...
-    let neuron = gov
-        .neuron_store
-        .with_active_neurons_iter(|mut iter| iter.find(|x| x.controller() == np_pid_2).unwrap());
+    let neuron = gov.neuron_store.with_active_neurons_iter(|mut iter| {
+        iter.find(|x| x.controller() == np_pid_2)
+            .map(|n| n.deref().clone())
+            .unwrap()
+    });
     assert_eq!(neuron.stake_e8s(), 99_999_999);
     // Find the transaction in the ledger...
     driver.assert_account_contains(