feat(nns): Split the active_neurons_in_stable_memory into 2 flags (#3312)

# Why

The previous set up for the feature flag
`active_neurons_in_stable_memory` does not allow us to safely rollback
if necessary, because the rolled back version would only look for active
neurons in the heap memory, while some would still be in the stable
memory until the timer picks them up (if the reverse migration is
implemented).

# What

Use `allow_active_neurons_in_stable_memory` for whether the canister
allows active neurons to be in stable memory, by checking stable memory
when iterating/finding active neurons.

Use `migrate_active_neurons_to_stable_memory` for whether the active
neurons should be in stable memory, so that the migration task in the
timer would move them accordingly.

## Launch/Rollback Plan

With the new setup, we can:

1. Turn on `allow_active_neurons_in_stable_memory`. If any problems
occur, we can roll it back.
2. If there is no problem found, we can turn on
`migrate_active_neurons_to_stable_memory`. If any problems occur, we can
roll back `migrate_active_neurons_to_stable_memory` and a later PR
(which should be merged before launching) will do the reverse migration.
We cannot rollback the 2 flags simultaneously (refer to the "Why"
section above).
3. If we have to rollback `allow_active_neurons_in_stable_memory` at
this point (however unlikely), we can rollback
`migrate_active_neurons_to_stable_memory` first, and do another release
to rollback `allow_active_neurons_in_stable_memory` after we confirm
that the reverse migration has been completed.

Author: jasonz-dfinity

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

@@ -7,9 +7,11 @@ use crate::{
         KnownNeuron, ListNeurons, Neuron as NeuronProto, ProposalData, Topic, Vote,
         VotingPowerEconomics,
     },
-    temporarily_disable_active_neurons_in_stable_memory,
+    temporarily_disable_allow_active_neurons_in_stable_memory,
+    temporarily_disable_migrate_active_neurons_to_stable_memory,
     temporarily_disable_stable_memory_following_index,
-    temporarily_enable_active_neurons_in_stable_memory,
+    temporarily_enable_allow_active_neurons_in_stable_memory,
+    temporarily_enable_migrate_active_neurons_to_stable_memory,
     temporarily_enable_stable_memory_following_index,
     test_utils::{MockEnvironment, StubCMC, StubIcpLedger},
 };
