Fix Stuck Collator Funds

cumulus/pallets/collator-selection/src/lib.rs

@@ -121,7 +121,7 @@ pub mod pallet {
 	use sp_std::vec::Vec;
 
 	/// The in-code storage version.
-	const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
+	const STORAGE_VERSION: StorageVersion = StorageVersion::new(2);
 
 	type BalanceOf<T> =
 		<<T as Config>::Currency as Currency<<T as SystemConfig>::AccountId>>::Balance;

cumulus/pallets/collator-selection/src/migration.rs

@@ -17,9 +17,115 @@
 //! A module that is responsible for migration of storage for Collator Selection.
 
 use super::*;
-use frame_support::traits::OnRuntimeUpgrade;
+use frame_support::traits::{OnRuntimeUpgrade, UncheckedOnRuntimeUpgrade};
 use log;
 
+/// Migrate to v2. Should have been part of https://github.com/paritytech/polkadot-sdk/pull/1340
+pub mod v2 {
+	use super::*;
+	use frame_support::{
+		pallet_prelude::*,
+		storage_alias,
+		traits::{Currency, ReservableCurrency},
+	};
+	use sp_runtime::traits::{Saturating, Zero};
+	#[cfg(feature = "try-runtime")]
+	use sp_std::vec::Vec;
+
+	/// [`UncheckedMigrationToV2`] wrapped in a
+	/// [`VersionedMigration`](frame_support::migrations::VersionedMigration), ensuring the
+	/// migration is only performed when on-chain version is 1.
+	pub type MigrationToV2<T> = frame_support::migrations::VersionedMigration<
+		1,
+		2,
+		UncheckedMigrationToV2<T>,
+		Pallet<T>,
+		<T as frame_system::Config>::DbWeight,
+	>;
+
+	#[storage_alias]
+	pub type Candidates<T: Config> = StorageValue<
+		Pallet<T>,
+		BoundedVec<CandidateInfo<<T as frame_system::Config>::AccountId, <<T as Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance>, <T as Config>::MaxCandidates>,
+		ValueQuery,
+	>;
+
+	/// Migrate to V2.
+	///
+	/// Note: This should have been in https://github.com/paritytech/polkadot-sdk/pull/1340.
+	pub struct UncheckedMigrationToV2<T>(sp_std::marker::PhantomData<T>);
+	impl<T: Config> UncheckedOnRuntimeUpgrade for UncheckedMigrationToV2<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let on_chain_version = Pallet::<T>::on_chain_storage_version();
+			if on_chain_version == 1 {
+				let mut weight = Weight::zero();
+				let mut count: u32 = 0;
+				// candidates who exist under the old `Candidates` key
+				let candidates = Candidates::<T>::take();
+
+				// New candidates who have registered since the upgrade. Under normal circumstances,
+				// this should not exist because the migration should be applied when the upgrade
+				// happens. But in Polkadot/Kusama we messed this up, and people registered under
+				// `CandidateList` while their funds were locked in `Candidates`.
+				let new_candidate_list = CandidateList::<T>::get();
+				if new_candidate_list.len().is_zero() {
+					// The new list is empty, so this is essentially being applied correctly. We
+					// just put the candidates into the new storage item.
+					CandidateList::<T>::put(&candidates);
+					// 1 write for the new list
+					weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 1));
+				} else {
+					// Oops, the runtime upgraded without the migration. There are new candidates in
+					// `CandidateList`. So, let's just refund the old ones and assume they have
+					// already started participating in the new system.
+					for candidate in candidates {
+						let _err = T::Currency::unreserve(&candidate.who, candidate.deposit);
+						count.saturating_inc();
+					}
+					// TODO: set each accrue to weight of `unreserve`
+					weight.saturating_accrue(T::DbWeight::get().reads_writes(0, count as u64));
+				}
+
+				StorageVersion::new(2).put::<Pallet<T>>();
+				log::info!(
+					target: LOG_TARGET,
+					"Unreserved locked bond of {} candidates, upgraded storage to version 2",
+					count,
+				);
+				// 1 read/write for storage version
+				weight.saturating_accrue(T::DbWeight::get().reads_writes(3, 2));
+				weight
+			} else {
+				log::info!(
+					target: LOG_TARGET,
+					"Migration did not execute. This probably should be removed"
+				);
+				T::DbWeight::get().reads(1)
+			}
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::DispatchError> {
+			let number_of_candidates = Candidates::<T>::get().to_vec().len();
+			Ok((number_of_candidates as u32).encode())
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_number_of_candidates: Vec<u8>) -> Result<(), sp_runtime::DispatchError> {
+			let new_number_of_candidates = Candidates::<T>::get().to_vec().len();
+			assert_eq!(
+				new_number_of_candidates, 0 as usize,
+				"after migration, the candidates map should be empty"
+			);
+
+			let on_chain_version = Pallet::<T>::on_chain_storage_version();
+			frame_support::ensure!(on_chain_version >= 1, "must_upgrade");
+
+			Ok(())
+		}
+	}
+}
+
 /// Version 1 Migration
 /// This migration ensures that any existing `Invulnerables` storage lists are sorted.
 pub mod v1 {
@@ -90,3 +196,54 @@ pub mod v1 {
 		}
 	}
 }
