PR feedback

rs/nns/governance/src/governance.rs

@@ -2318,7 +2318,7 @@ impl Governance {
                     // 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.recent_ballots();
+                    proto.recent_ballots = neuron.sorted_recent_ballots();
                     full_neurons.push(proto);
                 }
             });

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

@@ -77,6 +77,7 @@ pub struct Neuron {
     /// Map `Topic` to followees. The key is represented by an integer as
     /// Protobuf does not support enum keys in maps.
     pub followees: HashMap<i32, Followees>,
+    // TODO DO NOT MERGE make this private
     /// Information about how this neuron voted in the recent past. It
     /// only contains proposals that the neuron voted yes or no on.
     pub recent_ballots: Vec<BallotInfo>,
@@ -436,18 +437,20 @@ impl Neuron {
         std::cmp::min(ad_stake, u64::MAX as u128) as u64
     }
 
-    pub(crate) fn recent_ballots(&self) -> Vec<BallotInfo> {
-        if self.recent_ballots_next_entry_index.is_none() {
-            self.recent_ballots.clone()
-        } else {
+    /// 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;
-            let index = self.recent_ballots_next_entry_index.unwrap();
             recent_ballots[index..]
                 .iter()
                 .chain(recent_ballots[..index].iter())
                 .rev()
                 .cloned()
                 .collect()
+        } else {
+            self.recent_ballots.clone()
         }
     }
 
@@ -525,11 +528,12 @@ impl Neuron {
         }
 
         // Data migration for updating the recent ballots so we can use a circular buffer here.
-        if self.recent_ballots_next_entry_index.is_none() {
+        let next_entry_index = if let Some(index) = self.recent_ballots_next_entry_index {
+            index
+        } else {
             self.recent_ballots.reverse();
-            self.recent_ballots_next_entry_index =
-                Some(self.recent_ballots.len() % MAX_NEURON_RECENT_BALLOTS);
-        }
+            self.recent_ballots.len() % MAX_NEURON_RECENT_BALLOTS
+        };
 
         let ballot_info = BallotInfo {
             proposal_id: Some(*proposal_id),
@@ -538,13 +542,13 @@ impl Neuron {
 
         // Vector is full
         if self.recent_ballots.len() >= MAX_NEURON_RECENT_BALLOTS {
-            self.recent_ballots[self.recent_ballots_next_entry_index.unwrap()] = ballot_info;
+            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((self.recent_ballots_next_entry_index.unwrap() + 1) % MAX_NEURON_RECENT_BALLOTS);
+            Some((next_entry_index + 1) % MAX_NEURON_RECENT_BALLOTS);
     }
 
     pub(crate) fn refresh_voting_power(&mut self, now_seconds: u64) {
@@ -906,7 +910,7 @@ impl Neuron {
             || self.visibility() == Some(Visibility::Public)
             || self.is_hotkey_or_controller(&requester);
         if show_full {
-            recent_ballots.append(&mut self.recent_ballots());
+            recent_ballots.append(&mut self.sorted_recent_ballots());
             joined_community_fund_timestamp_seconds = self.joined_community_fund_timestamp_seconds;
         }
 
@@ -1394,7 +1398,15 @@ 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),
+            // 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 {
@@ -1484,6 +1496,7 @@ 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),
         }
     }

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

@@ -748,7 +748,7 @@ fn test_recent_ballots_accessor_pre_and_post_migration() {
     .build();
     neuron.recent_ballots_next_entry_index = None;
 
-    assert_eq!(neuron.recent_ballots(), recent_ballots);
+    assert_eq!(neuron.sorted_recent_ballots(), recent_ballots);
 
     neuron.register_recent_ballot(Topic::NetworkEconomics, &ProposalId { id: 100 }, Vote::No);
     assert_eq!(neuron.recent_ballots_next_entry_index, Some(1));
@@ -766,5 +766,5 @@ fn test_recent_ballots_accessor_pre_and_post_migration() {
         recent_ballots
     };
 
-    assert_eq!(neuron.recent_ballots(), expected_updated_ballots);
+    assert_eq!(neuron.sorted_recent_ballots(), expected_updated_ballots);
 }

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

@@ -16,7 +16,7 @@ use maplit::hashmap;
 use prost::Message;
 use std::{
     borrow::Cow,
-    collections::{BTreeMap as HeapBTreeMap, BTreeSet as HeapBTreeSet, HashMap},
+    collections::{btree_map::Entry, BTreeMap as HeapBTreeMap, HashMap},
     iter::Peekable,
     ops::{Bound as RangeBound, RangeBounds},
 };
@@ -897,15 +897,21 @@ fn update_range<Key, Value, Memory>(
     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);
-        }
+        match new_entries.entry(key) {
+            // If our new entries do not include a key in existing, we remove it from existing.
+            Entry::Vacant(_) => {
+                to_remove.push(key.clone());
+            }
+            Entry::Occupied(entry) => {
+                // If our new entries have the same value as what exists, we do not want to insert
+                // it, but instead remove it from the list of new entries, since it's present.
+                if *entry.get() == value {
+                    entry.remove();
+                }
+            }
+        };
     }
 
     for (new_key, new_value) in new_entries {