@@ -392,8 +394,9 @@ fn make_neuron(
 
 #[bench(raw)]
 fn cascading_vote_stable_neurons_with_heap_index() -> BenchResult {
-    let _a = temporarily_enable_active_neurons_in_stable_memory();
+    let _a = temporarily_enable_allow_active_neurons_in_stable_memory();
     let _b = temporarily_disable_stable_memory_following_index();
+    let _c = temporarily_enable_migrate_active_neurons_to_stable_memory();
 
     cast_vote_cascade_helper(
         SetUpStrategy::Chain {
@@ -406,8 +409,9 @@ fn cascading_vote_stable_neurons_with_heap_index() -> BenchResult {
 
 #[bench(raw)]
 fn cascading_vote_stable_everything() -> BenchResult {
-    let _a = temporarily_enable_active_neurons_in_stable_memory();
+    let _a = temporarily_enable_allow_active_neurons_in_stable_memory();
     let _b = temporarily_enable_stable_memory_following_index();
+    let _c = temporarily_enable_migrate_active_neurons_to_stable_memory();
 
     cast_vote_cascade_helper(
         SetUpStrategy::Chain {
@@ -420,8 +424,9 @@ fn cascading_vote_stable_everything() -> BenchResult {
 
 #[bench(raw)]
 fn cascading_vote_all_heap() -> BenchResult {
-    let _a = temporarily_disable_active_neurons_in_stable_memory();
+    let _a = temporarily_disable_allow_active_neurons_in_stable_memory();
     let _b = temporarily_disable_stable_memory_following_index();
+    let _c = temporarily_disable_migrate_active_neurons_to_stable_memory();
 
     cast_vote_cascade_helper(
         SetUpStrategy::Chain {
@@ -434,8 +439,9 @@ fn cascading_vote_all_heap() -> BenchResult {
 
 #[bench(raw)]
 fn cascading_vote_heap_neurons_stable_index() -> BenchResult {
-    let _a = temporarily_disable_active_neurons_in_stable_memory();
+    let _a = temporarily_disable_allow_active_neurons_in_stable_memory();
     let _b = temporarily_enable_stable_memory_following_index();
+    let _c = temporarily_disable_migrate_active_neurons_to_stable_memory();
 
     cast_vote_cascade_helper(
         SetUpStrategy::Chain {
@@ -448,8 +454,9 @@ fn cascading_vote_heap_neurons_stable_index() -> BenchResult {
 
 #[bench(raw)]
 fn single_vote_all_stable() -> BenchResult {
-    let _a = temporarily_enable_active_neurons_in_stable_memory();
+    let _a = temporarily_enable_allow_active_neurons_in_stable_memory();
     let _b = temporarily_enable_stable_memory_following_index();
+    let _c = temporarily_enable_migrate_active_neurons_to_stable_memory();
 
     cast_vote_cascade_helper(
         SetUpStrategy::SingleVote { num_neurons: 151 },
@@ -459,8 +466,9 @@ fn single_vote_all_stable() -> BenchResult {
 
 #[bench(raw)]
 fn centralized_following_all_stable() -> BenchResult {
-    let _a = temporarily_enable_active_neurons_in_stable_memory();
+    let _a = temporarily_enable_allow_active_neurons_in_stable_memory();
     let _b = temporarily_enable_stable_memory_following_index();
+    let _c = temporarily_enable_migrate_active_neurons_to_stable_memory();
 
     cast_vote_cascade_helper(
         SetUpStrategy::Centralized { num_neurons: 151 },
@@ -472,7 +480,8 @@ fn centralized_following_all_stable() -> BenchResult {
 fn compute_ballots_for_new_proposal_with_stable_neurons() -> BenchResult {
     let now_seconds = 1732817584;
 
-    let _f = temporarily_enable_active_neurons_in_stable_memory();
+    let _a = temporarily_enable_allow_active_neurons_in_stable_memory();
+    let _b = temporarily_enable_migrate_active_neurons_to_stable_memory();
     let neurons = (0..100)
         .map(|id| {
             (
@@ -558,13 +567,15 @@ fn list_neurons_benchmark() -> BenchResult {
 /// Benchmark list_neurons
 #[bench(raw)]
 fn list_neurons_stable() -> BenchResult {
-    let _f = temporarily_enable_active_neurons_in_stable_memory();
+    let _a = temporarily_enable_allow_active_neurons_in_stable_memory();
+    let _b = temporarily_enable_migrate_active_neurons_to_stable_memory();
     list_neurons_benchmark()
 }
 
 /// Benchmark list_neurons
 #[bench(raw)]
 fn list_neurons_heap() -> BenchResult {
-    let _f = temporarily_disable_active_neurons_in_stable_memory();
+    let _a = temporarily_disable_allow_active_neurons_in_stable_memory();
+    let _b = temporarily_disable_migrate_active_neurons_to_stable_memory();
     list_neurons_benchmark()
 }

rs/nns/governance/src/lib.rs

@@ -209,9 +209,11 @@ thread_local! {
 
     static ARE_SET_VISIBILITY_PROPOSALS_ENABLED: Cell<bool> = const { Cell::new(true) };
 
-    static ACTIVE_NEURONS_IN_STABLE_MEMORY_ENABLED: Cell<bool> = const { Cell::new(cfg!(feature = "test")) };
+    static ALLOW_ACTIVE_NEURONS_IN_STABLE_MEMORY: Cell<bool> = const { Cell::new(cfg!(feature = "test")) };
 
     static USE_STABLE_MEMORY_FOLLOWING_INDEX: Cell<bool> = const { Cell::new(cfg!(feature = "test")) };
+
+    static MIGRATE_ACTIVE_NEURONS_TO_STABLE_MEMORY: Cell<bool> = const { Cell::new(cfg!(feature = "test")) };
 }
 
 thread_local! {
@@ -285,20 +287,20 @@ pub fn temporarily_disable_set_visibility_proposals() -> Temporary {
     Temporary::new(&ARE_SET_VISIBILITY_PROPOSALS_ENABLED, false)
 }
 
-pub fn is_active_neurons_in_stable_memory_enabled() -> bool {
-    ACTIVE_NEURONS_IN_STABLE_MEMORY_ENABLED.with(|ok| ok.get())
+pub fn allow_active_neurons_in_stable_memory() -> bool {
+    ALLOW_ACTIVE_NEURONS_IN_STABLE_MEMORY.with(|ok| ok.get())
 }
 
 /// Only integration tests should use this.
 #[cfg(any(test, feature = "canbench-rs", feature = "test"))]
-pub fn temporarily_enable_active_neurons_in_stable_memory() -> Temporary {
-    Temporary::new(&ACTIVE_NEURONS_IN_STABLE_MEMORY_ENABLED, true)
+pub fn temporarily_enable_allow_active_neurons_in_stable_memory() -> Temporary {
+    Temporary::new(&ALLOW_ACTIVE_NEURONS_IN_STABLE_MEMORY, true)
 }
 
 /// Only integration tests should use this.
 #[cfg(any(test, feature = "canbench-rs", feature = "test"))]
-pub fn temporarily_disable_active_neurons_in_stable_memory() -> Temporary {
-    Temporary::new(&ACTIVE_NEURONS_IN_STABLE_MEMORY_ENABLED, false)
+pub fn temporarily_disable_allow_active_neurons_in_stable_memory() -> Temporary {
+    Temporary::new(&ALLOW_ACTIVE_NEURONS_IN_STABLE_MEMORY, false)
 }
 
 pub fn use_stable_memory_following_index() -> bool {
@@ -317,6 +319,22 @@ pub fn temporarily_disable_stable_memory_following_index() -> Temporary {
     Temporary::new(&USE_STABLE_MEMORY_FOLLOWING_INDEX, false)
 }
 
+pub fn migrate_active_neurons_to_stable_memory() -> bool {
+    MIGRATE_ACTIVE_NEURONS_TO_STABLE_MEMORY.with(|ok| ok.get())
+}
+
+/// Only integration tests should use this.
+#[cfg(any(test, feature = "canbench-rs", feature = "test"))]
+pub fn temporarily_enable_migrate_active_neurons_to_stable_memory() -> Temporary {
+    Temporary::new(&MIGRATE_ACTIVE_NEURONS_TO_STABLE_MEMORY, true)
+}
+
+/// Only integration tests should use this.
+#[cfg(any(test, feature = "canbench-rs", feature = "test"))]
+pub fn temporarily_disable_migrate_active_neurons_to_stable_memory() -> Temporary {
+    Temporary::new(&MIGRATE_ACTIVE_NEURONS_TO_STABLE_MEMORY, false)
+}
+
 pub fn decoder_config() -> DecoderConfig {
     let mut config = DecoderConfig::new();
     config.set_skipping_quota(DEFAULT_SKIPPING_QUOTA);

rs/nns/governance/src/neuron_data_validation.rs

@@ -1,5 +1,5 @@
 use crate::{
-    is_active_neurons_in_stable_memory_enabled,
+    allow_active_neurons_in_stable_memory,
     neuron::Neuron,
     neuron_store::NeuronStore,
     pb::v1::Topic,
@@ -225,7 +225,7 @@ impl ValidationInProgress {
             KnownNeuronIndexValidator,
         >::new()));
 
-        if !is_active_neurons_in_stable_memory_enabled() {
+        if !allow_active_neurons_in_stable_memory() {
             tasks.push_back(Box::new(StableNeuronStoreValidator::new(
                 INACTIVE_NEURON_VALIDATION_CHUNK_SIZE,
             )));
@@ -688,7 +688,7 @@ mod tests {
         neuron::{DissolveStateAndAge, NeuronBuilder},
         pb::v1::{neuron::Followees, KnownNeuronData},
         storage::{with_stable_neuron_indexes_mut, with_stable_neuron_store_mut},
-        temporarily_disable_active_neurons_in_stable_memory,
+        temporarily_disable_allow_active_neurons_in_stable_memory,
     };
 
     thread_local! {
@@ -996,7 +996,7 @@ mod tests {
 
     #[test]
     fn test_validator_invalid_issues_active_neuron_in_stable() {
-        let _t = temporarily_disable_active_neurons_in_stable_memory();
+        let _t = temporarily_disable_allow_active_neurons_in_stable_memory();
 
         // Step 1: Cause an issue with active neuron in stable storage.
         // Step 1.1 First create it as an inactive neuron so it can be added to stable storage.

rs/nns/governance/src/neuron_store.rs

@@ -1,8 +1,9 @@
 use crate::{
+    allow_active_neurons_in_stable_memory,
     governance::{
         Environment, TimeWarp, LOG_PREFIX, MIN_DISSOLVE_DELAY_FOR_VOTE_ELIGIBILITY_SECONDS,
     },
-    is_active_neurons_in_stable_memory_enabled,
+    migrate_active_neurons_to_stable_memory,
     neuron::types::Neuron,
     neurons_fund::neurons_fund_neuron::pick_most_important_hotkeys,
     pb::v1::{
@@ -319,10 +320,14 @@ pub struct NeuronStore {
     // NeuronStore) implements additional traits. Therefore, more elaborate wrapping is needed.
     clock: Box<dyn PracticalClock>,
 
-    // Whether to use stable memory for all neurons. This is a temporary flag to change the mode
-    // of operation for the NeuronStore.  Once all neurons are in stable memory, this will be
-    // removed, as well as heap_neurons.
-    use_stable_memory_for_all_neurons: bool,
+    // Whether to allow active neurons in stable memory. When this is true, when finding/iterating
+    // through active neurons, we need to check both heap and stable memory.
+    allow_active_neurons_in_stable_memory: bool,
+
+    /// Whether to migrate active neurons to stable memory. This is a temporary flag to change the
+    /// mode of operation for the NeuronStore. Once all neurons are in stable memory, this will be
+    /// removed.
+    migrate_active_neurons_to_stable_memory: bool,
 
     // Temporary flag to determine which following index to use
     use_stable_following_index: bool,
@@ -338,8 +343,9 @@ impl PartialEq for NeuronStore {
             heap_neurons,
             topic_followee_index,
             clock: _,
-            use_stable_memory_for_all_neurons: _,
+            allow_active_neurons_in_stable_memory: _,
             use_stable_following_index: _,
+            migrate_active_neurons_to_stable_memory: _,
         } = self;
 
         *heap_neurons == other.heap_neurons && *topic_followee_index == other.topic_followee_index
@@ -352,8 +358,9 @@ impl Default for NeuronStore {
             heap_neurons: BTreeMap::new(),
             topic_followee_index: HeapNeuronFollowingIndex::new(BTreeMap::new()),
             clock: Box::new(IcClock::new()),
-            use_stable_memory_for_all_neurons: false,
+            allow_active_neurons_in_stable_memory: false,
             use_stable_following_index: false,
+            migrate_active_neurons_to_stable_memory: false,
         }
     }
 }
@@ -368,8 +375,9 @@ impl NeuronStore {
             heap_neurons: BTreeMap::new(),
             topic_followee_index: HeapNeuronFollowingIndex::new(BTreeMap::new()),
             clock: Box::new(IcClock::new()),
-            use_stable_memory_for_all_neurons: is_active_neurons_in_stable_memory_enabled(),
+            allow_active_neurons_in_stable_memory: allow_active_neurons_in_stable_memory(),
             use_stable_following_index: use_stable_memory_following_index(),
+            migrate_active_neurons_to_stable_memory: migrate_active_neurons_to_stable_memory(),
         };
 
         // Adds the neurons one by one into neuron store.
@@ -400,8 +408,9 @@ impl NeuronStore {
                 .collect(),
             topic_followee_index: proto_to_heap_topic_followee_index(topic_followee_index),
             clock,
-            use_stable_memory_for_all_neurons: is_active_neurons_in_stable_memory_enabled(),
+            allow_active_neurons_in_stable_memory: allow_active_neurons_in_stable_memory(),
             use_stable_following_index: use_stable_memory_following_index(),
+            migrate_active_neurons_to_stable_memory: migrate_active_neurons_to_stable_memory(),
         }
     }
 
@@ -515,7 +524,7 @@ impl NeuronStore {
     // has changed, e.g. after an upgrade (2) the neuron was active, but becomes inactive due to
     // passage of time.
     fn target_storage_location(&self, neuron: &Neuron) -> StorageLocation {
-        if self.use_stable_memory_for_all_neurons || neuron.is_inactive(self.now()) {
+        if self.migrate_active_neurons_to_stable_memory || neuron.is_inactive(self.now()) {
             StorageLocation::Stable
         } else {
             StorageLocation::Heap
@@ -851,7 +860,7 @@ impl NeuronStore {
         callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Cow<Neuron>> + 'b>) -> R,
         sections: NeuronSections,
     ) -> R {
-        if self.use_stable_memory_for_all_neurons {
+        if self.allow_active_neurons_in_stable_memory {
             // Note, during migration, we still need heap_neurons, so we chain them onto the iterator
             with_stable_neuron_store(|stable_store| {
                 let now = self.now();
@@ -1218,10 +1227,10 @@ impl NeuronStore {
         neuron_id: &NeuronId,
         modify: impl FnOnce(u64) -> Result<u64, String>,
     ) -> Result<(), NeuronStoreError> {
-        // When `use_stable_memory_for_all_neurons` is true, all the neurons SHOULD be in the stable
+        // When `allow_active_neurons_in_stable_memory` is true, all the neurons SHOULD be in the stable
         // neuron store. Therefore, there is no need to move the neuron between heap/stable as it
         // might become active/inactive due to the change of maturity.
-        if self.use_stable_memory_for_all_neurons {
+        if self.allow_active_neurons_in_stable_memory {
             // The validity of this approach is based on the assumption that none of the neuron
             // indexes can be affected by its maturity.
             if self.heap_neurons.contains_key(&neuron_id.id) {
@@ -1342,7 +1351,7 @@ impl NeuronStore {
 
                     let is_neuron_inactive = neuron.is_inactive(self.now());
 
-                    if self.use_stable_memory_for_all_neurons || is_neuron_inactive {
+                    if self.allow_active_neurons_in_stable_memory || is_neuron_inactive {
                         None
                     } else {
                         // An active neuron in stable neuron store is invalid.