Skip to content

Commit

Permalink
add new migration method and more testing
Browse files Browse the repository at this point in the history
  • Loading branch information
@camfairchild
camfairchild committed Sep 13, 2024
1 parent 773670c commit 9f4b90c
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 37 deletions.
65 changes: 44 additions & 21 deletions pallets/subtensor/src/migrations/migrate_fix_pending_emission.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ fn get_account_id_from_ss58<T: Config>(ss58_str: &str) -> Result<T::AccountId, c
fn migrate_pending_emissions_including_null_stake<T: Config>(
old_hotkey: &T::AccountId,
new_hotkey: &T::AccountId,
migration_account: &T::AccountId,
) -> Weight {
let mut weight = T::DbWeight::get().reads(0);
let null_account = &DefaultAccount::<T>::get();
Expand All @@ -30,10 +31,10 @@ fn migrate_pending_emissions_including_null_stake<T: Config>(
weight.saturating_accrue(T::DbWeight::get().reads(1));

// Get the stake for the 0x000 key
let null_stake = Stake::<T>::get(&old_hotkey, null_account);
let null_stake = Stake::<T>::get(old_hotkey, null_account);
weight.saturating_accrue(T::DbWeight::get().reads(1));
// Remove
Stake::<T>::remove(&old_hotkey, null_account);
Stake::<T>::remove(old_hotkey, null_account);
weight.saturating_accrue(T::DbWeight::get().writes(1));

let new_total_coldkey_stake =
Expand All @@ -45,29 +46,35 @@ fn migrate_pending_emissions_including_null_stake<T: Config>(
}
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));

let new_total_hotkey_stake = TotalHotkeyStake::<T>::get(old_hotkey).saturating_sub(null_stake);
if new_total_hotkey_stake == 0 {
TotalHotkeyStake::<T>::remove(old_hotkey);
} else {
TotalHotkeyStake::<T>::insert(old_hotkey, new_total_hotkey_stake);
}
let new_staking_hotkeys = StakingHotkeys::<T>::get(null_account);
let new_staking_hotkeys = new_staking_hotkeys
.into_iter()
.filter(|hk| hk != old_hotkey)
.collect::<Vec<_>>();
StakingHotkeys::<T>::insert(null_account, new_staking_hotkeys);
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));

// Remove the stake from the total stake and total issuance (since it is re-emitted)
TotalStake::<T>::put(TotalStake::<T>::get().saturating_sub(null_stake));
TotalIssuance::<T>::put(TotalIssuance::<T>::get().saturating_sub(null_stake));
weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
// Insert the stake from the null account to the MIGRATION account under the OLD hotkey
Stake::<T>::insert(old_hotkey, migration_account, null_stake);
TotalColdkeyStake::<T>::insert(
migration_account,
TotalColdkeyStake::<T>::get(migration_account).saturating_add(null_stake),
);
let mut new_staking_hotkeys = StakingHotkeys::<T>::get(migration_account);
if !new_staking_hotkeys.contains(old_hotkey) {
new_staking_hotkeys.push(old_hotkey.clone());
}
StakingHotkeys::<T>::insert(migration_account, new_staking_hotkeys);
weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 3));

// Get the pending emissions for the NEW hotkey
let pending_emissions_new: u64 = PendingdHotkeyEmission::<T>::get(new_hotkey);
weight.saturating_accrue(T::DbWeight::get().reads(1));

// Add stake to the pending emissions for the new hotkey and the old hotkey
// Add the pending emissions for the new hotkey and the old hotkey
PendingdHotkeyEmission::<T>::insert(
new_hotkey,
pending_emissions_new
.saturating_add(pending_emissions_old)
.saturating_add(null_stake),
pending_emissions_new.saturating_add(pending_emissions_old),
);
weight.saturating_accrue(T::DbWeight::get().writes(1));