+
+#[cfg(all(feature = "try-runtime", test))]
+mod tests {
+	use super::*;
+	use crate::{
+		migration::v2::Candidates,
+		mock::{new_test_ext, Balances, Test},
+	};
+	use frame_support::{
+		traits::{Currency, ReservableCurrency, StorageVersion},
+		BoundedVec,
+	};
+	use sp_runtime::traits::ConstU32;
+
+	#[test]
+	fn migrate_to_v2() {
+		new_test_ext().execute_with(|| {
+			let storage_version = StorageVersion::new(1);
+			storage_version.put::<Pallet<Test>>();
+
+			let who = 1u64;
+
+			// Set balance to 100
+			Balances::make_free_balance_be(&who, 100u64);
+			// Reserve 20 -> 10 for the "old" candidacy and 10 for the "new"
+			Balances::reserve(&who, 20u64).unwrap();
+			// Candidate info
+			let candidate = CandidateInfo { who, deposit: 10u64 };
+			let bounded_candidates =
+				BoundedVec::<CandidateInfo<u64, u64>, ConstU32<20>>::try_from(vec![
+					candidate.clone()
+				])
+				.expect("it works");
+			// Candidate is in the old `Candidates` and the new `CandidateList`
+			Candidates::<Test>::put(bounded_candidates.clone());
+			CandidateList::<Test>::put(bounded_candidates);
+			// Sanity check
+			assert_eq!(Balances::free_balance(who), 80);
+
+			// Run migration
+			v2::MigrationToV2::<Test>::on_runtime_upgrade();
+
+			// 10 should have been unreserved from the old candidacy
+			assert_eq!(Balances::free_balance(who), 90);
+			// The storage item should be gone
+			assert!(Candidates::<Test>::get().is_empty());
+			// The new storage item should be preserved
+			assert_eq!(CandidateList::<Test>::get(), vec![candidate]);
+		});
+	}
+}

cumulus/parachains/runtimes/assets/asset-hub-rococo/src/lib.rs

@@ -973,10 +973,10 @@ pub type UncheckedExtrinsic =
 /// Migrations to apply on runtime upgrade.
 #[allow(deprecated)]
 pub type Migrations = (
-	pallet_collator_selection::migration::v1::MigrateToV1<Runtime>,
 	InitStorageVersions,
 	// unreleased
 	cumulus_pallet_xcmp_queue::migration::v4::MigrationToV4<Runtime>,
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	// permanent
 	pallet_xcm::migration::MigrateToLatestXcmVersion<Runtime>,
 );

cumulus/parachains/runtimes/assets/asset-hub-westend/src/lib.rs

