node: add PoS migration to fix total active stake

Author: tzemanovic

crates/node/src/shell/mod.rs

@@ -818,6 +818,22 @@ where
             _ => None,
         };
 
+        // Temporary migration code to fix the total active stake value
+        if let Some(height) =
+            std::env::var("NAMADA_MIGRATION_HEIGHT").ok().map(|height| {
+                <BlockHeight as std::str::FromStr>::from_str(&height.trim())
+                    .expect(
+                        "Invalid block height set in `NAMADA_MIGRATION_HEIGHT` env var",
+                    )
+            })
+        {
+            if height == height_to_commit {
+                proof_of_stake::fix_total_active_stake::<_, governance::Store<_>>(
+                    &mut self.state,
+                ).expect("Must be able to fix total active stake")
+            }
+        }
+
         self.state
             .commit_block()
             .expect("Encountered a storage error while committing a block");

crates/proof_of_stake/src/lib.rs

@@ -3277,6 +3277,47 @@ fn prune_old_delegations(
     Ok(())
 }
 
+/// Temporary migration code to fix the total active stake value
+pub fn fix_total_active_stake<S, Gov>(storage: &mut S) -> Result<()>
+where
+    S: StorageRead + StorageWrite,
+    Gov: governance::Read<S>,
+{
+    let epoch = storage.get_block_epoch()?;
+    let validators = storage::read_all_validator_addresses(storage, epoch)?;
+    let mut total_stake = token::Amount::zero();
+    let params = read_pos_params::<_, Gov>(storage)?;
+    for validator in validators {
+        let state = storage::read_validator_state::<_, Gov>(
+            storage, &validator, epoch,
+        )?;
+        let is_active = matches!(
+            state,
+            Some(
+                ValidatorState::Consensus
+                    | ValidatorState::BelowCapacity
+                    | ValidatorState::BelowThreshold,
+            )
+        );
+        if is_active {
+            let stake =
+                read_validator_stake(storage, &params, &validator, epoch)?;
+            total_stake =
+                total_stake.checked_add(stake).expect("Must not overflow");
+        }
+    }
+
+    let total_active_deltas = storage::total_active_deltas_handle();
+    total_active_deltas.set::<S, Gov>(
+        storage,
+        total_stake.change(),
+        epoch,
+        0,
+    )?;
+
+    Ok(())
+}
+
 #[cfg(any(test, feature = "testing"))]
 /// PoS related utility functions to help set up tests.
 pub mod test_utils {