Merge branch 'jason/NNS1-2991-2' into 'master'

feat(nns): NNS1-2991 Implement migration code to normalize neuron state & age

Go through all the neurons with legacy states and normalize them into validate states, during post upgrade.

More context is described in https://forum.dfinity.org/t/simplify-neuron-state-age/30527

Closes NNS1-2991 

Closes NNS1-2991

See merge request dfinity-lab/public/ic!19299

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

@@ -2234,8 +2234,6 @@ message Governance {
     // Migrates neuron indexes to stable storage.
     Migration neuron_indexes_migration = 1;
     Migration copy_inactive_neurons_to_stable_memory_migration = 2;
-
-    // TODO(NNS1-2533): Migration delete_inactive_neurons_from_heap = 3;
   }
 
   // Migration related data.
@@ -2555,10 +2553,12 @@ message AuditEvent {
   uint64 timestamp_seconds = 1;
 
   oneof payload {
-    // Reset aging timestamps for https://forum.dfinity.org/t/icp-neuron-age-is-52-years/21261/26
+    // Reset aging timestamps (https://forum.dfinity.org/t/icp-neuron-age-is-52-years/21261/26).
     ResetAging reset_aging = 2;
-    // Restore aging timestamp that were incorrectly reset.
+    // Restore aging timestamp that were incorrectly reset (https://forum.dfinity.org/t/restore-neuron-age-in-proposal-129394/29840).
     RestoreAging restore_aging = 3;
+    // Normalize neuron dissolve state and age (https://forum.dfinity.org/t/simplify-neuron-state-age/30527)
+    NormalizeDissolveStateAndAge normalize_dissolve_state_and_age = 4;
   }
 
   message ResetAging {
@@ -2600,6 +2600,30 @@ message AuditEvent {
     // Neuron's stake at the time of restore.
     optional uint64 neuron_stake_e8s = 6;
   }
+
+  enum NeuronLegacyCase {
+    NEURON_LEGACY_CASE_UNSPECIFIED = 0;
+    // Neuron is dissolving or dissolved but with a non-zero age.
+    NEURON_LEGACY_CASE_DISSOLVING_OR_DISSOLVED = 1;
+    // Neuron is dissolved with DissolveDelaySeconds(0).
+    NEURON_LEGACY_CASE_DISSOLVED = 2;
+    // Neuron has a None dissolve state.
+    NEURON_LEGACY_CASE_NONE_DISSOLVE_STATE = 3;
+  }
+
+  message NormalizeDissolveStateAndAge {
+    // The neuron id whose dissolve state and age were normalized.
+    optional uint64 neuron_id = 1;
+
+    // Which legacy case the neuron falls into.
+    NeuronLegacyCase neuron_legacy_case = 2;
+
+    // Previous when_dissolved_timestamp_seconds if the neuron was dissolving or dissolved.
+    optional uint64 previous_when_dissolved_timestamp_seconds = 3;
+
+    // Previous aging_since_timestamp_seconds.
+    optional uint64 previous_aging_since_timestamp_seconds = 4;
+  }
 }
 
 // The summary of the restore aging event.

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

@@ -3080,7 +3080,7 @@ pub struct AuditEvent {
     /// The timestamp of the event.
     #[prost(uint64, tag = "1")]
     pub timestamp_seconds: u64,
-    #[prost(oneof = "audit_event::Payload", tags = "2, 3")]
+    #[prost(oneof = "audit_event::Payload", tags = "2, 3, 4")]
     pub payload: ::core::option::Option<audit_event::Payload>,
 }
 /// Nested message and enum types in `AuditEvent`.
@@ -3157,14 +3157,85 @@ pub mod audit_event {
     }
     #[derive(candid::CandidType, candid::Deserialize, serde::Serialize, comparable::Comparable)]
     #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct NormalizeDissolveStateAndAge {
+        /// The neuron id whose dissolve state and age were normalized.
+        #[prost(uint64, optional, tag = "1")]
+        pub neuron_id: ::core::option::Option<u64>,
+        /// Which legacy case the neuron falls into.
+        #[prost(enumeration = "NeuronLegacyCase", tag = "2")]
+        pub neuron_legacy_case: i32,
+        /// Previous when_dissolved_timestamp_seconds if the neuron was dissolving or dissolved.
+        #[prost(uint64, optional, tag = "3")]
+        pub previous_when_dissolved_timestamp_seconds: ::core::option::Option<u64>,
+        /// Previous aging_since_timestamp_seconds.
+        #[prost(uint64, optional, tag = "4")]
+        pub previous_aging_since_timestamp_seconds: ::core::option::Option<u64>,
+    }
+    #[derive(
+        candid::CandidType,
+        candid::Deserialize,
+        serde::Serialize,
+        comparable::Comparable,
+        Clone,
+        Copy,
+        Debug,
+        PartialEq,
+        Eq,
+        Hash,
+        PartialOrd,
+        Ord,
+        ::prost::Enumeration,
+    )]
+    #[repr(i32)]
+    pub enum NeuronLegacyCase {
+        Unspecified = 0,
+        /// Neuron is dissolving or dissolved but with a non-zero age.
+        DissolvingOrDissolved = 1,
+        /// Neuron is dissolved with DissolveDelaySeconds(0).
+        Dissolved = 2,
+        /// Neuron has a None dissolve state.
+        NoneDissolveState = 3,
+    }
+    impl NeuronLegacyCase {
+        /// String value of the enum field names used in the ProtoBuf definition.
+        ///
+        /// The values are not transformed in any way and thus are considered stable
+        /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+        pub fn as_str_name(&self) -> &'static str {
+            match self {
+                NeuronLegacyCase::Unspecified => "NEURON_LEGACY_CASE_UNSPECIFIED",
+                NeuronLegacyCase::DissolvingOrDissolved => {
+                    "NEURON_LEGACY_CASE_DISSOLVING_OR_DISSOLVED"
+                }
+                NeuronLegacyCase::Dissolved => "NEURON_LEGACY_CASE_DISSOLVED",
+                NeuronLegacyCase::NoneDissolveState => "NEURON_LEGACY_CASE_NONE_DISSOLVE_STATE",
+            }
+        }
+        /// Creates an enum from field names used in the ProtoBuf definition.
+        pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+            match value {
+                "NEURON_LEGACY_CASE_UNSPECIFIED" => Some(Self::Unspecified),
+                "NEURON_LEGACY_CASE_DISSOLVING_OR_DISSOLVED" => Some(Self::DissolvingOrDissolved),
+                "NEURON_LEGACY_CASE_DISSOLVED" => Some(Self::Dissolved),
+                "NEURON_LEGACY_CASE_NONE_DISSOLVE_STATE" => Some(Self::NoneDissolveState),
+                _ => None,
+            }
+        }
+    }
+    #[derive(candid::CandidType, candid::Deserialize, serde::Serialize, comparable::Comparable)]
+    #[allow(clippy::derive_partial_eq_without_eq)]
     #[derive(Clone, PartialEq, ::prost::Oneof)]
     pub enum Payload {
-        /// Reset aging timestamps for <https://forum.dfinity.org/t/icp-neuron-age-is-52-years/21261/26>
+        /// Reset aging timestamps (<https://forum.dfinity.org/t/icp-neuron-age-is-52-years/21261/26>).
         #[prost(message, tag = "2")]
         ResetAging(ResetAging),
-        /// Restore aging timestamp that were incorrectly reset.
+        /// Restore aging timestamp that were incorrectly reset (<https://forum.dfinity.org/t/restore-neuron-age-in-proposal-129394/29840>).
         #[prost(message, tag = "3")]
         RestoreAging(RestoreAging),
+        /// Normalize neuron dissolve state and age (<https://forum.dfinity.org/t/simplify-neuron-state-age/30527>)
+        #[prost(message, tag = "4")]
+        NormalizeDissolveStateAndAge(NormalizeDissolveStateAndAge),
     }
 }
 /// The summary of the restore aging event.

