diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 8d713d76a..f7824e2a3 100644 --- a/pallets/subtensor/src/lib.rs +++ b/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, >; + #[pallet::storage] // --- Storage for migration run status + pub type HasMigrationRun = StorageMap<_, Identity, Vec, bool, ValueQuery>; + /// ================== /// ==== Genesis ===== /// ================== diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs index 72a86ea4e..20e91237f 100644 --- a/pallets/subtensor/src/migration.rs +++ b/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() -> 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() -> 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::::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::()); - // Cant fail on insert. - StorageVersion::new(next_storage_version).put::>(); - weight.saturating_accrue(T::DbWeight::get().writes(1)); + // Check if the migration has already run + if HasMigrationRun::::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::()); + + // Mark the migration as completed + HasMigrationRun::::insert(&migration_name, true); + weight = weight.saturating_add(T::DbWeight::get().writes(1)); + + // Set the storage version to 7 + StorageVersion::new(7).put::>(); + 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. /// diff --git a/pallets/subtensor/tests/migration.rs b/pallets/subtensor/tests/migration.rs index 360568235..8323cab8a 100644 --- a/pallets/subtensor/tests/migration.rs +++ b/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::::insert(coldkey, 0); StakingHotkeys::::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]); Stake::::insert(U256::from(1), U256::from(0), 10000); Stake::::insert(U256::from(2), U256::from(0), 10000); Stake::::insert(U256::from(3), U256::from(0), 10000); - pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::(); + + let weight = run_migration_and_check(migration_name); + assert!(weight != Weight::zero()); + assert_eq!(TotalColdkeyStake::::get(coldkey), 30000); + + let second_weight = run_migration_and_check(migration_name); + assert_eq!(second_weight, Weight::zero()); assert_eq!(TotalColdkeyStake::::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::::insert(coldkey, 100000000); StakingHotkeys::::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]); Stake::::insert(U256::from(1), U256::from(0), 10000); Stake::::insert(U256::from(2), U256::from(0), 10000); Stake::::insert(U256::from(3), U256::from(0), 10000); - pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::(); + + let weight = run_migration_and_check(migration_name); + assert!(weight != Weight::zero()); assert_eq!(TotalColdkeyStake::::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::::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]); Stake::::insert(U256::from(1), U256::from(0), 10000); Stake::::insert(U256::from(2), U256::from(0), 10000); Stake::::insert(U256::from(3), U256::from(0), 10000); - pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::(); + + let weight = run_migration_and_check(migration_name); + assert!(weight != Weight::zero()); assert_eq!(TotalColdkeyStake::::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::::insert(coldkey, 100000000); StakingHotkeys::::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]); - pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::(); + + let weight = run_migration_and_check(migration_name); + assert!(weight != Weight::zero()); assert_eq!(TotalColdkeyStake::::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::::insert(coldkey, 100000000); StakingHotkeys::::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]); Stake::::insert(U256::from(1), U256::from(0), 10000); Stake::::insert(U256::from(2), U256::from(0), 10000); - pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::(); + + let weight = run_migration_and_check(migration_name); + assert!(weight != Weight::zero()); assert_eq!(TotalColdkeyStake::::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::::insert(coldkey, 0); + StakingHotkeys::::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]); + Stake::::insert(U256::from(1), U256::from(0), 10000); + Stake::::insert(U256::from(2), U256::from(0), 10000); + Stake::::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::::get(coldkey), 30000); + + // Second run + let second_weight = run_migration_and_check(migration_name); + assert_eq!(second_weight, Weight::zero()); + assert_eq!(TotalColdkeyStake::::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::(); + + // Check if the migration has been marked as completed + assert!(HasMigrationRun::::get( + migration_name.as_bytes().to_vec() + )); + + // Return the weight of the executed migration + weight +}