@@ -963,7 +963,7 @@ pub type Migrations = (
 	// v9420
 	pallet_nfts::migration::v1::MigrateToV1<Runtime>,
 	// unreleased
-	pallet_collator_selection::migration::v1::MigrateToV1<Runtime>,
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	// unreleased
 	pallet_multisig::migrations::v1::MigrateToV1<Runtime>,
 	// unreleased

cumulus/parachains/runtimes/bridge-hubs/bridge-hub-rococo/src/lib.rs

@@ -139,7 +139,7 @@ pub type UncheckedExtrinsic =
 
 /// Migrations to apply on runtime upgrade.
 pub type Migrations = (
-	pallet_collator_selection::migration::v1::MigrateToV1<Runtime>,
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	pallet_multisig::migrations::v1::MigrateToV1<Runtime>,
 	InitStorageVersions,
 	cumulus_pallet_xcmp_queue::migration::v4::MigrationToV4<Runtime>,

cumulus/parachains/runtimes/bridge-hubs/bridge-hub-westend/src/lib.rs

@@ -118,7 +118,7 @@ pub type UncheckedExtrinsic =
 
 /// Migrations to apply on runtime upgrade.
 pub type Migrations = (
-	pallet_collator_selection::migration::v1::MigrateToV1<Runtime>,
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	pallet_multisig::migrations::v1::MigrateToV1<Runtime>,
 	InitStorageVersions,
 	// unreleased

cumulus/parachains/runtimes/collectives/collectives-westend/src/lib.rs

@@ -722,7 +722,7 @@ pub type UncheckedExtrinsic =
 /// `OnRuntimeUpgrade`. Included migrations must be idempotent.
 type Migrations = (
 	// unreleased
-	pallet_collator_selection::migration::v1::MigrateToV1<Runtime>,
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	// unreleased
 	cumulus_pallet_xcmp_queue::migration::v4::MigrationToV4<Runtime>,
 	// permanent

cumulus/parachains/runtimes/contracts/contracts-rococo/src/lib.rs

@@ -98,6 +98,8 @@ pub type UncheckedExtrinsic =
 
 /// Migrations to apply on runtime upgrade.
 pub type Migrations = (
+	pallet_collator_selection::migration::v1::MigrateToV1<Runtime>,
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	cumulus_pallet_parachain_system::migration::Migration<Runtime>,
 	cumulus_pallet_xcmp_queue::migration::v2::MigrationToV2<Runtime>,
 	cumulus_pallet_xcmp_queue::migration::v3::MigrationToV3<Runtime>,

cumulus/parachains/runtimes/coretime/coretime-rococo/src/lib.rs

@@ -108,6 +108,7 @@ pub type UncheckedExtrinsic =
 
 /// Migrations to apply on runtime upgrade.
 pub type Migrations = (
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	cumulus_pallet_xcmp_queue::migration::v4::MigrationToV4<Runtime>,
 	pallet_broker::migration::MigrateV0ToV1<Runtime>,
 	// permanent

cumulus/parachains/runtimes/coretime/coretime-westend/src/lib.rs

@@ -108,6 +108,7 @@ pub type UncheckedExtrinsic =
 
 /// Migrations to apply on runtime upgrade.
 pub type Migrations = (
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	cumulus_pallet_xcmp_queue::migration::v4::MigrationToV4<Runtime>,
 	pallet_broker::migration::MigrateV0ToV1<Runtime>,
 	// permanent

cumulus/parachains/runtimes/people/people-rococo/src/lib.rs

@@ -102,6 +102,7 @@ pub type UncheckedExtrinsic =
 
 /// Migrations to apply on runtime upgrade.
 pub type Migrations = (
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	// permanent
 	pallet_xcm::migration::MigrateToLatestXcmVersion<Runtime>,
 );

cumulus/parachains/runtimes/people/people-westend/src/lib.rs

@@ -102,6 +102,7 @@ pub type UncheckedExtrinsic =
 
 /// Migrations to apply on runtime upgrade.
 pub type Migrations = (
+	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
 	// permanent
 	pallet_xcm::migration::MigrateToLatestXcmVersion<Runtime>,
 );