fix(nns): Improve performance of stable neuron recent ballots recordi…

…ng (#2697)

This improves the performance of recording recent ballots in stable
memory by over 90% (in the worst case).

This PR contains migration code for ballots which cannot be rolled back
after it is deployed.

[Prev](#2670) |
[Next](#2701)

rs/nns/governance/canbench/canbench_results.yml

@@ -1,61 +1,61 @@
 benches:
   add_neuron_active_maximum:
     total:
-      instructions: 35992698
+      instructions: 36149096
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_active_typical:
     total:
-      instructions: 1838632
+      instructions: 1834158
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_maximum:
     total:
-      instructions: 98666560
+      instructions: 96285952
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_typical:
     total:
-      instructions: 7577551
+      instructions: 7381170
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_all_heap:
     total:
-      instructions: 32302982
+      instructions: 31887311
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_heap_neurons_stable_index:
     total:
-      instructions: 54494918
+      instructions: 54233977
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_stable_everything:
     total:
-      instructions: 1969690886
+      instructions: 160313292
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_stable_neurons_with_heap_index:
     total:
-      instructions: 1947188278
+      instructions: 138075846
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   centralized_following_all_stable:
     total:
-      instructions: 2047356165
+      instructions: 66264382
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   compute_ballots_for_new_proposal_with_stable_neurons:
     total:
-      instructions: 1743731
+      instructions: 1690958
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -85,31 +85,31 @@ benches:
     scopes: {}
   neuron_metrics_calculation_heap:
     total:
-      instructions: 845137
+      instructions: 837546
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   neuron_metrics_calculation_stable:
     total:
-      instructions: 1835769
+      instructions: 1843015
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   range_neurons_performance:
     total:
-      instructions: 47339494
+      instructions: 47360359
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   single_vote_all_stable:
     total:
-      instructions: 13161102
+      instructions: 364240
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
-  update_recent_ballots:
+  update_recent_ballots_stable_memory:
     total:
-      instructions: 16208293
+      instructions: 236573
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}

rs/nns/governance/proto/ic_nns_governance/pb/v1/governance.proto

@@ -458,6 +458,11 @@ message Neuron {
   // refreshed, this will be set to 2024-11-05T00:00:01 UTC (1730764801 seconds
   // after the UNIX epoch).
   optional uint64 voting_power_refreshed_timestamp_seconds = 24;
+
+  // The index of the next entry in the recent_ballots list that will be
+  // used for the circular buffer. This is used to determine which entry
+  // to overwrite next.
+  optional uint32 recent_ballots_next_entry_index = 25;
 }
 
 // Subset of Neuron that has no collections or big fields that might not exist in most neurons, and
@@ -484,6 +489,7 @@ message AbridgedNeuron {
   optional NeuronType neuron_type = 22;
   optional Visibility visibility = 23;
   optional uint64 voting_power_refreshed_timestamp_seconds = 24;
+  optional uint32 recent_ballots_next_entry_index = 25;
 
   reserved 1;
   reserved "id";

rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs

@@ -306,6 +306,11 @@ pub struct Neuron {
     /// after the UNIX epoch).
     #[prost(uint64, optional, tag = "24")]
     pub voting_power_refreshed_timestamp_seconds: ::core::option::Option<u64>,
+    /// The index of the next entry in the recent_ballots list that will be
+    /// used for the circular buffer. This is used to determine which entry
+    /// to overwrite next.
+    #[prost(uint32, optional, tag = "25")]
+    pub recent_ballots_next_entry_index: ::core::option::Option<u32>,
     /// At any time, at most one of `when_dissolved` and
     /// `dissolve_delay` are specified.
     ///
@@ -433,6 +438,8 @@ pub struct AbridgedNeuron {
     pub visibility: ::core::option::Option<i32>,
     #[prost(uint64, optional, tag = "24")]
     pub voting_power_refreshed_timestamp_seconds: ::core::option::Option<u64>,
+    #[prost(uint32, optional, tag = "25")]
+    pub recent_ballots_next_entry_index: ::core::option::Option<u32>,
     #[prost(oneof = "abridged_neuron::DissolveState", tags = "9, 10")]
     pub dissolve_state: ::core::option::Option<abridged_neuron::DissolveState>,
 }

rs/nns/governance/src/governance.rs

@@ -2313,7 +2313,13 @@ impl Governance {
                             && neuron.visibility() == Some(Visibility::Public)
                         );
                 if let_caller_read_full_neuron {
-                    full_neurons.push(NeuronProto::from(neuron.clone()));
+                    let mut proto = NeuronProto::from(neuron.clone());
+                    // We get the recent_ballots from the neuron itself, because
+                    // we are using a circular buffer to store them.  This solution is not ideal, but
+                    // we need to do a larger refactoring to use the correct API types instead of the internal
+                    // governance proto at this level.
+                    proto.recent_ballots = neuron.sorted_recent_ballots();
+                    full_neurons.push(proto);
                 }
             });
         }

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

@@ -125,6 +125,10 @@ pub struct Neuron {
     /// month. After that, following is cleared, except for ManageNeuron
     /// proposals.
     voting_power_refreshed_timestamp_seconds: u64,
+    /// This field is used to store the index of the most recent ballot in the
+    /// `recent_ballots` circular buffer. This is used to optimize insertions
+    /// into stable memory, to avoid rewriting all the data.
+    pub recent_ballots_next_entry_index: Option<usize>,
 }
 
 /// This is mostly the same as the version of PartialEq generated by derive. The
@@ -160,6 +164,7 @@ impl PartialEq for Neuron {
             voting_power_refreshed_timestamp_seconds: &'a u64,
 
             visibility: Visibility,
+            recent_ballots_next_entry_index: Option<usize>,
         }
 
         impl<'a> Normalized<'a> {
@@ -188,6 +193,7 @@ impl PartialEq for Neuron {
                     voting_power_refreshed_timestamp_seconds,
 
                     visibility: _,
+                    recent_ballots_next_entry_index,
                 } = src;
 
                 let visibility = src.visibility().unwrap_or(Visibility::Private);
@@ -216,6 +222,7 @@ impl PartialEq for Neuron {
                     voting_power_refreshed_timestamp_seconds,
 
                     visibility,
+                    recent_ballots_next_entry_index: *recent_ballots_next_entry_index,
                 }
             }
         }