rs/nns/governance/src/governance/merge_neurons.rs

@@ -1605,11 +1605,11 @@ mod tests {
             } else {
                 panic!("Source neuron should not stop dissolving after merging");
             }
-            let target_state_and_age = effect.target_neuron_dissolve_state_and_age;
+            let target_dissolve_state_and_age = effect.target_neuron_dissolve_state_and_age;
             if let DissolveStateAndAge::NotDissolving {
                 dissolve_delay_seconds,
                 aging_since_timestamp_seconds,
-            } = target_state_and_age
+            } = target_dissolve_state_and_age
             {
                 prop_assert!(dissolve_delay_seconds >= target_dissolve_delay_seconds);
                 // The resulted age should be between the source and target ages.

rs/nns/governance/src/neuron/dissolve_state_and_age.rs

@@ -1,4 +1,7 @@
-use crate::{governance::MAX_DISSOLVE_DELAY_SECONDS, pb::v1::NeuronState};
+use crate::{
+    governance::MAX_DISSOLVE_DELAY_SECONDS,
+    pb::v1::{audit_event::NeuronLegacyCase, NeuronState},
+};
 
 /// An enum to represent different combinations of a neurons dissolve_state and
 /// aging_since_timestamp_seconds. Currently, the back-and-forth conversions should make sure the
@@ -38,6 +41,16 @@ pub enum DissolveStateAndAge {
 }
 
 impl DissolveStateAndAge {
+    /// Returns true if the neuron is in a legacy state.
+    pub fn is_legacy(self) -> bool {
+        match self {
+            Self::NotDissolving { .. } | Self::DissolvingOrDissolved { .. } => false,
+            Self::LegacyDissolvingOrDissolved { .. }
+            | Self::LegacyDissolved { .. }
+            | Self::LegacyNoneDissolveState { .. } => true,
+        }
+    }
+
     /// Returns the current state given the current time. Mainly for differentiating between
     /// dissolving and dissolved neurons.
     pub fn current_state(self, now_seconds: u64) -> NeuronState {
@@ -323,6 +336,41 @@ impl DissolveStateAndAge {
             },
         }
     }
+
+    /// Returns the normalized neuron dissolve state and age. If the neuron is in a legacy state, it
+    /// returns the valid state, as well as the audit event log for logging purposes. Otherwise it
+    /// returns the existing state and None.
+    pub fn normalize(self, created_timestamp_seconds: u64) -> Option<(Self, NeuronLegacyCase)> {
+        match self {
+            Self::NotDissolving { .. } | Self::DissolvingOrDissolved { .. } => None,
+
+            Self::LegacyDissolvingOrDissolved {
+                when_dissolved_timestamp_seconds,
+                aging_since_timestamp_seconds: _,
+            } => Some((
+                Self::DissolvingOrDissolved {
+                    when_dissolved_timestamp_seconds,
+                },
+                NeuronLegacyCase::DissolvingOrDissolved,
+            )),
+
+            Self::LegacyDissolved { .. } => Some((
+                Self::DissolvingOrDissolved {
+                    when_dissolved_timestamp_seconds: created_timestamp_seconds,
+                },
+                NeuronLegacyCase::Dissolved,
+            )),
+
+            // This case should be impossible, but treating it the same way as LegacyDissolved is
+            // also reasonable.
+            Self::LegacyNoneDissolveState { .. } => Some((
+                Self::DissolvingOrDissolved {
+                    when_dissolved_timestamp_seconds: created_timestamp_seconds,
+                },
+                NeuronLegacyCase::NoneDissolveState,
+            )),
+        }
+    }
 }
 
 #[cfg(test)]
