wip

rs/nns/governance/canbench/canbench_results.yml

@@ -1,91 +1,91 @@
 benches:
   add_neuron_active_maximum:
     total:
-      instructions: 35887686
+      instructions: 36126615
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_active_typical:
     total:
-      instructions: 1833590
+      instructions: 1833344
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_maximum:
     total:
-      instructions: 98494555
+      instructions: 96225749
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_typical:
     total:
-      instructions: 7593880
+      instructions: 7376879
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_all_heap:
     total:
-      instructions: 32187827
+      instructions: 31810085
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_heap_neurons_stable_index:
     total:
-      instructions: 54527506
+      instructions: 54166563
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_stable_everything:
     total:
-      instructions: 1978888007
+      instructions: 451171207
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_stable_neurons_with_heap_index:
     total:
-      instructions: 1956237656
+      instructions: 428881617
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   centralized_following_all_stable:
     total:
-      instructions: 2057421983
+      instructions: 367625528
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   compute_ballots_for_new_proposal_with_stable_neurons:
     total:
-      instructions: 1685116
+      instructions: 1693073
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   neuron_metrics_calculation_heap:
     total:
-      instructions: 840837
+      instructions: 837546
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   neuron_metrics_calculation_stable:
     total:
-      instructions: 1834053
+      instructions: 1843815
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   range_neurons_performance:
     total:
-      instructions: 48615662
+      instructions: 47586359
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   single_vote_all_stable:
     total:
-      instructions: 13224969
+      instructions: 2308192
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
-  update_recent_ballots:
+  update_recent_ballots_stable_memory:
     total:
-      instructions: 16275553
+      instructions: 4797610
       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/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,
                 }
             }
         }
@@ -500,20 +507,28 @@ impl Neuron {
         if topic == Topic::ExchangeRate {
             return;
         }
+
+        // Data migration for updating the recent ballots so we can use a circular buffer here.
+        if self.recent_ballots_next_entry_index.is_none() {
+            self.recent_ballots.reverse();
+            self.recent_ballots_next_entry_index =
+                Some(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[self.recent_ballots_next_entry_index.unwrap()] = ballot_info;
+        } else {
+            self.recent_ballots.push(ballot_info);
         }
+        // Advance the index
+        self.recent_ballots_next_entry_index =
+            Some((self.recent_ballots_next_entry_index.unwrap() + 1) % MAX_NEURON_RECENT_BALLOTS);
     }
 
     pub(crate) fn refresh_voting_power(&mut self, now_seconds: u64) {
@@ -1122,6 +1137,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 +1174,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 +1207,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 +1253,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 +1347,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 +1379,7 @@ impl TryFrom<Neuron> for DecomposedNeuron {
             dissolve_state,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index: recent_ballots_next_entry_index.map(|x| x as u32),
         };
 
         Ok(Self {
@@ -1409,6 +1430,7 @@ impl From<DecomposedNeuron> for Neuron {
             dissolve_state,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index,
         } = main;
 
         let subaccount =
@@ -1447,6 +1469,7 @@ impl From<DecomposedNeuron> 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),
         }
     }
 }
@@ -1693,6 +1716,9 @@ impl NeuronBuilder {
         #[cfg(not(test))]
         let known_neuron_data = None;
 
+        let recent_ballots_next_entry_index =
+            Some(recent_ballots.len() % MAX_NEURON_RECENT_BALLOTS);
+
         Neuron {
             id,
             subaccount,
@@ -1716,6 +1742,7 @@ impl NeuronBuilder {
             neuron_type,
             visibility,
             voting_power_refreshed_timestamp_seconds,
+            recent_ballots_next_entry_index,
         }
     }
 }

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

@@ -178,12 +178,13 @@ fn add_neuron_inactive_typical() -> BenchResult {
 }
 
 #[bench(raw)]
