Fix issues with Subtensor storage versions

pallets/subtensor/src/lib.rs

@@ -84,7 +84,7 @@ pub mod pallet {
 
     /// Tracks version for migrations. Should be monotonic with respect to the
     /// order of migrations. (i.e. always increasing)
-    const STORAGE_VERSION: StorageVersion = StorageVersion::new(6);
+    const STORAGE_VERSION: StorageVersion = StorageVersion::new(7);
 
     /// Minimum balance required to perform a coldkey swap
     pub const MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP: u64 = 100_000_000; // 0.1 TAO in RAO
@@ -1140,6 +1140,9 @@ pub mod pallet {
         DefaultBonds<T>,
     >;
 
+    #[pallet::storage] // --- Storage for migration run status
+    pub type HasMigrationRun<T: Config> = StorageMap<_, Identity, Vec<u8>, bool, ValueQuery>;
+
     /// ==================
     /// ==== Genesis =====
     /// ==================

pallets/subtensor/src/migration.rs

@@ -1,4 +1,5 @@
 use super::*;
+use alloc::string::String;
 use frame_support::traits::DefensiveResult;
 use frame_support::{
     pallet_prelude::{Identity, OptionQuery},
@@ -56,30 +57,53 @@ pub fn do_migrate_fix_total_coldkey_stake<T: Config>() -> Weight {
     }
     weight
 }
-// Public migrate function to be called by Lib.rs on upgrade.
+
+/// Migrates and fixes the total coldkey stake.
+///
+/// This function checks if the migration has already run, and if not, it performs the migration
+/// to fix the total coldkey stake. It also marks the migration as completed after running.
+///
+/// # Returns
+/// The weight of the migration process.
 pub fn migrate_fix_total_coldkey_stake<T: Config>() -> Weight {
-    let current_storage_version: u16 = 7;
-    let next_storage_version: u16 = 8;
+    let migration_name = b"fix_total_coldkey_stake_v7".to_vec();
 
     // Initialize the weight with one read operation.
     let mut weight = T::DbWeight::get().reads(1);
 
-    // Grab the current on-chain storage version.
-    // Cant fail on retrieval.
-    let onchain_version = Pallet::<T>::on_chain_storage_version();
-
-    // Only run this migration on storage version 6.
-    if onchain_version == current_storage_version {
-        weight = weight.saturating_add(do_migrate_fix_total_coldkey_stake::<T>());
-        // Cant fail on insert.
-        StorageVersion::new(next_storage_version).put::<Pallet<T>>();
-        weight.saturating_accrue(T::DbWeight::get().writes(1));
+    // Check if the migration has already run
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            "Migration '{:?}' has already run. Skipping.",
+            migration_name
+        );
+        return Weight::zero();
     }
 
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    // Run the migration
+    weight = weight.saturating_add(do_migrate_fix_total_coldkey_stake::<T>());
+
+    // Mark the migration as completed
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    // Set the storage version to 7
+    StorageVersion::new(7).put::<Pallet<T>>();
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{:?}' completed. Storage version set to 7.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
     // Return the migration weight.
     weight
 }
-
 /// Performs migration to update the total issuance based on the sum of stakes and total balances.
 /// This migration is applicable only if the current storage version is 5, after which it updates the storage version to 6.
 ///

pallets/subtensor/tests/migration.rs

@@ -1,7 +1,7 @@
 #![allow(clippy::unwrap_used)]
 
 mod mock;
-use frame_support::assert_ok;
+use frame_support::{assert_ok, weights::Weight};
 use frame_system::Config;
 use mock::*;
 use pallet_subtensor::*;
@@ -278,17 +278,25 @@ fn test_migration_delete_subnet_21() {
     })
 }
 
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake --exact --nocapture
 // SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake --exact --nocapture
 #[test]
 fn test_migrate_fix_total_coldkey_stake() {
     new_test_ext(1).execute_with(|| {
+        let migration_name = "fix_total_coldkey_stake_v7";
         let coldkey = U256::from(0);
         TotalColdkeyStake::<Test>::insert(coldkey, 0);
         StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
         Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
         Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
         Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
-        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+
+        let weight = run_migration_and_check(migration_name);
+        assert!(weight != Weight::zero());
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+
+        let second_weight = run_migration_and_check(migration_name);
+        assert_eq!(second_weight, Weight::zero());
         assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
     })
 }
