Fix TotalIssuance on migration and TotalStake check

pallets/subtensor/src/macros/hooks.rs

@@ -82,7 +82,9 @@ mod hooks {
 
         #[cfg(feature = "try-runtime")]
         fn try_state(_n: BlockNumberFor<T>) -> Result<(), sp_runtime::TryRuntimeError> {
-            Self::check_accounting_invariants()?;
+            Self::check_total_issuance()?;
+            // Disabled: https://github.com/opentensor/subtensor/pull/1166
+            // Self::check_total_stake()?;
             Ok(())
         }
     }

pallets/subtensor/src/migrations/migrate_init_total_issuance.rs

@@ -1,6 +1,6 @@
 use super::*;
-use frame_support::pallet_prelude::OptionQuery;
-use frame_support::{pallet_prelude::Identity, storage_alias};
+use frame_support::pallet_prelude::{Identity, OptionQuery, Weight};
+use frame_support::storage_alias;
 use sp_std::vec::Vec;
 
 // TODO: Implement comprehensive tests for this migration
@@ -14,10 +14,46 @@ pub mod deprecated_loaded_emission_format {
         StorageMap<Pallet<T>, Identity, u16, Vec<(AccountIdOf<T>, u64)>, OptionQuery>;
 }
 
+pub(crate) fn migrate_init_total_issuance<T: Config>() -> Weight {
+    // Calculate the total locked tokens across all subnets
+    let subnets_len = crate::SubnetLocked::<T>::iter().count() as u64;
+    let total_subnet_locked: u64 =
+        crate::SubnetLocked::<T>::iter().fold(0, |acc, (_, v)| acc.saturating_add(v));
+
+    // Retrieve the total balance of all accounts
+    let total_account_balances = <<T as crate::Config>::Currency as fungible::Inspect<
+        <T as frame_system::Config>::AccountId,
+    >>::total_issuance();
+
+    // Get the total stake from the system
+    let total_stake = crate::TotalStake::<T>::get();
+
+    // Retrieve the previous total issuance for logging purposes
+    let prev_total_issuance = crate::TotalIssuance::<T>::get();
+
+    // Calculate the new total issuance
+    let new_total_issuance = total_account_balances
+        .saturating_add(total_stake)
+        .saturating_add(total_subnet_locked);
+
+    // Update the total issuance in storage
+    crate::TotalIssuance::<T>::put(new_total_issuance);
+
+    // Log the change in total issuance
+    log::info!(
+        "Subtensor Pallet Total Issuance Updated: previous: {:?}, new: {:?}",
+        prev_total_issuance,
+        new_total_issuance
+    );
+
+    // Return the weight of the operation
+    // We performed subnets_len + 5 reads and 1 write
+    <T as frame_system::Config>::DbWeight::get().reads_writes(subnets_len.saturating_add(5), 1)
+}
+
 pub mod initialise_total_issuance {
     use frame_support::pallet_prelude::Weight;
-    use frame_support::traits::{fungible, OnRuntimeUpgrade};
-    use sp_core::Get;
+    use frame_support::traits::OnRuntimeUpgrade;
 
     use crate::*;
 
@@ -33,41 +69,7 @@ pub mod initialise_total_issuance {
         ///
         /// Returns the weight of the migration operation.
         fn on_runtime_upgrade() -> Weight {
-            // Calculate the total locked tokens across all subnets
-            let subnets_len = crate::SubnetLocked::<T>::iter().count() as u64;
-            let total_subnet_locked: u64 =
-                crate::SubnetLocked::<T>::iter().fold(0, |acc, (_, v)| acc.saturating_add(v));
-
-            // Retrieve the total balance of all accounts
-            let total_account_balances = <<T as crate::Config>::Currency as fungible::Inspect<
-                <T as frame_system::Config>::AccountId,
-            >>::total_issuance();
-
-            // Get the total stake from the system
-            let total_stake = crate::TotalStake::<T>::get();
-
-            // Retrieve the previous total issuance for logging purposes
-            let prev_total_issuance = crate::TotalIssuance::<T>::get();
-
-            // Calculate the new total issuance
-            let new_total_issuance = total_account_balances
-                .saturating_add(total_stake)
-                .saturating_add(total_subnet_locked);
-
-            // Update the total issuance in storage
-            crate::TotalIssuance::<T>::put(new_total_issuance);
-
-            // Log the change in total issuance
-            log::info!(
-                "Subtensor Pallet Total Issuance Updated: previous: {:?}, new: {:?}",
-                prev_total_issuance,
-                new_total_issuance
-            );
-
-            // Return the weight of the operation
-            // We performed subnets_len + 5 reads and 1 write
-            <T as frame_system::Config>::DbWeight::get()
-                .reads_writes(subnets_len.saturating_add(5), 1)
+            super::migrate_init_total_issuance::<T>()
         }
 
         /// Performs post-upgrade checks to ensure the migration was successful.
@@ -76,7 +78,7 @@ pub mod initialise_total_issuance {
         #[cfg(feature = "try-runtime")]
         fn post_upgrade(_state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
             // Verify that all accounting invariants are satisfied after the migration
-            crate::Pallet::<T>::check_accounting_invariants()?;
+            crate::Pallet::<T>::check_total_issuance()?;
             Ok(())
         }
     }

pallets/subtensor/src/migrations/migrate_rao.rs

@@ -1,11 +1,13 @@
-use super::*;
 use alloc::string::String;
+
 use frame_support::IterableStorageMap;
 use frame_support::{traits::Get, weights::Weight};
-use log;
 use sp_runtime::format;
 use substrate_fixed::types::U64F64;
 
+use super::*;
+use crate::subnets::subnet::POOL_INITIAL_TAO;
+
 pub fn migrate_rao<T: Config>() -> Weight {
     let migration_name = b"migrate_rao".to_vec();
 
@@ -69,12 +71,12 @@ pub fn migrate_rao<T: Config>() -> Weight {
             TokenSymbol::<T>::insert(netuid, Pallet::<T>::get_symbol_for_subnet(0));
             continue;
         }
-        let owner: T::AccountId = SubnetOwner::<T>::get(netuid);
-        let lock: u64 = SubnetLocked::<T>::get(netuid);
+        let owner = SubnetOwner::<T>::get(netuid);
+        let lock = SubnetLocked::<T>::get(netuid);
 
         // Put initial TAO from lock into subnet TAO and produce numerically equal amount of Alpha
         // The initial TAO is the locked amount, with a minimum of 1 RAO and a cap of 100 TAO.
-        let pool_initial_tao = 100_000_000_000.min(lock.max(1));
+        let pool_initial_tao = POOL_INITIAL_TAO.min(lock.max(1));
 
         let remaining_lock = lock.saturating_sub(pool_initial_tao);
         // Refund the owner for the remaining lock.
@@ -127,6 +129,10 @@ pub fn migrate_rao<T: Config>() -> Weight {
         // TargetStakesPerInterval::<T>::put(10); (DEPRECATED)
     }
 
+    // update `TotalIssuance`, because currency issuance (`T::Currency`) has changed due to lock
+    // refunds above
+    weight = weight.saturating_add(migrate_init_total_issuance::migrate_init_total_issuance::<T>());
+
     // Mark the migration as completed
     HasMigrationRun::<T>::insert(&migration_name, true);
     weight = weight.saturating_add(T::DbWeight::get().writes(1));

pallets/subtensor/src/subnets/subnet.rs

@@ -2,6 +2,8 @@ use super::*;
 use frame_support::IterableStorageMap;
 use sp_core::Get;
 
+pub(crate) const POOL_INITIAL_TAO: u64 = 100_000_000_000;
+
 impl<T: Config> Pallet<T> {
     /// Retrieves the unique identifier (UID) for the root network.
     ///
@@ -235,7 +237,7 @@ impl<T: Config> Pallet<T> {
 
         // Put initial TAO from lock into subnet TAO and produce numerically equal amount of Alpha
         // The initial TAO is the locked amount, with a minimum of 1 RAO and a cap of 100 TAO.
-        let pool_initial_tao = 100_000_000_000.min(actual_tao_lock_amount.max(1));
+        let pool_initial_tao = POOL_INITIAL_TAO.min(actual_tao_lock_amount.max(1));
 
         let actual_tao_lock_amount_less_pool_tao =
             actual_tao_lock_amount.saturating_sub(pool_initial_tao);

pallets/subtensor/src/utils/mod.rs

@@ -2,4 +2,5 @@ use super::*;
 pub mod identity;
 pub mod misc;
 pub mod rate_limiting;
+#[cfg(feature = "try-runtime")]
 pub mod try_state;

pallets/subtensor/src/utils/try_state.rs

@@ -1,59 +1,72 @@
+use frame_support::traits::fungible::Inspect;
+
 use super::*;
+use crate::subnets::subnet::POOL_INITIAL_TAO;
 
 impl<T: Config> Pallet<T> {
-    /// Checks if the accounting invariants for [`TotalStake`], [`TotalSubnetLocked`], and [`TotalIssuance`] are correct.
-    ///
-    /// This function verifies that:
-    /// 1. The sum of all stakes matches the [`TotalStake`].
-    /// 2. The [`TotalSubnetLocked`] is correctly calculated.
-    /// 3. The [`TotalIssuance`] equals the sum of currency issuance, total stake, and total subnet locked.
-    ///
-    /// # Returns
-    ///
-    /// Returns `Ok(())` if all invariants are correct, otherwise returns an error.
-    #[cfg(feature = "try-runtime")]
-    pub fn check_accounting_invariants() -> Result<(), sp_runtime::TryRuntimeError> {
-        use frame_support::traits::fungible::Inspect;
-
-        // Calculate the total staked amount
-        let mut total_staked: u64 = 0;
-        for (_hotkey, _coldkey, stake) in Stake::<T>::iter() {
-            total_staked = total_staked.saturating_add(stake);
-        }
-
-        // Verify that the calculated total stake matches the stored TotalStake
-        ensure!(
-            total_staked == TotalStake::<T>::get(),
-            "TotalStake does not match total staked",
-        );
-
+    /// Checks [`TotalIssuance`] equals the sum of currency issuance, total stake, and total subnet
+    /// locked.
+    pub(crate) fn check_total_issuance() -> Result<(), sp_runtime::TryRuntimeError> {
         // Get the total subnet locked amount
-        let total_subnet_locked: u64 = Self::get_total_subnet_locked();
+        let total_subnet_locked = Self::get_total_subnet_locked();
 
         // Get the total currency issuance
-        let currency_issuance: u64 = T::Currency::total_issuance();
+        let currency_issuance = T::Currency::total_issuance();
 
         // Calculate the expected total issuance
-        let expected_total_issuance: u64 = currency_issuance
-            .saturating_add(total_staked)
+        let expected_total_issuance = currency_issuance
+            .saturating_add(TotalStake::<T>::get())
             .saturating_add(total_subnet_locked);
 
         // Verify the diff between calculated TI and actual TI is less than delta
         //
         // These values can be off slightly due to float rounding errors.
         // They are corrected every runtime upgrade.
-        const DELTA: u64 = 1000;
-        let diff = if TotalIssuance::<T>::get() > expected_total_issuance {
-            TotalIssuance::<T>::get().checked_sub(expected_total_issuance)
+        let delta = 1000;
+        let total_issuance = TotalIssuance::<T>::get();
+
+        let diff = if total_issuance > expected_total_issuance {
+            total_issuance.checked_sub(expected_total_issuance)
         } else {
-            expected_total_issuance.checked_sub(TotalIssuance::<T>::get())
+            expected_total_issuance.checked_sub(total_issuance)
         }
         .expect("LHS > RHS");
+
         ensure!(
-            diff <= DELTA,
+            diff <= delta,
             "TotalIssuance diff greater than allowable delta",
         );
 
         Ok(())
     }
+
+    /// Checks the sum of all stakes matches the [`TotalStake`].
+    #[allow(dead_code)]
+    pub(crate) fn check_total_stake() -> Result<(), sp_runtime::TryRuntimeError> {
+        // Calculate the total staked amount
+        let total_staked = SubnetTAO::<T>::iter().fold(0u64, |acc, (netuid, stake)| {
+            let acc = acc.saturating_add(stake);
+
+            if netuid == Self::get_root_netuid() {
+                // root network doesn't have initial pool TAO
+                acc
+            } else {
+                acc.saturating_sub(POOL_INITIAL_TAO)
+            }
+        });
+
+        log::warn!(
+            "total_staked: {}, TotalStake: {}",
+            total_staked,
+            TotalStake::<T>::get()
+        );
+
+        // Verify that the calculated total stake matches the stored TotalStake
+        ensure!(
+            total_staked == TotalStake::<T>::get(),
+            "TotalStake does not match total staked",
+        );
+
+        Ok(())
+    }
 }

runtime/src/lib.rs

@@ -220,7 +220,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     //   `spec_version`, and `authoring_version` are the same between Wasm and native.
     // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use
     //   the compatible custom types.
-    spec_version: 224,
+    spec_version: 225,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,

scripts/try-runtime-upgrade.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+
+# Tries runtime upgrade (via try-runtime).
+#
+# Usage:
+#   try-runtime-upgrade.sh [-p <runtime-path>] [-u <live-chain-url>] [-s <snapshot-path>]
+#
+# Dependencies:
+#   - rust toolchain
+#   - try-runtime-cli
+
+set -eou pipefail
+
+runtime_wasm_path="./target/release/wbuild/node-subtensor-runtime/node_subtensor_runtime.compact.wasm"
+live_chain_url="wss://dev.chain.opentensor.ai:443"
+snapshot_path=""
+
+parse_args() {
+  u_provided=false
+
+  while getopts "r:u:s:" opt; do
+    case "${opt}" in
+    r) runtime_wasm_path="${OPTARG}" ;;
+    u)
+      live_chain_url="${OPTARG}"
+      u_provided=true
+      ;;
+    s) snapshot_path="${OPTARG}" ;;
+    *) echo "Usage: $(basename "$0") [-r <runtime-path>] [-u <live-chain-url>] [-s <snapshot-path>]" && exit 1 ;;
+    esac
+  done
+
+  # Prevent specifying URI if snapshot is specified
+  if [ -n "$snapshot_path" ] && [ "$u_provided" = true ]; then
+    echo "Error: Either live URI or snapshot path should be specified, but not both."
+    exit 1
+  fi
+}
+
+build_runtime() {
+  cargo build -p node-subtensor-runtime --release --features "metadata-hash,try-runtime"
+}
+
+do_try_runtime() {
+  if [ -n "$snapshot_path" ]; then
+    chain_state="snap --path $snapshot_path"
+  else
+    chain_state="live --uri $live_chain_url"
+  fi
+
+  eval "try-runtime --runtime $runtime_wasm_path on-runtime-upgrade \
+    --no-weight-warnings --disable-spec-version-check --disable-idempotency-checks --checks=all \
+    $chain_state"
+}
+
+parse_args "$@"
+build_runtime
+do_try_runtime