Expand All @@ -80,19 +87,28 @@ pub fn do_migrate_fix_pending_emission<T: Config>() -> Weight {

let taostats_old_hotkey = "5Hddm3iBFD2GLT5ik7LZnT3XJUnRnN8PoeCFgGQgawUVKNm8";
let taostats_new_hotkey = "5GKH9FPPnWSUoeeTJp19wVtd84XqFW4pyK2ijV2GsFbhTrP1";
let migration_coldkey = "5D65DoFbapkYzJK17VRQo3HFs7FmMeicbaQern28UNDPypCT";

let taostats_old_hk_account = get_account_id_from_ss58::<T>(taostats_old_hotkey);
let taostats_new_hk_account = get_account_id_from_ss58::<T>(taostats_new_hotkey);

match (taostats_old_hk_account, taostats_new_hk_account) {
(Ok(taostats_old_hk_acct), Ok(taostats_new_hk_acct)) => {
let migration_ck_account = get_account_id_from_ss58::<T>(migration_coldkey);

match (
taostats_old_hk_account,
taostats_new_hk_account,
migration_ck_account.clone(),
) {
(Ok(taostats_old_hk_acct), Ok(taostats_new_hk_acct), Ok(migration_ck_account)) => {
weight.saturating_accrue(migrate_pending_emissions_including_null_stake::<T>(
&taostats_old_hk_acct,
&taostats_new_hk_acct,
&migration_ck_account,
));
log::info!("Migrated pending emissions from taostats old hotkey to new hotkey");
}
_ => {
log::warn!("Failed to get account id from ss58 for taostats hotkeys");
return weight;
}
}

Expand All @@ -102,15 +118,22 @@ pub fn do_migrate_fix_pending_emission<T: Config>() -> Weight {
let datura_old_hk_account = get_account_id_from_ss58::<T>(datura_old_hotkey);
let datura_new_hk_account = get_account_id_from_ss58::<T>(datura_new_hotkey);

match (datura_old_hk_account, datura_new_hk_account) {
(Ok(datura_old_hk_acct), Ok(datura_new_hk_acct)) => {
match (
datura_old_hk_account,
datura_new_hk_account,
migration_ck_account,
) {
(Ok(datura_old_hk_acct), Ok(datura_new_hk_acct), Ok(migration_ck_account)) => {
weight.saturating_accrue(migrate_pending_emissions_including_null_stake::<T>(
&datura_old_hk_acct,
&datura_new_hk_acct,
&migration_ck_account,
));
log::info!("Migrated pending emissions from datura old hotkey to new hotkey");
}
_ => {
log::warn!("Failed to get account id from ss58 for datura hotkeys");
return weight;
}
}

Expand Down
68 changes: 52 additions & 16 deletions pallets/subtensor/tests/migration.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -463,20 +463,22 @@ fn test_migrate_fix_pending_emissions() {
new_test_ext(1).execute_with(|| {
let migration_name = "fix_pending_emission";

let null_account = U256::from(0); // The null account
let rand_coldkeys = [U256::from(1), U256::from(2), U256::from(3), U256::from(4)];
let null_account = &U256::from(0); // The null account

let taostats_old_hotkey = "5Hddm3iBFD2GLT5ik7LZnT3XJUnRnN8PoeCFgGQgawUVKNm8";
let taostats_new_hotkey = "5GKH9FPPnWSUoeeTJp19wVtd84XqFW4pyK2ijV2GsFbhTrP1";

let taostats_old_hk_account: AccountId = get_account_id_from_ss58(taostats_old_hotkey);
let taostats_new_hk_account: AccountId = get_account_id_from_ss58(taostats_new_hotkey);
let taostats_old_hk_account: &AccountId = &get_account_id_from_ss58(taostats_old_hotkey);
let taostats_new_hk_account: &AccountId = &get_account_id_from_ss58(taostats_new_hotkey);

let datura_old_hotkey = "5FKstHjZkh4v3qAMSBa1oJcHCLjxYZ8SNTSz1opTv4hR7gVB";
let datura_new_hotkey = "5GP7c3fFazW9GXK8Up3qgu2DJBk8inu4aK9TZy3RuoSWVCMi";

let datura_old_hk_account: AccountId = get_account_id_from_ss58(datura_old_hotkey);
let datura_new_hk_account: AccountId = get_account_id_from_ss58(datura_new_hotkey);
let datura_old_hk_account: &AccountId = &get_account_id_from_ss58(datura_old_hotkey);
let datura_new_hk_account: &AccountId = &get_account_id_from_ss58(datura_new_hotkey);

let migration_coldkey = "5D65DoFbapkYzJK17VRQo3HFs7FmMeicbaQern28UNDPypCT";
let migration_account: &AccountId = &get_account_id_from_ss58(migration_coldkey);

// "Issue" the TAO we're going to insert to stake
let null_stake_datura = 123_456_789;
Expand All @@ -491,18 +493,17 @@ fn test_migrate_fix_pending_emissions() {
// Setup the old Datura hotkey with a pending emission
PendingdHotkeyEmission::<Test>::insert(datura_old_hk_account, 10_000);
// Setup the NEW Datura hotkey with a pending emission
PendingdHotkeyEmission::<Test>::insert(datura_new_hk_account, null_stake_datura);
Stake::<Test>::insert(datura_old_hk_account, null_account, 123_456_789);
let expected_datura_new_hk_pending_emission: u64 = 123_456_789 + 10_000 + null_stake_datura;
PendingdHotkeyEmission::<Test>::insert(datura_new_hk_account, 123_456_789);
Stake::<Test>::insert(datura_old_hk_account, null_account, null_stake_datura);
let expected_datura_new_hk_pending_emission: u64 = 123_456_789 + 10_000;

// Setup the old TaoStats hotkey with a pending emission
PendingdHotkeyEmission::<Test>::insert(taostats_old_hk_account, 987_654);
// Setup the new TaoStats hotkey with a pending emission
PendingdHotkeyEmission::<Test>::insert(taostats_new_hk_account, 100_000);
// Setup the old TaoStats hotkey with a null-key stake entry
Stake::<Test>::insert(taostats_old_hk_account, null_account, null_stake_tao_stats);
let expected_taostats_new_hk_pending_emission: u64 =
987_654 + 100_000 + null_stake_tao_stats;
let expected_taostats_new_hk_pending_emission: u64 = 987_654 + 100_000;

let total_issuance_before = SubtensorModule::get_total_issuance();

Expand Down Expand Up @@ -537,17 +538,52 @@ fn test_migrate_fix_pending_emissions() {
assert_eq!(Stake::<Test>::get(datura_old_hk_account, null_account), 0);
assert_eq!(Stake::<Test>::get(taostats_old_hk_account, null_account), 0);

// Check the total issuance is decreased by the null stake removed
let expected_total_issuance = total_issuance_before - null_stake_total;
// Check the total issuance is the SAME following migration (no TAO issued)
let expected_total_issuance = total_issuance_before;
assert_eq!(
SubtensorModule::get_total_issuance(),
expected_total_issuance
);

// Check total stake is decreased by the null stake removed
// Check total stake is the SAME following the migration (no new TAO staked)
assert_eq!(TotalStake::<Test>::get(), expected_total_issuance);
// Check the total stake maps are updated following the migration (removal of old null_account stake entries)
assert_eq!(TotalColdkeyStake::<Test>::get(null_account), 0);
assert_eq!(TotalHotkeyStake::<Test>::get(datura_old_hk_account), 0);
assert_eq!(TotalHotkeyStake::<Test>::get(taostats_old_hk_account), 0);
assert_eq!(
SubtensorModule::get_stake_for_coldkey_and_hotkey(null_account, datura_old_hk_account),
0
);
assert_eq!(
SubtensorModule::get_stake_for_coldkey_and_hotkey(
null_account,
taostats_old_hk_account
),
0
);

// Check staking hotkeys is updated
assert_eq!(StakingHotkeys::<Test>::get(null_account), vec![]);

// Check the migration key has stake with both *old* hotkeys
assert_eq!(
SubtensorModule::get_stake_for_coldkey_and_hotkey(
migration_account,
datura_old_hk_account
),
null_stake_datura
);
assert_eq!(
SubtensorModule::get_stake_for_coldkey_and_hotkey(
migration_account,
taostats_old_hk_account
),
null_stake_tao_stats
);
assert_eq!(
TotalColdkeyStake::<Test>::get(migration_account),
null_stake_total
);
assert!(StakingHotkeys::<Test>::get(migration_account).contains(&datura_old_hk_account));
assert!(StakingHotkeys::<Test>::get(migration_account).contains(&taostats_old_hk_account));
})
}

0 comments on commit 9f4b90c

Please sign in to comment.