Skip to content

Commit

Permalink
make check idempot
Browse files Browse the repository at this point in the history
  • Loading branch information
Cameron Fairchild committed Sep 17, 2024
1 parent 67aad02 commit de41a34
Showing 1 changed file with 80 additions and 34 deletions.
114 changes: 80 additions & 34 deletions pallets/subtensor/src/migrations/migrate_fix_pending_emission.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -148,22 +148,20 @@ pub fn do_migrate_fix_pending_emission<T: Config>() -> Weight {
/// Required so we can read values in the v0 storage format during the migration.
#[cfg(feature = "try-runtime")]
mod v0 {
use frame_support::storage_alias;
use subtensor_macros::freeze_struct;

#[freeze_struct("7adbed79987ae83d")]
#[freeze_struct("2228babfc0580c62")]
#[derive(codec::Encode, codec::Decode, Clone, PartialEq, Debug)]
pub struct OldStorage {
pub total_issuance_before: u64,
pub total_stake_before: u64,
pub expected_taostats_new_hk_pending_emission: u64,
pub expected_datura_new_hk_pending_emission: u64,
pub old_migration_stake_taostats: u64,
pub old_null_stake_taostats: u64,
pub old_migration_stake_datura: u64,
pub old_null_stake_datura: u64,
}

/// V0 type for [`crate::Value`].
#[storage_alias]
pub type Value<T: crate::Config> = StorageValue<crate::Pallet<T>, OldStorage>;
}

impl<T: Config> Pallet<T> {
Expand All @@ -184,7 +182,9 @@ impl<T: Config> Pallet<T> {
let old = old_storage;
let null_stake_total = old
.old_null_stake_taostats
.saturating_add(old.old_null_stake_datura);
.saturating_add(old.old_null_stake_datura)
.saturating_add(old.old_migration_stake_taostats)
.saturating_add(old.old_migration_stake_datura);

match (
taostats_old_hk_account,
Expand All @@ -209,10 +209,19 @@ impl<T: Config> Pallet<T> {
0
);

// Check the total hotkey stake is the same
assert_eq!(
TotalHotkeyStake::<T>::get(taostats_old_hk_acct),
old.old_null_stake_taostats + old.old_migration_stake_taostats
);

let new_null_stake_taostats =
Self::get_stake_for_coldkey_and_hotkey(migration_ck_acct, taostats_old_hk_acct);

assert_eq!(new_null_stake_taostats, old.old_null_stake_taostats);
assert_eq!(
new_null_stake_taostats,
old.old_null_stake_taostats + old.old_migration_stake_taostats
);
}
_ => {
log::warn!("Failed to get account id from ss58 for taostats hotkeys");
Expand Down Expand Up @@ -251,10 +260,19 @@ impl<T: Config> Pallet<T> {
0
);

// Check the total hotkey stake is the same
assert_eq!(
TotalHotkeyStake::<T>::get(datura_old_hk_acct),
old.old_null_stake_datura + old.old_migration_stake_datura
);

let new_null_stake_datura =
Self::get_stake_for_coldkey_and_hotkey(migration_ck_acct, datura_old_hk_acct);

assert_eq!(new_null_stake_datura, old.old_null_stake_datura);
assert_eq!(
new_null_stake_datura,
old.old_null_stake_datura + old.old_migration_stake_datura
);
}
_ => {
log::warn!("Failed to get account id from ss58 for datura hotkeys");
Expand All @@ -278,10 +296,11 @@ impl<T: Config> Pallet<T> {

// Check the total issuance is the SAME following migration (no TAO issued)
let expected_total_issuance = old.total_issuance_before;
let expected_total_stake = old.total_stake_before;
assert_eq!(Self::get_total_issuance(), expected_total_issuance);

// Check total stake is the SAME following the migration (no new TAO staked)
assert_eq!(TotalStake::<T>::get(), expected_total_issuance);
assert_eq!(TotalStake::<T>::get(), expected_total_stake);
// Check the total stake maps are updated following the migration (removal of old null_account stake entries)
assert_eq!(TotalColdkeyStake::<T>::get(null_account), 0);

Expand All @@ -303,7 +322,11 @@ pub mod migration {

#[cfg(feature = "try-runtime")]
fn get_old_storage_values<T: Config>() -> Result<v0::OldStorage, sp_runtime::TryRuntimeError> {
log::info!("Getting old storage values for migration");

let null_account = &DefaultAccount::<T>::get();
let migration_coldkey = "5GeRjQYsobRWFnrbBmGe5ugme3rfnDVF69N45YtdBpUFsJG8";
let migration_account = &get_account_id_from_ss58::<T>(migration_coldkey);

let taostats_old_hotkey = "5Hddm3iBFD2GLT5ik7LZnT3XJUnRnN8PoeCFgGQgawUVKNm8";
let taostats_new_hotkey = "5GKH9FPPnWSUoeeTJp19wVtd84XqFW4pyK2ijV2GsFbhTrP1";
Expand All @@ -314,20 +337,27 @@ pub mod migration {
let total_issuance_before = crate::Pallet::<T>::get_total_issuance();
let mut expected_taostats_new_hk_pending_emission: u64 = 0;
let mut expected_datura_new_hk_pending_emission: u64 = 0;
let old_null_stake_taostats: u64 = match (taostats_old_hk_account, taostats_new_hk_account)
{
(Ok(taostats_old_hk_acct), Ok(taostats_new_hk_acct)) => {
let (old_null_stake_taostats, old_migration_stake_taostats) = match (
taostats_old_hk_account,
taostats_new_hk_account,
migration_account,
) {
(Ok(taostats_old_hk_acct), Ok(taostats_new_hk_acct), Ok(migration_acct)) => {
expected_taostats_new_hk_pending_emission =
expected_taostats_new_hk_pending_emission
.saturating_add(PendingdHotkeyEmission::<T>::get(taostats_old_hk_acct))
.saturating_add(PendingdHotkeyEmission::<T>::get(taostats_new_hk_acct));

Ok::<u64, sp_runtime::TryRuntimeError>(
Ok::<(u64, u64), sp_runtime::TryRuntimeError>((
crate::Pallet::<T>::get_stake_for_coldkey_and_hotkey(
null_account,
taostats_old_hk_acct,
),
)
crate::Pallet::<T>::get_stake_for_coldkey_and_hotkey(
migration_acct,
taostats_old_hk_acct,
),
))
}
_ => {
log::warn!("Failed to get account id from ss58 for taostats hotkeys");
Expand All @@ -341,33 +371,48 @@ pub mod migration {
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);

let old_null_stake_datura: u64 = match (datura_old_hk_account, datura_new_hk_account) {
(Ok(datura_old_hk_acct), Ok(datura_new_hk_acct)) => {
let (old_null_stake_datura, old_migration_stake_datura) = match (
datura_old_hk_account,
datura_new_hk_account,
migration_account,
) {
(Ok(datura_old_hk_acct), Ok(datura_new_hk_acct), Ok(migration_acct)) => {
expected_datura_new_hk_pending_emission = expected_datura_new_hk_pending_emission
.saturating_add(PendingdHotkeyEmission::<T>::get(datura_old_hk_acct))
.saturating_add(PendingdHotkeyEmission::<T>::get(datura_new_hk_acct));

Ok::<u64, sp_runtime::TryRuntimeError>(
Ok::<(u64, u64), sp_runtime::TryRuntimeError>((
crate::Pallet::<T>::get_stake_for_coldkey_and_hotkey(
null_account,
datura_old_hk_acct,
),
)
crate::Pallet::<T>::get_stake_for_coldkey_and_hotkey(
migration_acct,
datura_old_hk_acct,
),
))
}
_ => {
log::warn!("Failed to get account id from ss58 for datura hotkeys");
Err("Failed to get account id from ss58 for datura hotkeys".into())
}
}?;

let total_stake_before: u64 = crate::Pallet::<T>::get_total_stake();

let result = v0::OldStorage {
total_issuance_before,
total_stake_before,
expected_taostats_new_hk_pending_emission,
expected_datura_new_hk_pending_emission,
old_migration_stake_taostats,
old_null_stake_taostats,
old_migration_stake_datura,
old_null_stake_datura,
};

log::info!("Got old storage values for migration");

Ok(result)
}

Expand All @@ -378,10 +423,18 @@ pub mod migration {
use codec::Encode;

// Get the old storage values
let old_storage = get_old_storage_values::<T>()?;
match get_old_storage_values::<T>() {
Ok(old_storage) => {
log::info!("Successfully got old storage values for migration");
let encoded = old_storage.encode();

// Return it as an encoded `Vec<u8>`
Ok(old_storage.encode())
Ok(encoded)
}
Err(e) => {
log::error!("Failed to get old storage values for migration: {:?}", e);
Err("Failed to get old storage values for migration".into())
}
}
}

// Runs the migrate function for the fix_pending_emission migration
Expand Down Expand Up @@ -430,21 +483,14 @@ pub mod migration {
fn post_upgrade(state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
use codec::Decode;

let maybe_old_value =
Option::<v0::OldStorage>::decode(&mut &state[..]).map_err(|_| {
let old_storage: v0::OldStorage =
v0::OldStorage::decode(&mut &state[..]).map_err(|_| {
sp_runtime::TryRuntimeError::Other("Failed to decode old value from storage")
})?;

match maybe_old_value {
Some(old_value) => {
// Verify that all null stake invariants are satisfied after the migration
crate::Pallet::<T>::check_null_stake_invariants(old_value)?;
}
None => {
log::warn!("Failed to decode old value from storage");
return Err("Failed to decode old value from storage".into());
}
};
// Verify that all null stake invariants are satisfied after the migration
crate::Pallet::<T>::check_null_stake_invariants(old_storage)?;

Ok(())
}
}
Expand Down

0 comments on commit de41a34

Please sign in to comment.