Fix total stake maps via migration

pallets/subtensor/src/lib.rs

@@ -51,7 +51,7 @@ pub mod subnet_info;
 
 // apparently this is stabilized since rust 1.36
 extern crate alloc;
-mod migration;
+pub mod migration;
 
 #[frame_support::pallet]
 pub mod pallet {
@@ -983,10 +983,16 @@ pub mod pallet {
         }
 
         fn on_runtime_upgrade() -> frame_support::weights::Weight {
-            // --- Migrate to v2
+            // --- Migrate storage
             use crate::migration;
 
-            migration::migrate_to_v2_separate_emission::<T>()
+			let mut weight = frame_support::weights::Weight::from_ref_time(0);
+
+			weight = weight
+				.saturating_add( migration::migrate_to_v1_separate_emission::<T>() )
+				.saturating_add( migration::migrate_to_v2_fixed_total_stake::<T>() );
+
+			return weight;
         }
     }
 

pallets/subtensor/src/migration.rs

@@ -9,8 +9,6 @@ use frame_support::{
     inherent::Vec
 };
 
-// TODO (camfairchild): TEST MIGRATION
-
 const LOG_TARGET: &str = "loadedemissionmigration";
 
 pub mod deprecated_loaded_emission_format {
@@ -22,16 +20,16 @@ pub mod deprecated_loaded_emission_format {
     pub(super) type LoadedEmission<T:Config> = StorageMap< Pallet<T>, Identity, u16, Vec<(AccountIdOf<T>, u64)>, OptionQuery >;
 }
 
-pub fn migrate_to_v2_separate_emission<T: Config>() -> Weight {
+pub fn migrate_to_v1_separate_emission<T: Config>() -> Weight {
     use deprecated_loaded_emission_format as old;
      // Check storage version
     let mut weight = T::DbWeight::get().reads_writes(1, 0);
 
     // Grab current version
     let onchain_version =  Pallet::<T>::on_chain_storage_version();
 
-    // Only runs if we haven't already updated version to 2.
-    if onchain_version < 2 { 
+    // Only runs if we haven't already updated version to 1.
+    if onchain_version < 1 {
         info!(target: LOG_TARGET, ">>> Updating the LoadedEmission to a new format {:?}", onchain_version);
 
         // We transform the storage values from the old into the new format.
@@ -51,13 +49,13 @@ pub fn migrate_to_v2_separate_emission<T: Config>() -> Weight {
         // Translate the old storage values into the new format.
         LoadedEmission::<T>::translate::<Vec<(AccountIdOf<T>, u64)>, _>(
             |netuid: u16, netuid_emissions: Vec<(AccountIdOf<T>, u64)>| -> Option<Vec<(AccountIdOf<T>, u64, u64)>> {
-                info!(target: LOG_TARGET, "     Do migration of netuid: {:?}...", netuid); 
-                
-                // We will assume all loaded emission is validator emissions, 
+                info!(target: LOG_TARGET, "     Do migration of netuid: {:?}...", netuid);
+
+                // We will assume all loaded emission is validator emissions,
                 //      so this will get distributed over delegatees (nominators), if there are any
-                //      This will NOT effect any servers that are not (also) a delegate validator. 
+                //      This will NOT effect any servers that are not (also) a delegate validator.
                 // server_emission will be 0 for any alread loaded emission.
-                
+
                 let mut new_netuid_emissions = Vec::new();
                 for (server, validator_emission) in netuid_emissions {
                     new_netuid_emissions.push((server, 0 as u64, validator_emission));
@@ -71,13 +69,81 @@ pub fn migrate_to_v2_separate_emission<T: Config>() -> Weight {
         );
 
         // Update storage version.
-        StorageVersion::new(1).put::<Pallet::<T>>(); // Update to version 2 so we don't run this again.
+        StorageVersion::new(1).put::<Pallet::<T>>(); // Update to version 1 so we don't run this again.
+        // One write to storage version
+        weight.saturating_accrue(T::DbWeight::get().writes(1));
+
+        weight
+    } else {
+        info!(target: LOG_TARGET, "Migration to v1 already done!");
+        Weight::zero()
+    }
+}
+
+
+const LOG_TARGET_1: &str = "fixtotalstakestorage";
+
+pub fn migrate_to_v2_fixed_total_stake<T: Config>() -> Weight {
+	let new_storage_version = 2;
+
+     // Check storage version
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Grab current version
+    let onchain_version =  Pallet::<T>::on_chain_storage_version();
+
+    // Only runs if we haven't already updated version past above new_storage_version.
+    if onchain_version < new_storage_version {
+        info!(target: LOG_TARGET_1, ">>> Fixing the TotalStake and TotalColdkeyStake storage {:?}", onchain_version);
+
+		// Stake and TotalHotkeyStake are known to be accurate
+		// TotalColdkeyStake is known to be inaccurate
+		// TotalStake is known to be inaccurate
+
+		TotalStake::<T>::put(0); // Set to 0
+		weight.saturating_accrue(T::DbWeight::get().writes(1));
+
+		// We iterate over TotalColdkeyStake keys and set them to 0
+		let total_coldkey_stake_keys = TotalColdkeyStake::<T>::iter_keys().collect::<Vec<_>>();
+		for coldkey in total_coldkey_stake_keys {
+			weight.saturating_accrue(T::DbWeight::get().reads(1));
+			TotalColdkeyStake::<T>::insert(coldkey, 0); // Set to 0
+			weight.saturating_accrue(T::DbWeight::get().writes(1));
+		}
+
+		// Now we iterate over the entire stake map, and sum each coldkey stake
+		//   We also track TotalStake
+		for (_, coldkey, stake) in Stake::<T>::iter() {
+			weight.saturating_accrue(T::DbWeight::get().reads(1));
+			// Get the current coldkey stake
+			let mut total_coldkey_stake = TotalColdkeyStake::<T>::get(coldkey.clone());
+			weight.saturating_accrue(T::DbWeight::get().reads(1));
+			// Add the stake to the coldkey stake
+			total_coldkey_stake = total_coldkey_stake.saturating_add(stake);
+			// Update the coldkey stake
+			TotalColdkeyStake::<T>::insert(coldkey, total_coldkey_stake);
+			weight.saturating_accrue(T::DbWeight::get().writes(1));
+
+			// Get the current total stake
+			let mut total_stake = TotalStake::<T>::get();
+			weight.saturating_accrue(T::DbWeight::get().reads(1));
+			// Add the stake to the total stake
+			total_stake = total_stake.saturating_add(stake);
+			// Update the total stake
+			TotalStake::<T>::put(total_stake);
+			weight.saturating_accrue(T::DbWeight::get().writes(1));
+		}
+
+		// Now both TotalStake and TotalColdkeyStake are accurate
+
+        // Update storage version.
+        StorageVersion::new(new_storage_version).put::<Pallet::<T>>(); // Update to version so we don't run this again.
         // One write to storage version
         weight.saturating_accrue(T::DbWeight::get().writes(1));
-        
+
         weight
     } else {
-        info!(target: LOG_TARGET, "Migration to v2 already done!");
+        info!(target: LOG_TARGET_1, "Migration to v2 already done!");
         Weight::zero()
     }
-}
+}

pallets/subtensor/tests/migration.rs

@@ -0,0 +1,87 @@
+mod mock;
+use mock::*;
+use sp_core::U256;
+
+#[test]
+fn test_migration_fix_total_stake_maps() {
+    new_test_ext().execute_with(|| {
+		let ck1 = U256::from(1);
+		let ck2 = U256::from(2);
+		let ck3 = U256::from(3);
+
+		let hk1 = U256::from(1 + 100);
+		let hk2 = U256::from(2 + 100);
+
+		let mut total_stake_amount = 0;
+
+		// Give each coldkey some stake in the maps
+		SubtensorModule::increase_stake_on_coldkey_hotkey_account(
+			&ck1,
+			&hk1,
+			100
+		);
+		total_stake_amount += 100;
+
+		SubtensorModule::increase_stake_on_coldkey_hotkey_account(
+			&ck2,
+			&hk1,
+			10_101
+		);
+		total_stake_amount += 10_101;
+
+		SubtensorModule::increase_stake_on_coldkey_hotkey_account(
+			&ck3,
+			&hk2,
+			100_000_000
+		);
+		total_stake_amount += 100_000_000;
+
+		SubtensorModule::increase_stake_on_coldkey_hotkey_account(
+			&ck1,
+			&hk2,
+			1_123_000_000
+		);
+		total_stake_amount += 1_123_000_000;
+
+		// Check that the total stake is correct
+		assert_eq!(SubtensorModule::get_total_stake(), total_stake_amount);
+
+		// Check that the total coldkey stake is correct
+		assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&ck1), 100 + 1_123_000_000);
+		assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&ck2), 10_101);
+		assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&ck3), 100_000_000);
+
+		// Check that the total hotkey stake is correct
+		assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hk1), 100 + 10_101);
+		assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hk2), 100_000_000 + 1_123_000_000);
+
+		// Mess up the total coldkey stake
+		pallet_subtensor::TotalColdkeyStake::<Test>::insert(ck1, 0);
+		// Verify that the total coldkey stake is now 0 for ck1
+		assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&ck1), 0);
+
+		// Mess up the total stake
+		pallet_subtensor::TotalStake::<Test>::put(123_456_789);
+		// Verify that the total stake is now wrong
+		assert_ne!(SubtensorModule::get_total_stake(), total_stake_amount);
+
+		// Run the migration to fix the total stake maps
+		pallet_subtensor::migration::migrate_to_v2_fixed_total_stake::<Test>();
+
+		// Verify that the total stake is now correct
+		assert_eq!(SubtensorModule::get_total_stake(), total_stake_amount);
+		// Verify that the total coldkey stake is now correct for each coldkey
+		assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&ck1), 100 + 1_123_000_000);
+		assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&ck2), 10_101);
+		assert_eq!(SubtensorModule::get_total_stake_for_coldkey(&ck3), 100_000_000);
+
+		// Verify that the total hotkey stake is STILL correct for each hotkey
+		assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hk1), 100 + 10_101);
+		assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hk2), 100_000_000 + 1_123_000_000);
+
+		// Verify that the Stake map has no extra entries
+		assert_eq!(pallet_subtensor::Stake::<Test>::iter().count(), 4); // 4 entries total
+		assert_eq!(pallet_subtensor::Stake::<Test>::iter_key_prefix(hk1).count(), 2); // 2 stake entries for hk1
+		assert_eq!(pallet_subtensor::Stake::<Test>::iter_key_prefix(hk2).count(), 2); // 2 stake entries for hk2
+    })
+}