Skip to content

Commit

Permalink
judgement proxy migration
Browse files Browse the repository at this point in the history
  • Loading branch information
liamaharon committed Jun 2, 2024
1 parent 38482ff commit 89a2f5b
Showing 1 changed file with 203 additions and 0 deletions.
203 changes: 203 additions & 0 deletions relay/kusama/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1927,6 +1927,208 @@ pub mod migrations {
}
}

/// Migration to remove deprecated judgement proxies.
mod clear_judgement_proxies {
use super::*;

use frame_support::{
pallet_prelude::ValueQuery,
storage_alias,
traits::{Currency, ReservableCurrency},
Twox64Concat,
};
use frame_system::pallet_prelude::BlockNumberFor;
use pallet_proxy::ProxyDefinition;
use sp_runtime::{BoundedVec, Saturating};

/// ProxyType including the deprecated `IdentityJudgement`.
#[derive(
Copy,
Clone,
Eq,
PartialEq,
Ord,
PartialOrd,
Encode,
Decode,
RuntimeDebug,
MaxEncodedLen,
TypeInfo,
)]
pub enum PrevProxyType {
Any,
NonTransfer,
Governance,
Staking,
IdentityJudgement,
CancelProxy,
Auction,
Society,
NominationPools,
}

type BalanceOf<T> = <<T as pallet_proxy::Config>::Currency as Currency<
<T as frame_system::Config>::AccountId,
>>::Balance;

type PrevProxiesValue<T> = (
BoundedVec<ProxyDefinition<AccountId, PrevProxyType, BlockNumberFor<T>>, MaxProxies>,
BalanceOf<T>,
);

/// Proxies including the deprecated `IdentityJudgement` type.
#[storage_alias]
pub type Proxies<T: pallet_proxy::Config> = StorageMap<
pallet_proxy::Pallet<T>,
Twox64Concat,
AccountId,
PrevProxiesValue<T>,
ValueQuery,
>;

pub struct Migration;
impl OnRuntimeUpgrade for Migration {
/// Compute the expected post-upgrade state for Proxies stroage, and the reserved value
/// for all accounts with a proxy.
#[cfg(feature = "try-runtime")]
fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
let mut expected_proxies: BTreeMap<AccountId, PrevProxiesValue<Runtime>> =
BTreeMap::new();
let mut expected_reserved_amounts: BTreeMap<AccountId, Balance> = BTreeMap::new();

for (who, (mut proxies, old_deposit)) in
Proxies::<Runtime>::iter().collect::<Vec<_>>()
{
let proxies_len_before = proxies.len() as u64;
proxies.retain(|proxy| proxy.proxy_type != PrevProxyType::IdentityJudgement);
let proxies_len_after = proxies.len() as u64;

let new_deposit =
pallet_proxy::Pallet::<Runtime>::deposit(proxies.len() as u32);

let current_reserved =
<Balances as ReservableCurrency<AccountId>>::reserved_balance(&who);

// Update the deposit only if proxies were removed and the deposit decreased.
if new_deposit < old_deposit && proxies_len_after < proxies_len_before {
expected_proxies.insert(who.clone(), (proxies, new_deposit));
expected_reserved_amounts.insert(
who,
current_reserved.saturating_sub(old_deposit - new_deposit),
);
} else {
// Nothing should change
expected_proxies.insert(who.clone(), (proxies, old_deposit));
expected_reserved_amounts.insert(who, current_reserved);
}
}

let pre_upgrade_state = (expected_proxies, expected_reserved_amounts);
Ok(pre_upgrade_state.encode())
}

fn on_runtime_upgrade() -> Weight {
let mut reads = 0u64;
let mut writes = 0u64;
let mut proxies_removed_total = 0u64;

Proxies::<Runtime>::translate(
|who: AccountId, (mut proxies, old_deposit): PrevProxiesValue<Runtime>| {
// Remove filter out IdentityJudgement proxies.
let proxies_len_before = proxies.len() as u64;
proxies.retain(|proxy| {
if proxy.proxy_type == PrevProxyType::IdentityJudgement {
false
} else {
true
}
});
let proxies_len_after = proxies.len() as u64;

let deposit = if proxies_len_before - proxies_len_after > 0 {
log::info!(
"Removing {} IdentityJudgement proxies for {:?}",
proxies_len_before - proxies_len_after,
&who
);
proxies_removed_total
.saturating_accrue(proxies_len_before - proxies_len_after);

let new_deposit =
pallet_proxy::Pallet::<Runtime>::deposit(proxies.len() as u32);

// Be kind and don't increase the deposit in case it increased (can
// happen if param change).
let deposit =
if new_deposit < old_deposit { new_deposit } else { old_deposit };
if deposit < old_deposit {
writes.saturating_inc();
<Balances as ReservableCurrency<AccountId>>::unreserve(
&who,
old_deposit - deposit,
);
}

deposit
} else {
// Nothing to do, use the old deposit.
old_deposit
};

reads.saturating_accrue(proxies_len_before as u64 + 1);
writes.saturating_accrue(proxies_len_after as u64 + 1);
Some((proxies, deposit))
},
);

log::info!("Removed {} IdentityJudgement proxies in total", proxies_removed_total);
<Runtime as frame_system::Config>::DbWeight::get().reads_writes(reads, writes)
}

#[cfg(feature = "try-runtime")]
fn post_upgrade(state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
use frame_support::ensure;
use sp_runtime::TryRuntimeError;

let (expected_proxies, expected_total_reserved): (
BTreeMap<AccountId, PrevProxiesValue<Runtime>>,
BTreeMap<AccountId, Balance>,
) = Decode::decode(&mut &state[..]).expect("Failed to decode pre-upgrade state");

// Check Proxies storage is as expected
for (who, (proxies, deposit)) in Proxies::<Runtime>::iter() {
match expected_proxies.get(&who) {
Some((expected_proxies, expected_deposit)) => {
ensure!(&proxies == expected_proxies, "Unexpected Proxy");
ensure!(&deposit == expected_deposit, "Unexpected deposit");
},
None => {
return Err(TryRuntimeError::Other("Missing Proxy".into()));
},
}
}

// Check the total reserved amounts for every account is as expected
for (who, expected_reserved) in expected_total_reserved.iter() {
let current_reserved =
<Balances as ReservableCurrency<AccountId>>::reserved_balance(who);

ensure!(current_reserved == *expected_reserved, "Reserved balance mismatch");
}

// Check there are no extra entries in the expected state that are not in the
// current state
for (who, _) in expected_proxies.iter() {
if !Proxies::<Runtime>::contains_key(who) {
return Err(TryRuntimeError::Other("Extra entry in expected state".into()));
}
}

Ok(())
}
}
}

/// Unreleased migrations. Add new ones here:
pub type Unreleased = (
frame_support::migrations::RemovePallet<StateTrieMigrationName, RocksDbWeight>,
Expand Down Expand Up @@ -1960,6 +2162,7 @@ pub mod migrations {
IdentityMigratorPalletName,
<Runtime as frame_system::Config>::DbWeight,
>,
clear_judgement_proxies::Migration,
);

/// Migrations/checks that do not need to be versioned and can run on every update.
Expand Down

0 comments on commit 89a2f5b

Please sign in to comment.