Introduce on_runtime_upgrade (paritytech#5058)

* Initial idea of `on_runtime_upgrade` * Runtime storage for module version * Gui shawntabrizi runtime upgrade (paritytech#5118) * adding unleash to ci (paritytech#5020) * adding unleash to ci * fixing formatting * with a dot please * alpha.3 now * do not publish testing helpers * remove old test-helpers cruft * fix cargo.lock * with alpha 4 * do not publish runtime-interface-test either * disable more test crates from publishing * switch to alpha.5 * replace tempdir with tempfile * update lru * switch to bytes 0.5 * release script fixes * switch on and to latest alpha * BUT THE SPACES * Fix: CI failing for some CLI tests (paritytech#5043) * Initial commit Forked at: 41bb219 Parent branch: origin/master * Increase killing grace period of CLI tests and display more info * Use --dev everywhere possible * Put pruning mode to its own params struct * Add pruning params to export-blocks command * Added missing file * Removed not-dev mode in tests * Add pruning mode to the revert command * Decrease killing grace period again * Move back unsafe_pruning to import_params * Applied proposed changes * aura: remove unused tx pool (paritytech#5046) * aura: remove unused transaction pool parameter * node-template: remove transaction pool from aura * aura: fix tests * Extend rust nightly detection in `wasm-builder` (paritytech#5021) Instead of just testing `cargo` and `rustup run nightly`, we now test the `CARGO` env variable and also scan non default nightlies. The user is also now able to select the toolchain with `WASM_BUILD_TOOLCHAIN`. * Add steps setting to benchmarking CLI (paritytech#5033) * Add steps setting to CLI, use max value to hit worst case. * Bump impl_version. * Apply review suggestion. * Remove indices from node-template (paritytech#5025) * Remove indices from node-template * Use identity lookup instead * Bump impl * clean cargo.toml * Fix documentation for "BlockBuilder::push_trusted" (paritytech#5051) * fix doc * rephrase * do not check unleash on every PR, only master and tags (paritytech#5054) * do not check unleash on every PR, only master and tags * move scripts folder * add signed-tag check to CI * remove publish-to-crates-io dependencies Co-authored-by: s3krit <[email protected]> * prepare version to alpha.1 (paritytech#5055) bump version to -alpha.1 * Sync: validate block responses for required data (paritytech#5052) * Less verbose state-db logging * Validate block responses for block bodies * Update client/network/src/protocol.rs Co-Authored-By: Bastian Köcher <[email protected]> * Added validation test * Disconnect on missing header as well * Typo Co-Authored-By: André Silva <[email protected]> Co-authored-by: Bastian Köcher <[email protected]> Co-authored-by: André Silva <[email protected]> * Make these chainspecs fields private (paritytech#5031) * Fix dockerfile (paritytech#5059) * Adds documentation for `wipe` and `commit` (paritytech#5053) * Adds documentation for `wipe` and `commit` This adds documentation to `wipe` and `commit` of `Externalities`. Besides that it removes the default implementation that would just panic and requires that all implementers of the trait implement the functions. * Update primitives/externalities/src/lib.rs Co-Authored-By: joe petrowski <[email protected]> Co-authored-by: joe petrowski <[email protected]> * Fix the issue with `trybuild`'s `ui` tests (paritytech#4992) * repro ui bug * fix the tests * test with the new image * test without CARGO_HOME * test without fixes * test again * fix trybuild old versions * bump CArgo.lock * fix trybuild newest versions * bump Cargo.lock * trying on the latest image * bump Cargo.lock * run with the old image * ci will be green on the image from 2020-02-19 [skip ci] * bump Cargo.lock * Activate publishing of draft releases... (paritytech#5062) * Activate publishing of draft releases... ... And fix the message sending (missing parameter). * publish_draft_release.sh now checks latest... ... release on github rather than just a tag * Fix/div by zero (paritytech#5041) * Handle gas_price being zero separately * Bump spec_version * Add a unit & integration tests for gas price = 0 * set missing metadata fields, prepping alpha.2 (paritytech#5067) * setting first batch of descriptions * fix what I just broke * next batch * and pallets, too * last batch * set cargo.lock * keep'em dev-deps * bump version to alpha.2 * Fix revalidation not revalidating multiple times (paritytech#5065) * removes use of sc_client::Client from sc_finality_grandpa (paritytech#5030) * removes use of sc_client::Client from sc_finality_grandpa * code formatting * code formatting * removes use of sc_client::Client from sc_finality_grandpa * Remove deprecated host functions (paritytech#5038) Sadly we need to keep one function `ext_blake2_256`. This function is manually defined in `sp-core`. * removes use of sc_client::Client from sc_basic_authorship (paritytech#5050) * removes use of sc-client from sc-basic-authorship * refactor use of ProposerFactory * correct dep path * pallet-transaction-payment clean up (paritytech#5070) * Formatting clean up * Introduce separate setters for the fees. * *: Rename prometheus-exporter crate to substrate-prometheus-end… (paritytech#5076) This patch renames the crate for the following two reasons: 1. The prometheus-exporter crate introduces native in-process Prometheus style instrumentation to the Substrate project. Within the Prometheus ecosystem the term "exporter" is used for external processes exposing metrics for e.g. the Linux Kernel. In-process exposition would be described via the term "endpoint". 2. "prometheus-exporter" is generic and ignores the fact that it is only usable within the context of Substrate. In addition the name "prometheus-exporter" is already taken on crates.io. * rename `browser-utils` to `substrate-browser-utils` (paritytech#5079) * prepping for Alpha.3 (paritytech#5080) * Bump to alpha.3 * update gitlab-ci * Propagate DispatchError for benchmarks. (paritytech#5075) * Propagate DispatchError for benchmarks. * Apply review suggestions. * Use RuntimeString. * fix expect Co-Authored-By: Bastian Köcher <[email protected]> Co-authored-by: Bastian Köcher <[email protected]> * Add options to overwrite range bounds in benchmark command. (paritytech#5072) * Add --mins --maxs to benchmark command. * Apply review suggestions. * Update yamux to version 0.4.4. (paritytech#5086) * Remove more instances of futures01 (paritytech#4633) * Start removing last few instances of futures01 * Use to_poll on wasm * Revert "Use to_poll on wasm" This reverts commit 1c61728. * Fix fg test * Upgrade network test futures * Update offchain hyper version * Update service test * bump tokio to 0.2.10 * Removed some unneeded tokios * fixes * fix run_until_all_full * Make service test debuggable * Update client/offchain/src/api/http.rs Co-Authored-By: Demi Obenour <[email protected]> * Add service_test to test-int output * nitpicking * Finally fix test * Give up and revert client/serviec/test * Revert gitlab ci too Co-authored-by: Demi Obenour <[email protected]> * Make export blocks default to json on stdout (paritytech#5090) * Make export blocks default to json on stdout * Multiline instead of single line to stay under 100 cols * Change --json flag to --binary, defaulting to json * Offence reporting returns a result (paritytech#5082) * Offence reporting returns a result * Bump spec_version * Use unwrap instead of assertions * Fix more review grumbles * Update to libp2p 0.16.2 (paritytech#5088) * Remove request ID from the new protocol (paritytech#5049) * Make sure we remove a peer on disconnect in gossip (paritytech#5104) * Make sure we remove peers on disconnect in gossip state machine * Clear up the code * Add a comment * Expose `state-db` memory info (paritytech#5110) This exposes memory statistics from the state-db. * Change extrinsic_count to extrinsic_index in pallet-utility (paritytech#5044) Co-authored-by: Benjamin Kampmann <[email protected]> * Add more metrics to prometheus (paritytech#5034) * Add a few things * Add finality_grandpa_round * fix fg tests * Nitpicks * Nitpicks * Fix name of prometheus crate * Update to SCALE 1.2.0 (paritytech#5113) This updates `parity-scale-codec` to `1.2.0`, which includes multiple performance improvements and a fix that bounds the capacity of a vector at decoding. * Lazy payouts (paritytech#4474) * TODOs * Remove superfluous: * partial implementation * full implementation * fix preferences * update comments * upgrade test WIP * fix more tests * fix cutoff * fix saturation * comment * upgrade mock * upgrade test * WIP migration * WIP migration * remove slot stake stuff * fix merge * migration of ledger * remove equalize from test * add test * fix * update doc * fix compilation * improve test readibility * improve doc * fix most todo * fix migration and test * remove println * WIP * add test and spec * weight * update doc * safer end_era * fix exposure of conversion * Revert "safer end_era" This reverts commit 72ff737. * fix useless put * exposure clipped * doc * fix payout with clipped * fix node runtime * add doc * pluggable and generalized staking module * remove print * update doc * refactor * improve documentation and implementation * fix test * Fix test * fix test * fix test * fix remove lowest stake from exposure, not biggest. * nomination index arguments in nominator_payout * add test * try to fix offence * apply slashed and bond eras until active era * doc * update spec version * add test upgrade from previous test environment * Apply suggestions from code review Co-Authored-By: Shawn Tabrizi <[email protected]> * nominators upgrade has been cleaned * dynamic history depth implementation * make current_era - history_depth included * Change equality check to start era to less than or equal * Use era specific validator prefs * Add print statement and comment about start era if < * fix next_reward overflow * make more check for bad era claim for zero cost * small refactor * code refactor + fix use of deprecated storage * fix wasm build * add comment * Fix tests * remove outdated comment * Apply suggestions from code review Co-Authored-By: Shawn Tabrizi <[email protected]> * gather active era information into one storage Co-authored-by: thiolliere <[email protected]> Co-authored-by: Shawn Tabrizi <[email protected]> * impl on_runtime_upgrade Co-authored-by: Benjamin Kampmann <[email protected]> Co-authored-by: Cecile Tonglet <[email protected]> Co-authored-by: André Silva <[email protected]> Co-authored-by: Bastian Köcher <[email protected]> Co-authored-by: Marcio Diaz <[email protected]> Co-authored-by: Nikolay Volf <[email protected]> Co-authored-by: s3krit <[email protected]> Co-authored-by: Arkadiy Paronyan <[email protected]> Co-authored-by: Pierre Krieger <[email protected]> Co-authored-by: Chevdor <[email protected]> Co-authored-by: joe petrowski <[email protected]> Co-authored-by: Denis Pisarev <[email protected]> Co-authored-by: Eric <[email protected]> Co-authored-by: Seun Lanlege <[email protected]> Co-authored-by: Sergei Pepyakin <[email protected]> Co-authored-by: Max Inden <[email protected]> Co-authored-by: Ashley <[email protected]> Co-authored-by: Toralf Wittner <[email protected]> Co-authored-by: Demi Obenour <[email protected]> Co-authored-by: pscott <[email protected]> Co-authored-by: Fedor Sakharov <[email protected]> Co-authored-by: Gavin Wood <[email protected]> Co-authored-by: thiolliere <[email protected]> * make compile * Add some tests * docs * Remove "useless" code * Fix merge and use n + 1 block number * Fix tests * unfix ui tests * Update on_initialize.stderr * fix test * Fix test * Bump spec * Remove `on_finalise` and `on_initialise` * Use bool for tracking runtime upgraded * typo * Support runtime upgrade with `set_storage` * Refactor migration code location * add trailing newlines * Remove old `IsUpgraded` flag * Update state root * Exhaustive match statement * Apply suggestions from code review Co-Authored-By: Kian Paimani <[email protected]> Co-authored-by: Benjamin Kampmann <[email protected]> Co-authored-by: Cecile Tonglet <[email protected]> Co-authored-by: André Silva <[email protected]> Co-authored-by: Bastian Köcher <[email protected]> Co-authored-by: Marcio Diaz <[email protected]> Co-authored-by: Nikolay Volf <[email protected]> Co-authored-by: s3krit <[email protected]> Co-authored-by: Arkadiy Paronyan <[email protected]> Co-authored-by: Pierre Krieger <[email protected]> Co-authored-by: Chevdor <[email protected]> Co-authored-by: joe petrowski <[email protected]> Co-authored-by: Denis Pisarev <[email protected]> Co-authored-by: Eric <[email protected]> Co-authored-by: Seun Lanlege <[email protected]> Co-authored-by: Sergei Pepyakin <[email protected]> Co-authored-by: Max Inden <[email protected]> Co-authored-by: Ashley <[email protected]> Co-authored-by: Toralf Wittner <[email protected]> Co-authored-by: Demi Obenour <[email protected]> Co-authored-by: pscott <[email protected]> Co-authored-by: Fedor Sakharov <[email protected]> Co-authored-by: Gavin Wood <[email protected]> Co-authored-by: thiolliere <[email protected]> Co-authored-by: Kian Paimani <[email protected]>
General-Beck · Mar 6, 2020 · 211b2fc · 211b2fc
1 parent 3179d01
commit 211b2fc
Show file tree

Hide file tree

Showing 21 changed files with 811 additions and 513 deletions.
diff --git a/frame/balances/src/lib.rs b/frame/balances/src/lib.rs
@@ -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

diff --git a/frame/balances/src/migration.rs b/frame/balances/src/migration.rs
@@ -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);
+}
diff --git a/frame/democracy/src/lib.rs b/frame/democracy/src/lib.rs
@@ -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);

diff --git a/frame/executive/src/lib.rs b/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!("8a22606e925c39bb0c8e8f6f5871c0aceab88a2fcff6b2d92660af8f6daff0b1").into(),
 					extrinsics_root: hex!("03170a2e7597b7b7e3d84c05391d139a62b157e78786d8c082f29dcf4c111314").into(),
 					digest: Digest { logs: vec![], },
 				},