General-Beck
diff --git a/‎frame/balances/src/lib.rs‎
Lines changed: 21 additions & 84 deletions b/‎frame/balances/src/lib.rs‎
Lines changed: 21 additions & 84 deletions
diff --git a/‎frame/balances/src/migration.rs‎
Lines changed: 88 additions & 0 deletions b/‎frame/balances/src/migration.rs‎
Lines changed: 88 additions & 0 deletions
diff --git a/‎frame/democracy/src/lib.rs‎
Lines changed: 2 additions & 2 deletions b/‎frame/democracy/src/lib.rs‎
Lines changed: 2 additions & 2 deletions
diff --git a/‎frame/executive/src/lib.rs‎
Lines changed: 14 additions & 3 deletions b/‎frame/executive/src/lib.rs‎
Lines changed: 14 additions & 3 deletions
@@ -157,6 +157,7 @@ mod tests_composite;
 mod tests;
 #[cfg(feature = "runtime-benchmarks")]
 mod benchmarking;
+mod migration;
 
 use sp_std::prelude::*;
 use sp_std::{cmp, result, mem, fmt::Debug, ops::BitOr, convert::Infallible};
@@ -348,6 +349,21 @@ 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)]
+enum Releases {
+	V1_0_0,
+	V2_0_0,
+}
+
+impl Default for Releases {
+	fn default() -> Self {
+		Releases::V1_0_0
+	}
+}
+
 decl_storage! {
 	trait Store for Module<T: Trait<I>, I: Instance=DefaultInstance> as Balances {
 		/// The total units issued in the system.
@@ -367,10 +383,10 @@ decl_storage! {
 		/// NOTE: Should only be accessed when setting, changing and freeing a lock.
 		pub Locks get(fn locks): map hasher(blake2_256) T::AccountId => Vec<BalanceLock<T::Balance>>;
 
-		/// True if network has been upgraded to this version.
+		/// Storage version of the pallet.
 		///
-		/// True for new networks.
-		IsUpgraded build(|_: &GenesisConfig<T, I>| true): bool;
+		/// This is set to v2.0.0 for new networks.
+		StorageVersion build(|_: &GenesisConfig<T, I>| Releases::V2_0_0): Releases;
 	}
 	add_extra_genesis {
 		config(balances): Vec<(T::AccountId, T::Balance)>;
@@ -518,11 +534,8 @@ decl_module! {
 			<Self as Currency<_>>::transfer(&transactor, &dest, value, KeepAlive)?;
 		}
 
-		fn on_initialize() {
-			if !IsUpgraded::<I>::get() {
-				IsUpgraded::<I>::put(true);
-				Self::do_upgrade();
-			}
+		fn on_runtime_upgrade() {
+			migration::on_runtime_upgrade::<T, I>();
 		}
 	}
 }
@@ -548,82 +561,6 @@ impl<Balance, BlockNumber> OldBalanceLock<Balance, BlockNumber> {
 impl<T: Trait<I>, I: Instance> Module<T, I> {
 	// PRIVATE MUTABLES
 
-	// Upgrade from the pre-#4649 balances/vesting into the new balances.
-	pub fn do_upgrade() {
-		sp_runtime::print("Upgrading account balances...");
-		// First, migrate from old FreeBalance to new Account.
-		// We also move all locks across since only accounts with FreeBalance values have locks.
-		// FreeBalance: map T::AccountId => T::Balance
-		for (hash, free) in StorageIterator::<T::Balance>::new(b"Balances", b"FreeBalance").drain() {
-			let mut account = AccountData { free, ..Default::default() };
-			// Locks: map T::AccountId => Vec<BalanceLock>
-			let old_locks = get_storage_value::<Vec<OldBalanceLock<T::Balance, T::BlockNumber>>>(b"Balances", b"Locks", &hash);
-			if let Some(locks) = old_locks {
-				let locks = locks.into_iter()
-					.map(|i| {
-						let (result, expiry) = i.upgraded();
-						if expiry != T::BlockNumber::max_value() {
-							// Any `until`s that are not T::BlockNumber::max_value come from
-							// democracy and need to be migrated over there.
-							// Democracy: Locks get(locks): map T::AccountId => Option<T::BlockNumber>;
-							put_storage_value(b"Democracy", b"Locks", &hash, expiry);
-						}
-						result
-					})
-					.collect::<Vec<_>>();
-				for l in locks.iter() {
-					if l.reasons == Reasons::All || l.reasons == Reasons::Misc {
-						account.misc_frozen = account.misc_frozen.max(l.amount);
-					}
-					if l.reasons == Reasons::All || l.reasons == Reasons::Fee {
-						account.fee_frozen = account.fee_frozen.max(l.amount);
-					}
-				}
-				put_storage_value(b"Balances", b"Locks", &hash, locks);
-			}
-			put_storage_value(b"Balances", b"Account", &hash, account);
-		}
-		// Second, migrate old ReservedBalance into new Account.
-		// ReservedBalance: map T::AccountId => T::Balance
-		for (hash, reserved) in StorageIterator::<T::Balance>::new(b"Balances", b"ReservedBalance").drain() {
-			let mut account = get_storage_value::<AccountData<T::Balance>>(b"Balances", b"Account", &hash).unwrap_or_default();
-			account.reserved = reserved;
-			put_storage_value(b"Balances", b"Account", &hash, account);
-		}
-
-		// Finally, migrate vesting and ensure locks are in place. We will be lazy and just lock
-		// for the maximum amount (i.e. at genesis). Users will need to call "vest" to reduce the
-		// lock to something sensible.
-		// pub Vesting: map T::AccountId => Option<VestingSchedule>;
-		for (hash, vesting) in StorageIterator::<(T::Balance, T::Balance, T::BlockNumber)>::new(b"Balances", b"Vesting").drain() {
-			let mut account = get_storage_value::<AccountData<T::Balance>>(b"Balances", b"Account", &hash).unwrap_or_default();
-			let mut locks = get_storage_value::<Vec<BalanceLock<T::Balance>>>(b"Balances", b"Locks", &hash).unwrap_or_default();
-			locks.push(BalanceLock {
-				id: *b"vesting ",
-				amount: vesting.0.clone(),
-				reasons: Reasons::Misc,
-			});
-			account.misc_frozen = account.misc_frozen.max(vesting.0.clone());
-			put_storage_value(b"Vesting", b"Vesting", &hash, vesting);
-			put_storage_value(b"Balances", b"Locks", &hash, locks);
-			put_storage_value(b"Balances", b"Account", &hash, account);
-		}
-
-		for (hash, balances) in StorageIterator::<AccountData<T::Balance>>::new(b"Balances", b"Account").drain() {
-			let nonce = take_storage_value::<T::Index>(b"System", b"AccountNonce", &hash).unwrap_or_default();
-			let mut refs: system::RefCount = 0;
-			// The items in Kusama that would result in a ref count being incremented.
-			if have_storage_value(b"Democracy", b"Proxy", &hash) { refs += 1 }
-			// We skip Recovered since it's being replaced anyway.
-			let mut prefixed_hash = twox_64(&b":session:keys"[..]).to_vec();
-			prefixed_hash.extend(&b":session:keys"[..]);
-			prefixed_hash.extend(&hash[..]);
-			if have_storage_value(b"Session", b"NextKeys", &prefixed_hash) { refs += 1 }
-			if have_storage_value(b"Staking", b"Bonded", &hash) { refs += 1 }
-			put_storage_value(b"System", b"Account", &hash, (nonce, refs, &balances));
-		}
-	}
-
 	/// Get the free balance of an account.
 	pub fn free_balance(who: impl sp_std::borrow::Borrow<T::AccountId>) -> T::Balance {
 		Self::account(who.borrow()).free
 
@@ -0,0 +1,88 @@
+use super::*;
+
+pub fn on_runtime_upgrade<T: Trait<I>, I: Instance>() {
+	match StorageVersion::<I>::get() {
+		Releases::V2_0_0 => return,
+		Releases::V1_0_0 => upgrade_v1_to_v2::<T, I>(),
+	}
+}
+
+// Upgrade from the pre-#4649 balances/vesting into the new balances.
+fn upgrade_v1_to_v2<T: Trait<I>, I: Instance>() {
+	sp_runtime::print("Upgrading account balances...");
+	// First, migrate from old FreeBalance to new Account.
+	// We also move all locks across since only accounts with FreeBalance values have locks.
+	// FreeBalance: map T::AccountId => T::Balance
+	for (hash, free) in StorageIterator::<T::Balance>::new(b"Balances", b"FreeBalance").drain() {
+		let mut account = AccountData { free, ..Default::default() };
+		// Locks: map T::AccountId => Vec<BalanceLock>
+		let old_locks = get_storage_value::<Vec<OldBalanceLock<T::Balance, T::BlockNumber>>>(b"Balances", b"Locks", &hash);
+		if let Some(locks) = old_locks {
+			let locks = locks.into_iter()
+				.map(|i| {
+					let (result, expiry) = i.upgraded();
+					if expiry != T::BlockNumber::max_value() {
+						// Any `until`s that are not T::BlockNumber::max_value come from
+						// democracy and need to be migrated over there.
+						// Democracy: Locks get(locks): map T::AccountId => Option<T::BlockNumber>;
+						put_storage_value(b"Democracy", b"Locks", &hash, expiry);
+					}
+					result
+				})
+				.collect::<Vec<_>>();
+			for l in locks.iter() {
+				if l.reasons == Reasons::All || l.reasons == Reasons::Misc {
+					account.misc_frozen = account.misc_frozen.max(l.amount);
+				}
+				if l.reasons == Reasons::All || l.reasons == Reasons::Fee {
+					account.fee_frozen = account.fee_frozen.max(l.amount);
+				}
+			}
+			put_storage_value(b"Balances", b"Locks", &hash, locks);
+		}
+		put_storage_value(b"Balances", b"Account", &hash, account);
+	}
+	// Second, migrate old ReservedBalance into new Account.
+	// ReservedBalance: map T::AccountId => T::Balance
+	for (hash, reserved) in StorageIterator::<T::Balance>::new(b"Balances", b"ReservedBalance").drain() {
+		let mut account = get_storage_value::<AccountData<T::Balance>>(b"Balances", b"Account", &hash).unwrap_or_default();
+		account.reserved = reserved;
+		put_storage_value(b"Balances", b"Account", &hash, account);
+	}
+
+	// Finally, migrate vesting and ensure locks are in place. We will be lazy and just lock
+	// for the maximum amount (i.e. at genesis). Users will need to call "vest" to reduce the
+	// lock to something sensible.
+	// pub Vesting: map T::AccountId => Option<VestingSchedule>;
+	for (hash, vesting) in StorageIterator::<(T::Balance, T::Balance, T::BlockNumber)>::new(b"Balances", b"Vesting").drain() {
+		let mut account = get_storage_value::<AccountData<T::Balance>>(b"Balances", b"Account", &hash).unwrap_or_default();
+		let mut locks = get_storage_value::<Vec<BalanceLock<T::Balance>>>(b"Balances", b"Locks", &hash).unwrap_or_default();
+		locks.push(BalanceLock {
+			id: *b"vesting ",
+			amount: vesting.0.clone(),
+			reasons: Reasons::Misc,
+		});
+		account.misc_frozen = account.misc_frozen.max(vesting.0.clone());
+		put_storage_value(b"Vesting", b"Vesting", &hash, vesting);
+		put_storage_value(b"Balances", b"Locks", &hash, locks);
+		put_storage_value(b"Balances", b"Account", &hash, account);
+	}
+
+	for (hash, balances) in StorageIterator::<AccountData<T::Balance>>::new(b"Balances", b"Account").drain() {
+		let nonce = take_storage_value::<T::Index>(b"System", b"AccountNonce", &hash).unwrap_or_default();
+		let mut refs: system::RefCount = 0;
+		// The items in Kusama that would result in a ref count being incremented.
+		if have_storage_value(b"Democracy", b"Proxy", &hash) { refs += 1 }
+		// We skip Recovered since it's being replaced anyway.
+		let mut prefixed_hash = twox_64(&b":session:keys"[..]).to_vec();
+		prefixed_hash.extend(&b":session:keys"[..]);
+		prefixed_hash.extend(&hash[..]);
+		if have_storage_value(b"Session", b"NextKeys", &prefixed_hash) { refs += 1 }
+		if have_storage_value(b"Staking", b"Bonded", &hash) { refs += 1 }
+		put_storage_value(b"System", b"Account", &hash, (nonce, refs, &balances));
+	}
+
+	take_storage_value::<T::Index>(b"Balances", b"IsUpgraded", &[]);
+
+	StorageVersion::<I>::put(Releases::V2_0_0);
+}
@@ -1555,7 +1555,7 @@ mod tests {
 	};
 	use sp_core::H256;
 	use sp_runtime::{
-		traits::{BlakeTwo256, IdentityLookup, Bounded, BadOrigin, OnInitialize},
+		traits::{BlakeTwo256, IdentityLookup, Bounded, BadOrigin, OnRuntimeUpgrade},
 		testing::Header, Perbill,
 	};
 	use pallet_balances::{BalanceLock, Error as BalancesError};
@@ -1713,7 +1713,7 @@ mod tests {
 		];
 		s.top = data.into_iter().collect();
 		sp_io::TestExternalities::new(s).execute_with(|| {
-			Balances::on_initialize(1);
+			Balances::on_runtime_upgrade();
 			assert_eq!(Balances::free_balance(1), 5);
 			assert_eq!(Balances::reserved_balance(1), 5);
 			assert_eq!(Balances::usable_balance(&1), 2);
 
@@ -75,13 +75,16 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 use sp_std::{prelude::*, marker::PhantomData};
-use frame_support::weights::{GetDispatchInfo, WeighBlock, DispatchInfo};
+use frame_support::{
+	storage::StorageValue,
+	weights::{GetDispatchInfo, WeighBlock, DispatchInfo}
+};
 use sp_runtime::{
 	generic::Digest,
 	ApplyExtrinsicResult,
 	traits::{
 		self, Header, Zero, One, Checkable, Applyable, CheckEqual, OnFinalize, OnInitialize,
-		NumberFor, Block as BlockT, OffchainWorker, Dispatchable, Saturating,
+		NumberFor, Block as BlockT, OffchainWorker, Dispatchable, Saturating, OnRuntimeUpgrade,
 	},
 	transaction_validity::TransactionValidity,
 };
@@ -110,6 +113,7 @@ impl<
 	Context: Default,
 	UnsignedValidator,
 	AllModules:
+		OnRuntimeUpgrade +
 		OnInitialize<System::BlockNumber> +
 		OnFinalize<System::BlockNumber> +
 		OffchainWorker<System::BlockNumber> +
@@ -135,6 +139,7 @@ impl<
 	Context: Default,
 	UnsignedValidator,
 	AllModules:
+		OnRuntimeUpgrade +
 		OnInitialize<System::BlockNumber> +
 		OnFinalize<System::BlockNumber> +
 		OffchainWorker<System::BlockNumber> +
@@ -176,6 +181,12 @@ where
 		extrinsics_root: &System::Hash,
 		digest: &Digest<System::Hash>,
 	) {
+		if frame_system::RuntimeUpgraded::take() {
+			<AllModules as OnRuntimeUpgrade>::on_runtime_upgrade();
+			<frame_system::Module<System>>::register_extra_weight_unchecked(
+				<AllModules as WeighBlock<System::BlockNumber>>::on_runtime_upgrade()
+			);
+		}
 		<frame_system::Module<System>>::initialize(
 			block_number,
 			parent_hash,
@@ -573,7 +584,7 @@ mod tests {
 				header: Header {
 					parent_hash: [69u8; 32].into(),
 					number: 1,
-					state_root: hex!("17caebd966d10cc6dc9659edf7fa3196511593f6c39f80f9b97cdbc3b0855cf3").into(),
+					state_root: hex!("8a22606e925c39bb0c8e8f6f5871c0aceab88a2fcff6b2d92660af8f6daff0b1").into(),
 					extrinsics_root: hex!("03170a2e7597b7b7e3d84c05391d139a62b157e78786d8c082f29dcf4c111314").into(),
 					digest: Digest { logs: vec![], },
 				},