feat(nns): Move inactive neurons in the heap to stable storage throug…

…h a timer (#2643)

# Why

This change has 2 purposes:

* While the target neuron location criteria doesn't change (inactive
neuron -> stable), some neurons can become inactive due to the passage
of time, this change migrates those neurons to stable storage in a timer
* In the future, when the stable storage for active neurons is enabled,
this change will also move the active neurons to stable storage in a
similar fashion

# What

* Removing a state machine test regarding neuron indexes. It is not
useful anymore because the migration is over.
* Adding a similar state machine test for the current use case
* Adding the capability for moving a batch of neurons from heap memory
to stable storage
* Schedule a timer to execute the migration code

rs/nns/governance/canister/canister.rs

@@ -158,6 +158,11 @@ fn set_governance(gov: Governance) {
         .expect("Error initializing the governance canister.");
 }
 
+fn schedule_timers() {
+    schedule_seeding(Duration::from_nanos(0));
+    schedule_adjust_neurons_storage(Duration::from_nanos(0), NeuronIdProto { id: 0 });
+}
+
 // Seeding interval seeks to find a balance between the need for rng secrecy, and
 // avoiding the overhead of frequent reseeding.
 const SEEDING_INTERVAL: Duration = Duration::from_secs(3600);
@@ -188,6 +193,29 @@ fn schedule_seeding(duration: Duration) {
     });
 }
 