@@ -568,7 +616,7 @@ mod tests {
 
     #[test]
     fn test_increase_dissolve_delay_for_dissolved_neurons() {
-        for dissolved_state_and_age in [
+        for dissolve_state_and_age in [
             // Valid cases
             DissolveStateAndAge::DissolvingOrDissolved {
                 when_dissolved_timestamp_seconds: NOW,
@@ -624,7 +672,7 @@ mod tests {
             },
         ] {
             assert_increase_dissolve_delay(
-                dissolved_state_and_age,
+                dissolve_state_and_age,
                 1,
                 DissolveStateAndAge::NotDissolving {
                     dissolve_delay_seconds: 1,
@@ -634,7 +682,7 @@ mod tests {
 
             // Test that the dissolve delay is capped at MAX_DISSOLVE_DELAY_SECONDS.
             assert_increase_dissolve_delay(
-                dissolved_state_and_age,
+                dissolve_state_and_age,
                 MAX_DISSOLVE_DELAY_SECONDS as u32 + 1000,
                 DissolveStateAndAge::NotDissolving {
                     dissolve_delay_seconds: MAX_DISSOLVE_DELAY_SECONDS,
@@ -910,4 +958,75 @@ mod tests {
             }
         );
     }
+
+    #[test]
+    fn test_normalize() {
+        let created_timestamp_seconds = 123_456_789;
+
+        {
+            let normal_cases = vec![
+                DissolveStateAndAge::NotDissolving {
+                    dissolve_delay_seconds: 1000,
+                    aging_since_timestamp_seconds: NOW - 100,
+                },
+                DissolveStateAndAge::DissolvingOrDissolved {
+                    when_dissolved_timestamp_seconds: NOW + 1000,
+                },
+                DissolveStateAndAge::DissolvingOrDissolved {
+                    when_dissolved_timestamp_seconds: NOW - 1000,
+                },
+            ];
+            for normal_case in normal_cases {
+                assert_eq!(normal_case.normalize(created_timestamp_seconds), None);
+            }
+        }
+
+        // In the legacy dissolving/dissolved case, we remove the age bonus.
+        let legacy_dissolving_or_dissolved = DissolveStateAndAge::LegacyDissolvingOrDissolved {
+            when_dissolved_timestamp_seconds: NOW + 1000,
+            aging_since_timestamp_seconds: NOW - 1000,
+        };
+        let (normalized, legacy_case) = legacy_dissolving_or_dissolved
+            .normalize(created_timestamp_seconds)
+            .unwrap();
+        assert_eq!(
+            normalized,
+            DissolveStateAndAge::DissolvingOrDissolved {
+                when_dissolved_timestamp_seconds: NOW + 1000
+            }
+        );
+        assert_eq!(legacy_case, NeuronLegacyCase::DissolvingOrDissolved);
+
+        // In the legacy dissolved case, we normalize it to dissolving/dissolved case while setting
+        // the dissolved timestamp to the creation timestamp.
+        let legacy_dissolved = DissolveStateAndAge::LegacyDissolved {
+            aging_since_timestamp_seconds: NOW - 1000,
+        };
+        let (normalized, legacy_case) = legacy_dissolved
+            .normalize(created_timestamp_seconds)
+            .unwrap();
+        assert_eq!(
+            normalized,
+            DissolveStateAndAge::DissolvingOrDissolved {
+                when_dissolved_timestamp_seconds: created_timestamp_seconds
+            }
+        );
+        assert_eq!(legacy_case, NeuronLegacyCase::Dissolved);
+
+        // In the legacy none-dissolve state case, we normalize it to dissolving/dissolved case while setting
+        // the dissolved timestamp to the creation timestamp.
+        let legacy_none_dissolve_state = DissolveStateAndAge::LegacyNoneDissolveState {
+            aging_since_timestamp_seconds: NOW - 1000,
+        };
+        let (normalized, legacy_case) = legacy_none_dissolve_state
+            .normalize(created_timestamp_seconds)
+            .unwrap();
+        assert_eq!(
+            normalized,
+            DissolveStateAndAge::DissolvingOrDissolved {
+                when_dissolved_timestamp_seconds: created_timestamp_seconds
+            }
+        );
+        assert_eq!(legacy_case, NeuronLegacyCase::NoneDissolveState);
+    }
 }

rs/nns/governance/src/neuron/mod.rs

@@ -356,8 +356,9 @@ impl Neuron {
     fn start_dissolving(&mut self, now_seconds: u64) -> Result<(), GovernanceError> {
         let dissolve_state_and_age = self.dissolve_state_and_age();
         if let DissolveStateAndAge::NotDissolving { .. } = dissolve_state_and_age {
-            let new_disolved_state_and_age = dissolve_state_and_age.start_dissolving(now_seconds);
-            self.set_dissolve_state_and_age(new_disolved_state_and_age);
+            let new_disolved_dissolve_state_and_age =
+                dissolve_state_and_age.start_dissolving(now_seconds);
+            self.set_dissolve_state_and_age(new_disolved_dissolve_state_and_age);
             Ok(())
         } else {
             Err(GovernanceError::new(ErrorType::RequiresNotDissolving))
@@ -369,11 +370,12 @@ impl Neuron {
     /// If the neuron is not dissolving, an error is returned.
     fn stop_dissolving(&mut self, now_seconds: u64) -> Result<(), GovernanceError> {
         let dissolve_state_and_age = self.dissolve_state_and_age();
-        let new_disolved_state_and_age = dissolve_state_and_age.stop_dissolving(now_seconds);
-        if new_disolved_state_and_age == dissolve_state_and_age {
+        let new_disolved_dissolve_state_and_age =
+            dissolve_state_and_age.stop_dissolving(now_seconds);
+        if new_disolved_dissolve_state_and_age == dissolve_state_and_age {
             Err(GovernanceError::new(ErrorType::RequiresDissolving))
         } else {
-            self.set_dissolve_state_and_age(new_disolved_state_and_age);
+            self.set_dissolve_state_and_age(new_disolved_dissolve_state_and_age);
             Ok(())
         }
     }
@@ -669,10 +671,10 @@ impl Neuron {
             self.cached_neuron_stake_e8s = new_stake_e8s;
 
             let new_aging_since_timestamp_seconds = now.saturating_sub(new_age_seconds);
-            let new_disolved_state_and_age = self
+            let new_disolved_dissolve_state_and_age = self
                 .dissolve_state_and_age()
                 .adjust_age(new_aging_since_timestamp_seconds);
-            self.set_dissolve_state_and_age(new_disolved_state_and_age);
+            self.set_dissolve_state_and_age(new_disolved_dissolve_state_and_age);
         }
     }