@@ -429,6 +436,23 @@ impl Neuron {
         std::cmp::min(ad_stake, u64::MAX as u128) as u64
     }
 
+    /// Get the recent ballots, with most recent ballots first
+    pub(crate) fn sorted_recent_ballots(&self) -> Vec<BallotInfo> {
+        if let Some(index) = self.recent_ballots_next_entry_index {
+            // We store ballots in a circular buffer with oldest first, so we need to reverse
+            // this as well as rearrange it before returning it.
+            let recent_ballots = &self.recent_ballots;
+            recent_ballots[index..]
+                .iter()
+                .chain(recent_ballots[..index].iter())
+                .rev()
+                .cloned()
+                .collect()
+        } else {
+            self.recent_ballots.clone()
+        }
+    }
+
     /// Given the specified `ballots`: determine how this neuron would
     /// vote on a proposal of `topic` based on which neurons this
     /// neuron follows on this topic (or on the default topic if this
@@ -489,6 +513,7 @@ impl Neuron {
     /// Register that this neuron has cast a ballot for a
     /// proposal. Don't include votes on "real time" topics (such as
     /// setting the ICP/SDR exchange rate).
+    /// TODO(NNS1-3479) delete this method after all neurons are migrated to stable memory
     pub(crate) fn register_recent_ballot(
         &mut self,
         topic: Topic,
@@ -500,20 +525,29 @@ impl Neuron {
         if topic == Topic::ExchangeRate {
             return;
         }
+
+        // Data migration for updating the recent ballots so we can use a circular buffer here.
+        let next_entry_index = if let Some(index) = self.recent_ballots_next_entry_index {
+            index
+        } else {
+            self.recent_ballots.reverse();
+            self.recent_ballots.len() % MAX_NEURON_RECENT_BALLOTS
+        };
+
         let ballot_info = BallotInfo {
             proposal_id: Some(*proposal_id),
             vote: vote as i32,
         };
-        // We would really like to have a circular buffer here. As
-        // we're dealing with a simple vector, we insert at the
-        // beginning and remove at the end once we have reached
-        // the maximum number of votes to keep track of.
-        self.recent_ballots.insert(0, ballot_info);
-        // Pop and discard elements from the end until we reach
-        // the maximum allowed length of the vector.
-        while self.recent_ballots.len() > MAX_NEURON_RECENT_BALLOTS {
-            self.recent_ballots.pop();
+
+        // Vector is full
+        if self.recent_ballots.len() >= MAX_NEURON_RECENT_BALLOTS {
+            self.recent_ballots[next_entry_index] = ballot_info;
+        } else {
+            self.recent_ballots.push(ballot_info);
         }
+        // Advance the index
+        self.recent_ballots_next_entry_index =
+            Some((next_entry_index + 1) % MAX_NEURON_RECENT_BALLOTS);
     }
 
     pub(crate) fn refresh_voting_power(&mut self, now_seconds: u64) {
@@ -875,7 +909,7 @@ impl Neuron {
             || self.visibility() == Some(Visibility::Public)
             || self.is_hotkey_or_controller(&requester);
         if show_full {
-            recent_ballots.append(&mut self.recent_ballots.clone());
+            recent_ballots.append(&mut self.sorted_recent_ballots());
             joined_community_fund_timestamp_seconds = self.joined_community_fund_timestamp_seconds;
         }
 
@@ -1098,7 +1132,6 @@ impl Neuron {
 impl From<Neuron> for NeuronProto {
     fn from(neuron: Neuron) -> Self {
         let visibility = neuron.visibility().map(|visibility| visibility as i32);
-
         let Neuron {
             id,
             subaccount,
@@ -1122,6 +1155,7 @@ impl From<Neuron> for NeuronProto {
             neuron_type,
             visibility: _,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index,
         } = neuron;
 
         let id = Some(id);
@@ -1158,6 +1192,7 @@ impl From<Neuron> for NeuronProto {
             neuron_type,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index: recent_ballots_next_entry_index.map(|x| x as u32),
         }
     }
 }
@@ -1190,6 +1225,7 @@ impl TryFrom<NeuronProto> for Neuron {
             neuron_type,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index,
         } = proto;
 
         let id = id.ok_or("Neuron ID is missing")?;
@@ -1235,6 +1271,7 @@ impl TryFrom<NeuronProto> for Neuron {
             neuron_type,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index: recent_ballots_next_entry_index.map(|x| x as usize),
         })
     }
 }
@@ -1328,6 +1365,7 @@ impl TryFrom<Neuron> for DecomposedNeuron {
             neuron_type,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index,
         } = source;
 
         let account = subaccount.to_vec();
@@ -1359,6 +1397,15 @@ impl TryFrom<Neuron> for DecomposedNeuron {
             dissolve_state,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            // Conversion to u32 is safe because the value cannot be very large.  If it overflowed
+            // u32 max, we would have spent 68GB on recent ballots for this single neuron.
+            recent_ballots_next_entry_index: recent_ballots_next_entry_index
+                .map(|x| {
+                    u32::try_from(x).map_err(|e| NeuronStoreError::InvalidData {
+                        reason: format!("Failed to convert recent_ballots_next_entry_index: {}", e),
+                    })
+                })
+                .transpose()?,
         };
 
         Ok(Self {
@@ -1409,6 +1456,7 @@ impl From<DecomposedNeuron> for Neuron {
             dissolve_state,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index,
         } = main;
 
         let subaccount =
@@ -1447,6 +1495,8 @@ impl From<DecomposedNeuron> for Neuron {
             neuron_type,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            // usize is always at least u32, so this is safe.
+            recent_ballots_next_entry_index: recent_ballots_next_entry_index.map(|x| x as usize),
         }
     }
 }
@@ -1716,6 +1766,7 @@ impl NeuronBuilder {
             neuron_type,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index: None,
         }
     }
 }