+// The interval before adjusting neuron storage for the next batch of neurons starting from last
+// neuron id scanned in the last batch.
+const ADJUST_NEURON_STORAGE_BATCH_INTERVAL: Duration = Duration::from_secs(5);
+// The interval before adjusting neuron storage for the next round starting from the smallest neuron
+// id.
+const ADJUST_NEURON_STORAGE_ROUND_INTERVAL: Duration = Duration::from_secs(3600);
+
+fn schedule_adjust_neurons_storage(delay: Duration, start_neuron_id: NeuronIdProto) {
+    ic_cdk_timers::set_timer(delay, move || {
+        let next_neuron_id = governance_mut().batch_adjust_neurons_storage(start_neuron_id);
+        match next_neuron_id {
+            Some(next_neuron_id) => schedule_adjust_neurons_storage(
+                ADJUST_NEURON_STORAGE_BATCH_INTERVAL,
+                next_neuron_id,
+            ),
+            None => schedule_adjust_neurons_storage(
+                ADJUST_NEURON_STORAGE_ROUND_INTERVAL,
+                NeuronIdProto { id: 0 },
+            ),
+        };
+    });
+}
+
 struct CanisterEnv {
     rng: Option<ChaCha20Rng>,
     time_warp: GovTimeWarp,
@@ -408,7 +436,7 @@ fn canister_init_(init_payload: ApiGovernanceProto) {
         init_payload.neurons.len()
     );
 
-    schedule_seeding(Duration::from_nanos(0));
+    schedule_timers();
     set_governance(Governance::new(
         InternalGovernanceProto::from(init_payload),
         Box::new(CanisterEnv::new()),
@@ -452,7 +480,7 @@ fn canister_post_upgrade() {
         restored_state.xdr_conversion_rate,
     );
 
-    schedule_seeding(Duration::from_nanos(0));
+    schedule_timers();
     set_governance(Governance::new_restored(
         restored_state,
         Box::new(CanisterEnv::new()),

rs/nns/governance/src/governance.rs

@@ -6967,6 +6967,11 @@ impl Governance {
         })
     }
 
+    pub fn batch_adjust_neurons_storage(&mut self, start_neuron_id: NeuronId) -> Option<NeuronId> {
+        self.neuron_store
+            .batch_adjust_neurons_storage(start_neuron_id)
+    }
+
     /// Recompute cached metrics once per day
     pub fn should_compute_cached_metrics(&self) -> bool {
         if let Some(metrics) = self.heap_data.metrics.as_ref() {

rs/nns/governance/src/neuron_store.rs

@@ -206,6 +206,7 @@ pub struct NeuronsFundNeuron {
     pub hotkeys: Vec<PrincipalId>,
 }
 
+#[derive(Eq, PartialEq)]
 enum StorageLocation {
     Heap,
     Stable,
@@ -475,6 +476,18 @@ impl NeuronStore {
         heap_len + stable_len
     }
 
+    // Returns the target storage location of a neuron. It might not be the actual storage location
+    // if the neuron already exists, for 2 possible reasons: (1) the target storage location logic
+    // 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()) {
+            StorageLocation::Stable
+        } else {
+            StorageLocation::Heap
+        }
+    }
+
     /// Add a new neuron
     pub fn add_neuron(&mut self, neuron: Neuron) -> Result<NeuronId, NeuronStoreError> {
         let neuron_id = neuron.id();
@@ -485,7 +498,7 @@ impl NeuronStore {
             return Err(NeuronStoreError::NeuronAlreadyExists(neuron_id));
         }
 
-        if self.use_stable_memory_for_all_neurons || neuron.is_inactive(self.now()) {
+        if self.target_storage_location(&neuron) == StorageLocation::Stable {
             // Write as primary copy in stable storage.
             with_stable_neuron_store_mut(|stable_neuron_store| {
                 stable_neuron_store.create(neuron.clone())
@@ -566,6 +579,72 @@ impl NeuronStore {
         self.remove_neuron_from_indexes(&neuron_to_remove);
     }
 
+    /// Adjusts the storage location of neurons, since active neurons might become inactive due to
+    /// passage of time.
+    pub fn batch_adjust_neurons_storage(&mut self, start_neuron_id: NeuronId) -> Option<NeuronId> {
+        static BATCH_SIZE_FOR_MOVING_NEURONS: usize = 1000;
+
+        #[cfg(target_arch = "wasm32")]
+        static MAX_NUM_INSTRUCTIONS_PER_BATCH: u64 = 5_000_000_000;
+
+        #[cfg(target_arch = "wasm32")]
+        let max_instructions_reached =
+            || ic_cdk::api::instruction_counter() >= MAX_NUM_INSTRUCTIONS_PER_BATCH;
+
+        #[cfg(not(target_arch = "wasm32"))]
+        let max_instructions_reached = || false;
+
+        self.adjust_neuron_storage_with_max_instructions(
+            start_neuron_id,
+            BATCH_SIZE_FOR_MOVING_NEURONS,
+            max_instructions_reached,
+        )
+    }
+
+    fn adjust_neuron_storage_with_max_instructions(
+        &mut self,
+        start_neuron_id: NeuronId,
+        max_batch_size: usize,
+        max_instructions_reached: impl Fn() -> bool,
+    ) -> Option<NeuronId> {
+        // We currently only move neurons from heap to stable storage, since it's impossible to have
+        // active neurons in stable storage. In the future, we might need to move neurons from
+        // stable storage to heap as a rollback mechanism, but it is not implemented here yet.
+        let neuron_ids: Vec<_> = self
+            .heap_neurons
+            .range(start_neuron_id.id..)
+            .take(max_batch_size)
+            .map(|(id, _)| NeuronId { id: *id })
+            .collect();
+        // We know it is the last batch if the number of neurons is less than the batch size.
+        let is_last_batch = neuron_ids.len() < max_batch_size;
+
+        if neuron_ids.is_empty() {
+            return None;
+        }
+
+        let mut next_neuron_id = Some(start_neuron_id);
+
+        for neuron_id in neuron_ids {
+            if max_instructions_reached() {
+                // We don't need to look at the `is_last_batch` because at least one neuron is
+                // skipped due to instruction limit.
+                return next_neuron_id;
+            }
+
+            // We don't modify the neuron, but the below just makes sure that the neuron is in the
+            // appropriate storage location given its state and the current time.
+            let _ = self.with_neuron_mut(&neuron_id, |_| {});
+            next_neuron_id = neuron_id.next();
+        }
+
+        if is_last_batch {
+            None
+        } else {
+            next_neuron_id
+        }
+    }
+
     fn remove_neuron_from_indexes(&mut self, neuron: &Neuron) {
         let neuron_id = neuron.id();
         if let Err(error) = with_stable_neuron_indexes_mut(|indexes| indexes.remove_neuron(neuron))
@@ -653,12 +732,7 @@ impl NeuronStore {
         new_neuron: Neuron,
         previous_location: StorageLocation,
     ) -> Result<(), NeuronStoreError> {
-        let target_location =
-            if self.use_stable_memory_for_all_neurons || new_neuron.is_inactive(self.now()) {
-                StorageLocation::Stable
-            } else {
-                StorageLocation::Heap
-            };
+        let target_location = self.target_storage_location(&new_neuron);
         let is_neuron_changed = *old_neuron != new_neuron;
 
         self.validate_neuron(&new_neuron)?;

rs/nns/governance/src/neuron_store/neuron_store_tests.rs

@@ -10,6 +10,7 @@ use ic_nns_constants::GOVERNANCE_CANISTER_ID;
 use maplit::{btreemap, hashmap, hashset};
 use num_traits::bounds::LowerBounded;
 use pretty_assertions::assert_eq;
+use std::cell::Cell;
 
 static CREATED_TIMESTAMP_SECONDS: u64 = 123_456_789;
 
@@ -695,6 +696,110 @@ fn test_get_non_empty_neuron_ids_readable_by_caller() {
     );
 }
 
+#[test]
+fn test_batch_adjust_neurons_storage() {
+    // This test doesn't make sense after neurons are migrated completely to stable memory.
+    let _f = temporarily_disable_active_neurons_in_stable_memory();
+
+    // Step 1.1: set up an empty neuron store.
+    let mut neuron_store = NeuronStore::new(BTreeMap::new());
+
+    // Step 1.2: set up 5 active neurons with stake
+    for i in 1..=5 {
+        let neuron = simple_neuron_builder(i)
+            .with_cached_neuron_stake_e8s(1)
+            .with_dissolve_state_and_age(DissolveStateAndAge::DissolvingOrDissolved {
+                when_dissolved_timestamp_seconds: neuron_store.now(),
+            })
+            .build();
+        neuron_store.add_neuron(neuron).unwrap();
+    }
+
+    // Step 1.3: set up 5 active neurons without stake, which will become inactive when the time is
+    // advanced.
+    for i in 6..=10 {
+        let neuron = active_neuron_builder(i, neuron_store.now()).build();
+        neuron_store.add_neuron(neuron).unwrap();
+    }
+
+    // Step 1.4: warp time so that the neuron becomes inactive without modification.
+    warp_time_to_make_neuron_inactive(&mut neuron_store);
+
+    // Step 1.5: define a lambda which always returns false, for checking instructions.
+    let always_false = || false;
+
+    // Step 1.6: make sure the counts of neurons in heap and stable are expected.
+    assert_eq!(neuron_store.heap_neuron_store_len(), 10);
+    assert_eq!(neuron_store.stable_neuron_store_len(), 0);
+
+    // Step 2: adjust the storage of neurons for the first 6 neurons and verifies the counts. Since
+    // the first 5 neurons are active because of their stake, only 1 neuron is moved.
+    let next_neuron_id = neuron_store.adjust_neuron_storage_with_max_instructions(
+        NeuronId { id: 0 },
+        6,
+        always_false,
+    );
+    assert_eq!(next_neuron_id, Some(NeuronId { id: 7 }));
+    assert_eq!(neuron_store.heap_neuron_store_len(), 9);
+    assert_eq!(neuron_store.stable_neuron_store_len(), 1);
+
+    // Step 3: adjust the storage of neurons for the rest of 4 neurons and verifies the counts.
+    let next_neuron_id = neuron_store.adjust_neuron_storage_with_max_instructions(
+        NeuronId { id: 7 },
+        6,
+        always_false,
+    );
+    assert_eq!(next_neuron_id, None);
+    assert_eq!(neuron_store.heap_neuron_store_len(), 5);
+    assert_eq!(neuron_store.stable_neuron_store_len(), 5);
+}
+
+#[test]
+fn test_batch_adjust_neurons_storage_exceeds_instructions_limit() {
+    // This test doesn't make sense after neurons are migrated completely to stable memory.
+    let _f = temporarily_disable_active_neurons_in_stable_memory();
+
+    // Step 1.1: set up an empty neuron store.
+    let mut neuron_store = NeuronStore::new(BTreeMap::new());
+
+    // Step 1.2: set up 5 active neurons without stake, which will become inactive when the time is
+    // advanced.
+    for i in 1..=5 {
+        let neuron = active_neuron_builder(i, neuron_store.now()).build();
+        neuron_store.add_neuron(neuron).unwrap();
+    }
+
+    // Step 1.4: warp time so that the neuron becomes inactive without modification.
+    warp_time_to_make_neuron_inactive(&mut neuron_store);
+
+    // Step 1.5: make sure the counts of neurons in heap and stable are expected.
+    assert_eq!(neuron_store.heap_neuron_store_len(), 5);
+    assert_eq!(neuron_store.stable_neuron_store_len(), 0);
+
+    // Step 2: adjust the storage of neurons for the first 10 neurons, however, the instruction
+    // limit checker returns true for the 4th time it's called, allowing moving only 3 neurons.
+    let counter = Cell::new(0);
+    let next_neuron_id =
+        neuron_store.adjust_neuron_storage_with_max_instructions(NeuronId { id: 0 }, 10, || {
+            counter.set(counter.get() + 1);
+            counter.get() > 3
+        });
+    assert_eq!(next_neuron_id, Some(NeuronId { id: 4 }));
+    assert_eq!(neuron_store.heap_neuron_store_len(), 2);
+    assert_eq!(neuron_store.stable_neuron_store_len(), 3);
+
+    // Step 3: adjust the storage of neurons for the rest of 4 neurons and verifies the counts.
+    let counter = Cell::new(0);
+    let next_neuron_id =
+        neuron_store.adjust_neuron_storage_with_max_instructions(NeuronId { id: 4 }, 10, || {
+            counter.set(counter.get() + 1);
+            counter.get() > 3
+        });
+    assert_eq!(next_neuron_id, None);
+    assert_eq!(neuron_store.heap_neuron_store_len(), 0);
+    assert_eq!(neuron_store.stable_neuron_store_len(), 5);
+}
+
 #[test]
 fn test_get_full_neuron() {
     let principal_id = PrincipalId::new_user_test_id(42);