Introduce on_runtime_upgrade

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/executive/src/lib.rs

@@ -81,7 +81,7 @@ use sp_runtime::{
 	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 +110,7 @@ impl<
 	Context: Default,
 	UnsignedValidator,
 	AllModules:
+		OnRuntimeUpgrade +
 		OnInitialize<System::BlockNumber> +
 		OnFinalize<System::BlockNumber> +
 		OffchainWorker<System::BlockNumber> +
@@ -135,6 +136,7 @@ impl<
 	Context: Default,
 	UnsignedValidator,
 	AllModules:
+		OnRuntimeUpgrade +
 		OnInitialize<System::BlockNumber> +
 		OnFinalize<System::BlockNumber> +
 		OffchainWorker<System::BlockNumber> +
@@ -176,13 +178,21 @@ where
 		extrinsics_root: &System::Hash,
 		digest: &Digest<System::Hash>,
 	) {
+		<AllModules as OnRuntimeUpgrade>::on_runtime_upgrade();
 		<frame_system::Module<System>>::initialize(
 			block_number,
 			parent_hash,
 			extrinsics_root,
 			digest,
 			frame_system::InitKind::Full,
 		);
+		let last_runtime_upgrade = <frame_system::Module<System>>::last_runtime_upgrade();
+		if last_runtime_upgrade.map(|n| n == *block_number).unwrap_or(false) {
+			<AllModules as OnRuntimeUpgrade>::on_runtime_upgrade();
+			<frame_system::Module<System>>::register_extra_weight_unchecked(
+				<AllModules as WeighBlock<System::BlockNumber>>::on_runtime_upgrade()
+			);
+		}
 		<AllModules as OnInitialize<System::BlockNumber>>::on_initialize(*block_number);
 		<frame_system::Module<System>>::register_extra_weight_unchecked(
 			<AllModules as WeighBlock<System::BlockNumber>>::on_initialize(*block_number)

frame/staking/src/lib.rs

@@ -927,10 +927,13 @@ decl_module! {
 
 		fn deposit_event() = default;
 
-		fn on_initialize() {
+		fn on_runtime_upgrade() {
 			Self::ensure_storage_upgraded();
 		}
 
+		fn on_initialize() {
+		}
+
 		fn on_finalize() {
 			// Set the start of the first era.
 			if let Some(mut active_era) = Self::active_era() {