perf(nns): Change from benchmarking listing neurons for unstake maturity to benchmarking the entire unstake maturity operation (#4183)

# Why

When trying to unstake maturity, not only the listing steps are
computationally intensive, the actual unstaking is too, because it calls
`with_neuron_mut` in a loop. We should cap the number of neurons to
unstake in a single message, and this benchmark prepares for the
improvement.

# What

* Refactor the `Governance::unstake_maturity_of_dissolved_neurons` to
`NeuronStore::unstake_maturity_of_dissolved_neurons` since only neuron
store is involved for this operation.
* Change benchmark from `list_neurons_ready_to_unstake_maturity` to
`unstake_maturity_of_dissolved_neurons`
* Add bench scopes for listing and unstaking, since the former is linear
to # of neurons and the latter depends on the available number of
neurons to unstake maturity (and those 2 numbers don't have a reasonable
ratio that we can assume)
* Update bench results

Author: jasonz-dfinity

rs/nns/governance/canbench/canbench_results.yml

@@ -101,18 +101,6 @@ benches:
       heap_increase: 9
       stable_memory_increase: 0
     scopes: {}
-  list_neurons_ready_to_unstake_maturity_heap:
-    total:
-      instructions: 158195
-      heap_increase: 0
-      stable_memory_increase: 0
-    scopes: {}
-  list_neurons_ready_to_unstake_maturity_stable:
-    total:
-      instructions: 43332558
-      heap_increase: 0
-      stable_memory_increase: 0
-    scopes: {}
   list_neurons_stable:
     total:
       instructions: 113665560
@@ -173,10 +161,38 @@ benches:
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
+  unstake_maturity_of_dissolved_neurons_heap:
+    total:
+      instructions: 2644396
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes:
+      list_neuron_ids:
+        instructions: 225923
+        heap_increase: 0
+        stable_memory_increase: 0
+      unstake_maturity:
+        instructions: 2416625
+        heap_increase: 0
+        stable_memory_increase: 0
+  unstake_maturity_of_dissolved_neurons_stable:
+    total:
+      instructions: 92282613
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes:
+      list_neuron_ids:
+        instructions: 44812672
+        heap_increase: 0
+        stable_memory_increase: 0
+      unstake_maturity:
+        instructions: 47468213
+        heap_increase: 0
+        stable_memory_increase: 0
   update_recent_ballots_stable_memory:
     total:
       instructions: 275035
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
-version: 0.1.8
+version: 0.1.9

rs/nns/governance/src/governance.rs

@@ -6562,30 +6562,11 @@ impl Governance {
         xdr_permyriad_per_icp / dec!(10_000)
     }
 
-    /// When a neuron is finally dissolved, if there is any staked maturity it is moved to regular maturity
-    /// which can be spawned (and is modulated).
+    /// Unstakes the maturity of neurons that have dissolved.
     pub fn unstake_maturity_of_dissolved_neurons(&mut self) {
         let now_seconds = self.env.now();
-        // Filter all the neurons that are currently in "dissolved" state and have some staked maturity.
-        // No neuron in stable storage should have staked maturity.
-        for neuron_id in self
-            .neuron_store
-            .list_neurons_ready_to_unstake_maturity(now_seconds)
-        {
-            let unstake_result = self
-                .neuron_store
-                .with_neuron_mut(&neuron_id, |neuron| neuron.unstake_maturity(now_seconds));
-
-            match unstake_result {
-                Ok(_) => {}
-                Err(e) => {
-                    println!(
-                        "{}Error in heartbeat when moving staked maturity for neuron {:?}: {:?}",
-                        LOG_PREFIX, neuron_id, e
-                    );
-                }
-            };
-        }
+        self.neuron_store
+            .unstake_maturity_of_dissolved_neurons(now_seconds);
     }
 
     fn can_spawn_neurons(&self) -> bool {

rs/nns/governance/src/neuron/types.rs

@@ -1706,7 +1706,7 @@ pub struct NeuronBuilder {
     recent_ballots_next_entry_index: Option<usize>,
     #[cfg(test)]
     transfer: Option<NeuronStakeTransfer>,
-    #[cfg(test)]
+    #[cfg(any(test, feature = "canbench-rs"))]
     staked_maturity_e8s_equivalent: Option<u64>,
     #[cfg(test)]
     known_neuron_data: Option<KnownNeuronData>,
@@ -1748,7 +1748,7 @@ impl NeuronBuilder {
             recent_ballots_next_entry_index: Some(0),
             #[cfg(test)]
             transfer: None,
-            #[cfg(test)]
+            #[cfg(any(test, feature = "canbench-rs"))]
             staked_maturity_e8s_equivalent: None,
             #[cfg(test)]
             known_neuron_data: None,
@@ -1850,7 +1850,7 @@ impl NeuronBuilder {
         self
     }
 
-    #[cfg(test)]
+    #[cfg(any(test, feature = "canbench-rs"))]
     pub fn with_staked_maturity_e8s_equivalent(
         mut self,
         staked_maturity_e8s_equivalent: u64,
@@ -1906,7 +1906,7 @@ impl NeuronBuilder {
             recent_ballots_next_entry_index,
             #[cfg(test)]
             transfer,
-            #[cfg(test)]
+            #[cfg(any(test, feature = "canbench-rs"))]
             staked_maturity_e8s_equivalent,
             #[cfg(test)]
             known_neuron_data,
@@ -1937,7 +1937,7 @@ impl NeuronBuilder {
         let recent_ballots_next_entry_index = Some(0);
         #[cfg(not(test))]
         let transfer = None;
-        #[cfg(not(test))]
+        #[cfg(not(any(test, feature = "canbench-rs")))]
         let staked_maturity_e8s_equivalent = None;
         #[cfg(not(test))]
         let known_neuron_data = None;

rs/nns/governance/src/neuron_store.rs

@@ -907,7 +907,7 @@ impl NeuronStore {
     }
 
     /// List all neuron ids whose neurons have staked maturity greater than 0.
-    pub fn list_neurons_ready_to_unstake_maturity(&self, now_seconds: u64) -> Vec<NeuronId> {
+    fn list_neurons_ready_to_unstake_maturity(&self, now_seconds: u64) -> Vec<NeuronId> {
         self.with_active_neurons_iter_sections(
             |iter| {
                 iter.filter(|neuron| neuron.ready_to_unstake_maturity(now_seconds))
@@ -981,6 +981,35 @@ impl NeuronStore {
         (ballots, deciding_voting_power, potential_voting_power)
     }
 
+    /// When a neuron is finally dissolved, if there is any staked maturity it is moved to regular maturity
+    /// which can be spawned (and is modulated).
+    pub fn unstake_maturity_of_dissolved_neurons(&mut self, now_seconds: u64) {
+        let neuron_ids = {
+            #[cfg(feature = "canbench-rs")]
+            let _scope_list = canbench_rs::bench_scope("list_neuron_ids");
+            self.list_neurons_ready_to_unstake_maturity(now_seconds)
+        };
+
+        #[cfg(feature = "canbench-rs")]
+        let _scope_unstake = canbench_rs::bench_scope("unstake_maturity");
+        // Filter all the neurons that are currently in "dissolved" state and have some staked maturity.
+        // No neuron in stable storage should have staked maturity.
+        for neuron_id in neuron_ids {
+            let unstake_result =
+                self.with_neuron_mut(&neuron_id, |neuron| neuron.unstake_maturity(now_seconds));
+
+            match unstake_result {
+                Ok(_) => {}
+                Err(e) => {
+                    println!(
+                        "{}Error when moving staked maturity for neuron {:?}: {:?}",
+                        LOG_PREFIX, neuron_id, e
+                    );
+                }
+            };
+        }
+    }
+
     /// Returns the full neuron if the given principal is authorized - either it can vote for the
     /// given neuron or any of its neuron managers.
     pub fn get_full_neuron(

rs/nns/governance/src/neuron_store/benches.rs

@@ -318,7 +318,7 @@ fn add_neuron_ready_to_unstake_maturity(
 ) {
     let id = rng.next_u64();
     let subaccount = subaccount_from_id(id);
-    let mut neuron = NeuronBuilder::new(
+    let neuron = NeuronBuilder::new(
         NeuronId { id: rng.next_u64() },
         subaccount,
         PrincipalId::new_user_test_id(id),
@@ -327,32 +327,36 @@ fn add_neuron_ready_to_unstake_maturity(
         },
         123_456_789,
     )
+    .with_staked_maturity_e8s_equivalent(1_000_000_000)
     .build();
-    neuron.staked_maturity_e8s_equivalent = Some(1_000_000_000);
     neuron_store.add_neuron(neuron).unwrap();
 }
 
 #[bench(raw)]
-fn list_neurons_ready_to_unstake_maturity_heap() -> BenchResult {
+fn unstake_maturity_of_dissolved_neurons_heap() -> BenchResult {
     let _a = temporarily_disable_allow_active_neurons_in_stable_memory();
     let _b = temporarily_disable_migrate_active_neurons_to_stable_memory();
     let mut rng = new_rng();
     let mut neuron_store = set_up_neuron_store(&mut rng, 1_000, 2_000);
-    add_neuron_ready_to_unstake_maturity(now_seconds(), &mut rng, &mut neuron_store);
+    for _ in 0..100 {
+        add_neuron_ready_to_unstake_maturity(now_seconds(), &mut rng, &mut neuron_store);
+    }
 
-    bench_fn(|| neuron_store.list_neurons_ready_to_unstake_maturity(now_seconds()))
+    bench_fn(|| neuron_store.unstake_maturity_of_dissolved_neurons(now_seconds()))
 }
 
 #[bench(raw)]
-fn list_neurons_ready_to_unstake_maturity_stable() -> BenchResult {
+fn unstake_maturity_of_dissolved_neurons_stable() -> BenchResult {
     let _a = temporarily_enable_allow_active_neurons_in_stable_memory();
     let _b = temporarily_enable_migrate_active_neurons_to_stable_memory();
     let mut rng = new_rng();
     let mut neuron_store = set_up_neuron_store(&mut rng, 1_000, 2_000);
-    add_neuron_ready_to_unstake_maturity(now_seconds(), &mut rng, &mut neuron_store);
+    for _ in 0..100 {
+        add_neuron_ready_to_unstake_maturity(now_seconds(), &mut rng, &mut neuron_store);
+    }
 
     bench_fn(|| {
-        neuron_store.list_neurons_ready_to_unstake_maturity(now_seconds());
+        neuron_store.unstake_maturity_of_dissolved_neurons(now_seconds());
     })
 }