pallet-balances: Fix inactive funds migration

frame/balances/src/lib.rs

@@ -246,8 +246,13 @@ pub mod pallet {
 		type ReserveIdentifier: Parameter + Member + MaxEncodedLen + Ord + Copy;
 	}
 
+	/// The current storage version.
+	const STORAGE_VERSION: frame_support::traits::StorageVersion =
+		frame_support::traits::StorageVersion::new(1);
+
 	#[pallet::pallet]
 	#[pallet::generate_store(pub(super) trait Store)]
+	#[pallet::storage_version(STORAGE_VERSION)]
 	pub struct Pallet<T, I = ()>(PhantomData<(T, I)>);
 
 	#[pallet::call]
@@ -556,13 +561,6 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
-	/// Storage version of the pallet.
-	///
-	/// This is set to v2.0.0 for new networks.
-	#[pallet::storage]
-	pub(super) type StorageVersion<T: Config<I>, I: 'static = ()> =
-		StorageValue<_, Releases, ValueQuery>;
-
 	#[pallet::genesis_config]
 	pub struct GenesisConfig<T: Config<I>, I: 'static = ()> {
 		pub balances: Vec<(T::AccountId, T::Balance)>,
@@ -581,8 +579,6 @@ pub mod pallet {
 			let total = self.balances.iter().fold(Zero::zero(), |acc: T::Balance, &(_, n)| acc + n);
 			<TotalIssuance<T, I>>::put(total);
 
-			<StorageVersion<T, I>>::put(Releases::V2_0_0);
-
 			for (_, balance) in &self.balances {
 				assert!(
 					*balance >= <T as Config<I>>::ExistentialDeposit::get(),
@@ -727,21 +723,6 @@ impl<Balance: Saturating + Copy + Ord> AccountData<Balance> {
 	}
 }
 
-// A value placed in storage that represents the current version of the Balances storage.
-// This value is used by the `on_runtime_upgrade` logic to determine whether we run
-// storage migration logic. This should match directly with the semantic versions of the Rust crate.
-#[derive(Encode, Decode, Clone, Copy, PartialEq, Eq, RuntimeDebug, MaxEncodedLen, TypeInfo)]
-enum Releases {
-	V1_0_0,
-	V2_0_0,
-}
-
-impl Default for Releases {
-	fn default() -> Self {
-		Releases::V1_0_0
-	}
-}
-
 pub struct DustCleaner<T: Config<I>, I: 'static = ()>(
 	Option<(T::AccountId, NegativeImbalance<T, I>)>,
 );

frame/balances/src/migration.rs

@@ -15,50 +15,57 @@
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
 use super::*;
-use frame_support::{pallet_prelude::*, traits::OnRuntimeUpgrade, weights::Weight};
+use frame_support::{
+	pallet_prelude::*,
+	traits::{OnRuntimeUpgrade, PalletInfoAccess},
+	weights::Weight,
+};
 
-// NOTE: This must be used alongside the account whose balance is expected to be inactive.
-// Generally this will be used for the XCM teleport checking account.
-pub struct MigrateToTrackInactive<T, A>(PhantomData<(T, A)>);
-impl<T: Config, A: Get<T::AccountId>> OnRuntimeUpgrade for MigrateToTrackInactive<T, A> {
-	fn on_runtime_upgrade() -> Weight {
-		let current_version = Pallet::<T>::current_storage_version();
-		let onchain_version = Pallet::<T>::on_chain_storage_version();
+fn migrate_v0_to_v1<T: Config<I>, I: 'static>(accounts: &[T::AccountId]) -> Weight {
+	let onchain_version = Pallet::<T, I>::on_chain_storage_version();
+
+	if onchain_version == 0 {
+		let total = accounts
+			.iter()
+			.map(|a| Pallet::<T, I>::total_balance(a))
+			.fold(T::Balance::zero(), |a, e| a.saturating_add(e));
+		Pallet::<T, I>::deactivate(total);
+
+		// Remove the old `StorageVersion` type.
+		frame_support::storage::unhashed::kill(&frame_support::storage::storage_prefix(
+			Pallet::<T, I>::name().as_bytes(),
+			"StorageVersion".as_bytes(),
+		));
 
-		if onchain_version == 0 && current_version == 1 {
-			let b = Pallet::<T>::total_balance(&A::get());
-			Pallet::<T>::deactivate(b);
-			current_version.put::<Pallet<T>>();
-			log::info!(target: "runtime::balances", "Storage to version {:?}", current_version);
-			T::DbWeight::get().reads_writes(4, 3)
-		} else {
-			log::info!(target: "runtime::balances",  "Migration did not execute. This probably should be removed");
-			T::DbWeight::get().reads(2)
-		}
+		// Set storage version to `1`.
+		StorageVersion::new(1).put::<Pallet<T, I>>();
+
+		log::info!(target: "runtime::balances", "Storage to version 1");
+		T::DbWeight::get().reads_writes(2 + accounts.len() as u64, 3)
+	} else {
+		log::info!(target: "runtime::balances",  "Migration did not execute. This probably should be removed");
+		T::DbWeight::get().reads(1)
 	}
 }
 
 // NOTE: This must be used alongside the account whose balance is expected to be inactive.
 // Generally this will be used for the XCM teleport checking account.
-pub struct MigrateManyToTrackInactive<T, A>(PhantomData<(T, A)>);
-impl<T: Config, A: Get<Vec<T::AccountId>>> OnRuntimeUpgrade for MigrateManyToTrackInactive<T, A> {
+pub struct MigrateToTrackInactive<T, A, I = ()>(PhantomData<(T, A, I)>);
+impl<T: Config<I>, A: Get<T::AccountId>, I: 'static> OnRuntimeUpgrade
+	for MigrateToTrackInactive<T, A, I>
+{
 	fn on_runtime_upgrade() -> Weight {
-		let current_version = Pallet::<T>::current_storage_version();
-		let onchain_version = Pallet::<T>::on_chain_storage_version();
+		migrate_v0_to_v1::<T, I>(&[A::get()])
+	}
+}
 
-		if onchain_version == 0 && current_version == 1 {
-			let accounts = A::get();
-			let total = accounts
-				.iter()
-				.map(|a| Pallet::<T>::total_balance(a))
-				.fold(T::Balance::zero(), |a, e| a.saturating_add(e));
-			Pallet::<T>::deactivate(total);
-			current_version.put::<Pallet<T>>();
-			log::info!(target: "runtime::balances", "Storage to version {:?}", current_version);
-			T::DbWeight::get().reads_writes(3 + accounts.len() as u64, 3)
-		} else {
-			log::info!(target: "runtime::balances",  "Migration did not execute. This probably should be removed");
-			T::DbWeight::get().reads(2)
-		}
+// NOTE: This must be used alongside the accounts whose balance is expected to be inactive.
+// Generally this will be used for the XCM teleport checking accounts.
+pub struct MigrateManyToTrackInactive<T, A, I = ()>(PhantomData<(T, A, I)>);
+impl<T: Config<I>, A: Get<Vec<T::AccountId>>, I: 'static> OnRuntimeUpgrade
+	for MigrateManyToTrackInactive<T, A, I>
+{
+	fn on_runtime_upgrade() -> Weight {
+		migrate_v0_to_v1::<T, I>(&A::get())
 	}
 }

frame/executive/src/lib.rs

@@ -951,13 +951,13 @@ mod tests {
 		block_import_works_inner(
 			new_test_ext_v0(1),
 			array_bytes::hex_n_into_unchecked(
-				"0d786e24c1f9e6ce237806a22c005bbbc7dee4edd6692b6c5442843d164392de",
+				"216e61b2689d1243eb56d89c9084db48e50ebebc4871d758db131432c675d7c0",
 			),
 		);
 		block_import_works_inner(
 			new_test_ext(1),
 			array_bytes::hex_n_into_unchecked(
-				"348485a4ab856467b440167e45f99b491385e8528e09b0e51f85f814a3021c93",
+				"4738b4c0aab02d6ddfa62a2a6831ccc975a9f978f7db8d7ea8e68eba8639530a",
 			),
 		);
 	}