@@ -297,13 +305,16 @@ fn test_migrate_fix_total_coldkey_stake() {
 #[test]
 fn test_migrate_fix_total_coldkey_stake_value_already_in_total() {
     new_test_ext(1).execute_with(|| {
+        let migration_name = "fix_total_coldkey_stake_v7";
         let coldkey = U256::from(0);
         TotalColdkeyStake::<Test>::insert(coldkey, 100000000);
         StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
         Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
         Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
         Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
-        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+
+        let weight = run_migration_and_check(migration_name);
+        assert!(weight != Weight::zero());
         assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
     })
 }
@@ -312,12 +323,15 @@ fn test_migrate_fix_total_coldkey_stake_value_already_in_total() {
 #[test]
 fn test_migrate_fix_total_coldkey_stake_no_entry() {
     new_test_ext(1).execute_with(|| {
+        let migration_name = "fix_total_coldkey_stake_v7";
         let coldkey = U256::from(0);
         StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
         Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
         Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
         Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
-        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+
+        let weight = run_migration_and_check(migration_name);
+        assert!(weight != Weight::zero());
         assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
     })
 }
@@ -326,10 +340,13 @@ fn test_migrate_fix_total_coldkey_stake_no_entry() {
 #[test]
 fn test_migrate_fix_total_coldkey_stake_no_entry_in_hotkeys() {
     new_test_ext(1).execute_with(|| {
+        let migration_name = "fix_total_coldkey_stake_v7";
         let coldkey = U256::from(0);
         TotalColdkeyStake::<Test>::insert(coldkey, 100000000);
         StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
-        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+
+        let weight = run_migration_and_check(migration_name);
+        assert!(weight != Weight::zero());
         assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 0);
     })
 }
@@ -338,12 +355,53 @@ fn test_migrate_fix_total_coldkey_stake_no_entry_in_hotkeys() {
 #[test]
 fn test_migrate_fix_total_coldkey_stake_one_hotkey_stake_missing() {
     new_test_ext(1).execute_with(|| {
+        let migration_name = "fix_total_coldkey_stake_v7";
         let coldkey = U256::from(0);
         TotalColdkeyStake::<Test>::insert(coldkey, 100000000);
         StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
         Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
         Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
-        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+
+        let weight = run_migration_and_check(migration_name);
+        assert!(weight != Weight::zero());
         assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 20000);
     })
 }
+
+// New test to check if migration runs only once
+#[test]
+fn test_migrate_fix_total_coldkey_stake_runs_once() {
+    new_test_ext(1).execute_with(|| {
+        let migration_name = "fix_total_coldkey_stake_v7";
+        let coldkey = U256::from(0);
+        TotalColdkeyStake::<Test>::insert(coldkey, 0);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
+
+        // First run
+        let first_weight = run_migration_and_check(migration_name);
+        assert!(first_weight != Weight::zero());
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+
+        // Second run
+        let second_weight = run_migration_and_check(migration_name);
+        assert_eq!(second_weight, Weight::zero());
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+    })
+}
+
+fn run_migration_and_check(migration_name: &'static str) -> frame_support::weights::Weight {
+    // Execute the migration and store its weight
+    let weight: frame_support::weights::Weight =
+        pallet_subtensor::migration::migrate_fix_total_coldkey_stake::<Test>();
+
+    // Check if the migration has been marked as completed
+    assert!(HasMigrationRun::<Test>::get(
+        migration_name.as_bytes().to_vec()
+    ));
+
+    // Return the weight of the executed migration
+    weight
+}