Merge pull request #4 from liamaharon/liam-proxy-migration

Judgement proxy migration

relay/kusama/src/lib.rs

@@ -1834,10 +1834,7 @@ pub mod migrations {
 			log::info!(target: "runtime::session_keys", "Collecting pre-upgrade session keys state");
 			let key_ids = SessionKeys::key_ids();
 			frame_support::ensure!(
-				key_ids
-					.into_iter()
-					.find(|&k| *k == sp_core::crypto::key_types::IM_ONLINE)
-					.is_none(),
+				!key_ids.iter().any(|k| *k == sp_core::crypto::key_types::IM_ONLINE),
 				"New session keys contain the ImOnline key that should have been removed",
 			);
 			let storage_key = pallet_session::QueuedKeys::<Runtime>::hashed_key();
@@ -1853,7 +1850,7 @@ pub mod migrations {
 					state.extend_from_slice(keys.get_raw(*key_id));
 				}
 			});
-			frame_support::ensure!(state.len() > 0, "Queued keys are not empty before upgrade");
+			frame_support::ensure!(!state.is_empty(), "Queued keys are not empty before upgrade");
 			Ok(state)
 		}
 
@@ -1884,7 +1881,10 @@ pub mod migrations {
 					new_state.extend_from_slice(keys.get_raw(*key_id));
 				}
 			});
-			frame_support::ensure!(new_state.len() > 0, "Queued keys are not empty after upgrade");
+			frame_support::ensure!(
+				!new_state.is_empty(),
+				"Queued keys are not empty after upgrade"
+			);
 			frame_support::ensure!(
 				old_state == new_state,
 				"Pre-upgrade and post-upgrade keys do not match!"
@@ -1929,6 +1929,214 @@ 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 {
+						// If there're no proxies left, they should be removed
+						if proxies.len() > 0 {
+							expected_proxies.insert(who.clone(), (proxies, new_deposit));
+						}
+						expected_reserved_amounts.insert(
+							who,
+							current_reserved.saturating_sub(old_deposit - new_deposit),
+						);
+					} else {
+						// Deposit should not change. If any proxies needed to be removed, this
+						// won't impact that.
+						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| proxy.proxy_type != PrevProxyType::IdentityJudgement);
+						let proxies_len_after = proxies.len() as u64;
+
+						let deposit = if proxies_len_before > proxies_len_after {
+							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 = new_deposit.min(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 + 1);
+						writes.saturating_accrue(proxies_len_after + 1);
+
+						// No need to keep the k/v around if there're no proxies left.
+						match proxies.is_empty() {
+							true => {
+								debug_assert_eq!(deposit, 0);
+								None
+							},
+							false => 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"));
+						},
+					}
+				}
+
+				// 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"));
+					}
+				}
+
+				Ok(())
+			}
+		}
+	}
+
 	/// Unreleased migrations. Add new ones here:
 	pub type Unreleased = (
 		frame_support::migrations::RemovePallet<StateTrieMigrationName, RocksDbWeight>,
@@ -1962,6 +2170,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.