Introduce on_runtime_upgrade

bin/node/runtime/src/lib.rs

@@ -82,7 +82,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	// and set impl_version to 0. If only runtime
 	// implementation changes and behavior does not, then leave spec_version as
 	// is and increment impl_version.
-	spec_version: 227,
+	spec_version: 228,
 	impl_version: 0,
 	apis: RUNTIME_API_VERSIONS,
 };

frame/balances/src/lib.rs

@@ -156,6 +156,7 @@ mod tests_composite;
 #[macro_use]
 mod tests;
 mod benchmarking;
+mod migration;
 
 use sp_std::prelude::*;
 use sp_std::{cmp, result, mem, fmt::Debug, ops::BitOr, convert::Infallible};
@@ -366,10 +367,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.
+		/// Module version number.
 		///
-		/// True for new networks.
-		IsUpgraded build(|_: &GenesisConfig<T, I>| true): bool;
+		/// Default to version 1 for new networks.
+		ModuleVersion build(|_: &GenesisConfig<T, I>| 1): u8;
 	}
 	add_extra_genesis {
 		config(balances): Vec<(T::AccountId, T::Balance)>;
@@ -517,11 +518,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>();
 		}
 	}
 }
@@ -547,82 +545,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

frame/balances/src/migration.rs

@@ -0,0 +1,87 @@
+use super::*;
+
+// Upgrade from the pre-#4649 balances/vesting into the new balances.
+pub fn on_runtime_upgrade<T: Trait<I>, I: Instance>() {
+	let current_version = ModuleVersion::<I>::get();
+
+	if current_version == 1 {
+		return
+	} else if current_version == 0 {
+		ModuleVersion::<I>::put(1);
+		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));
+		}
+
+		// Recursively call the upgrade function to apply all upgrades.
+		on_runtime_upgrade::<T, I>()
+	}
+}

frame/democracy/src/lib.rs

@@ -1249,7 +1249,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};
@@ -1407,7 +1407,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);

frame/executive/src/lib.rs

@@ -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!("c96987b52003684f590a6b9b55af9ded88b2e6f8d84441c456f46e59f5e73070").into(),
 					extrinsics_root: hex!("03170a2e7597b7b7e3d84c05391d139a62b157e78786d8c082f29dcf4c111314").into(),
 					digest: Digest { logs: vec![], },
 				},

frame/staking/src/lib.rs

@@ -927,7 +927,7 @@ decl_module! {
 
 		fn deposit_event() = default;
 
-		fn on_initialize() {
+		fn on_runtime_upgrade() {
 			Self::ensure_storage_upgraded();
 		}
 
@@ -1132,8 +1132,6 @@ decl_module! {
 		/// # </weight>
 		#[weight = SimpleDispatchInfo::FixedNormal(750_000)]
 		fn validate(origin, prefs: ValidatorPrefs) {
-			Self::ensure_storage_upgraded();
-
 			let controller = ensure_signed(origin)?;
 			let ledger = Self::ledger(&controller).ok_or(Error::<T>::NotController)?;
 			let stash = &ledger.stash;
@@ -1154,8 +1152,6 @@ decl_module! {
 		/// # </weight>
 		#[weight = SimpleDispatchInfo::FixedNormal(750_000)]
 		fn nominate(origin, targets: Vec<<T::Lookup as StaticLookup>::Source>) {
-			Self::ensure_storage_upgraded();
-
 			let controller = ensure_signed(origin)?;
 			let ledger = Self::ledger(&controller).ok_or(Error::<T>::NotController)?;
 			let stash = &ledger.stash;
@@ -1433,7 +1429,6 @@ decl_module! {
 		///
 		/// - `stash`: The stash account to reap. Its balance must be zero.
 		fn reap_stash(_origin, stash: T::AccountId) {
-			Self::ensure_storage_upgraded();
 			ensure!(T::Currency::total_balance(&stash).is_zero(), Error::<T>::FundedTarget);
 			Self::kill_stash(&stash)?;
 			T::Currency::remove_lock(STAKING_ID, &stash);
@@ -2108,7 +2103,6 @@ impl<T: Trait> Module<T> {
 /// some session can lag in between the newest session planned and the latest session started.
 impl<T: Trait> pallet_session::SessionManager<T::AccountId> for Module<T> {
 	fn new_session(new_index: SessionIndex) -> Option<Vec<T::AccountId>> {
-		Self::ensure_storage_upgraded();
 		Self::new_session(new_index)
 	}
 	fn start_session(start_index: SessionIndex) {