-fn update_recent_ballots() -> BenchResult {
+fn update_recent_ballots_stable_memory() -> BenchResult {
     let _a = temporarily_enable_active_neurons_in_stable_memory();
     let _b = temporarily_enable_stable_memory_following_index();
     let mut rng = new_rng();
     let mut neuron_store = set_up_neuron_store(&mut rng, 100, 200);
-    let neuron = build_neuron(&mut rng, NeuronLocation::Heap, NeuronSize::Maximum);
+    let mut neuron = build_neuron(&mut rng, NeuronLocation::Heap, NeuronSize::Maximum);
+    neuron.recent_ballots_next_entry_index = Some(0);
     let id = neuron.id();
 
     assert_eq!(neuron.recent_ballots.len(), MAX_NEURON_RECENT_BALLOTS);

rs/nns/governance/src/pb/conversions.rs

@@ -179,6 +179,8 @@ impl From<pb_api::Neuron> for pb::Neuron {
             dissolve_state: item.dissolve_state.map(|x| x.into()),
             visibility: item.visibility,
             voting_power_refreshed_timestamp_seconds: item.voting_power_refreshed_timestamp_seconds,
+            // This field is internal only and should not be read from API types.
+            recent_ballots_next_entry_index: None,
         }
     }
 }

rs/nns/governance/src/storage/neurons.rs

@@ -798,7 +798,7 @@ fn update_repeated_field<Element, Memory>(
     new_elements: Vec<Element>,
     map: &mut StableBTreeMap<(NeuronId, /* index */ u64), Element, Memory>,
 ) where
-    Element: Storable,
+    Element: Storable + PartialEq,
     Memory: ic_stable_structures::Memory,
 {
     let new_entries = new_elements
@@ -823,27 +823,30 @@ fn update_repeated_field<Element, Memory>(
 // TODO(NNS1-2513): To avoid the caller passing an incorrect range (e.g. too
 // small, or to big), derive range from NeuronId.
 fn update_range<Key, Value, Memory>(
-    new_entries: HeapBTreeMap<Key, Value>,
+    mut new_entries: HeapBTreeMap<Key, Value>,
     range: impl RangeBounds<Key>,
     map: &mut StableBTreeMap<Key, Value, Memory>,
 ) where
     Key: Storable + Ord + Clone,
-    Value: Storable,
+    Value: Storable + PartialEq,
     Memory: ic_stable_structures::Memory,
 {
     let new_keys = new_entries.keys().cloned().collect::<HeapBTreeSet<Key>>();
 
+    let mut to_remove = vec![];
+    for (key, value) in map.range(range) {
+        if !new_keys.contains(&key) {
+            to_remove.push(key.clone());
+        } else if new_entries.get(&key).expect("We just checked it's there") == &value {
+            new_entries.remove(&key);
+        }
+    }
+
     for (new_key, new_value) in new_entries {
         map.insert(new_key, new_value);
     }
 
-    let obsolete_keys = map
-        .range(range)
-        .filter(|(key, _value)| !new_keys.contains(key))
-        .map(|(key, _value)| key)
-        .collect::<Vec<_>>();
-
-    for obsolete_key in obsolete_keys {
+    for obsolete_key in to_remove {
         map.remove(&obsolete_key);
     }
 }

rs/nns/governance/src/storage/neurons/neurons_tests.rs

@@ -506,12 +506,13 @@ fn test_abridged_neuron_size() {
         dissolve_state: Some(DissolveState::WhenDissolvedTimestampSeconds(u64::MAX)),
         visibility: None,
         voting_power_refreshed_timestamp_seconds: Some(u64::MAX),
+        recent_ballots_next_entry_index: Some(100),
     };
 
     assert!(abridged_neuron.encoded_len() as u32 <= AbridgedNeuron::BOUND.max_size());
     // This size can be updated. This assertion is here to make sure we are very aware of growth.
     // Reminder: the amount we allocated for AbridgedNeuron is 380 bytes.
-    assert_eq!(abridged_neuron.encoded_len(), 196);
+    assert_eq!(abridged_neuron.encoded_len(), 199);
 }
 
 #[test]