migrations: VersionedRuntimeUpgrade (#14311)

* VersionedRuntimeUpgrade

* only require one version and add a pre-upgrade check

* add docs

* improve warning log

* improve comments

* fix log

* use associated constants

* allow passing from and to versions

* test versioned_runtime_upgrade

* fix typo

* improve docs

* docs

* docs

* remove event from dummy pallet

* remove pre_upgrade current storage version check

* derive_impl

* skip pre/post checks if the actual migration will not run

* improve variable naming

* docs

* fix post_upgrade 'should run' logic

* fix comments

* pre and post hook tests

* feature gate try-runtime stuff

* remove deprecated macro

* Update frame/support/src/migrations.rs

Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

* decode_all

* make experimental

* use rust generics

* add info log when running

* simplify tests

* improve log

* improve log

* cleaner pre_upgrade encoding

* Update frame/support/src/migrations.rs

Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>

* Update frame/support/src/migrations.rs

Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>

* Update frame/support/src/migrations.rs

Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>

* Update frame/support/src/migrations.rs

Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>

* Update frame/support/src/migrations.rs

Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>

* VersionedPostUpgradeData enum

* move versioned runtime upgrade tests to test/tests

* fix rust doc

* clarify comment

---------

Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>

substrate/frame/support/Cargo.toml

@@ -79,6 +79,7 @@ runtime-benchmarks = [
 try-runtime = [
 	"sp-debug-derive/force-debug"
 ]
+experimental = []
 # By default some types have documentation, `no-metadata-docs` allows to reduce the documentation
 # in the metadata.
 no-metadata-docs = ["frame-support-procedural/no-metadata-docs", "sp-api/no-metadata-docs"]

substrate/frame/support/src/migrations.rs

@@ -15,8 +15,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#[cfg(feature = "try-runtime")]
-use crate::storage::unhashed::contains_prefixed_key;
 use crate::{
 	traits::{GetStorageVersion, NoStorageVersionSet, PalletInfoAccess, StorageVersion},
 	weights::{RuntimeDbWeight, Weight},
@@ -25,9 +23,148 @@ use impl_trait_for_tuples::impl_for_tuples;
 use sp_core::Get;
 use sp_io::{hashing::twox_128, storage::clear_prefix, KillStorageResult};
 use sp_std::marker::PhantomData;
+
 #[cfg(feature = "try-runtime")]
 use sp_std::vec::Vec;
 
+#[cfg(feature = "experimental")]
+use crate::traits::OnRuntimeUpgrade;
+
+/// Make it easier to write versioned runtime upgrades.
+///
+/// [`VersionedRuntimeUpgrade`] allows developers to write migrations without worrying about
+/// checking and setting storage versions. Instead, the developer wraps their migration in this
+/// struct which takes care of version handling using best practices.
+///
+/// It takes 5 type parameters:
+/// - `From`: The version being upgraded from.
+/// - `To`: The version being upgraded to.
+/// - `Inner`: An implementation of `OnRuntimeUpgrade`.
+/// - `Pallet`: The Pallet being upgraded.
+/// - `Weight`: The runtime's RuntimeDbWeight implementation.
+///
+/// When a [`VersionedRuntimeUpgrade`] `on_runtime_upgrade`, `pre_upgrade`, or `post_upgrade`
+/// method is called, the on-chain version of the pallet is compared to `From`. If they match, the
+/// `Inner` equivalent is called and the pallets on-chain version is set to `To` after the
+/// migration. Otherwise, a warning is logged notifying the developer that the upgrade was a noop
+/// and should probably be removed.
+///
+/// ### Examples
+/// ```ignore
+/// // In file defining migrations
+/// pub struct VersionUncheckedMigrateV5ToV6<T>(sp_std::marker::PhantomData<T>);
+/// impl<T: Config> OnRuntimeUpgrade for VersionUncheckedMigrateV5ToV6<T> {
+/// 	// OnRuntimeUpgrade implementation...
+/// }
+///
+/// pub type VersionCheckedMigrateV5ToV6<Runtime, Pallet, DbWeight> =
+/// 	VersionedRuntimeUpgrade<5, 6, VersionUncheckedMigrateV5ToV6<Runtime>, Pallet, DbWeight>;
+///
+/// // Migrations tuple to pass to the Executive pallet:
+/// pub type Migrations = (
+/// 	// other migrations...
+/// 	VersionCheckedMigrateV5ToV6<Runtime, Balances, RuntimeDbWeight>,
+/// 	// other migrations...
+/// );
+/// ```
+#[cfg(feature = "experimental")]
+pub struct VersionedRuntimeUpgrade<const FROM: u16, const TO: u16, Inner, Pallet, Weight> {
+	_marker: PhantomData<(Inner, Pallet, Weight)>,
+}
+
+/// A helper enum to wrap the pre_upgrade bytes like an Option before passing them to post_upgrade.
+/// This enum is used rather than an Option to make the API clearer to the developer.
+#[cfg(feature = "experimental")]
+#[derive(codec::Encode, codec::Decode)]
+pub enum VersionedPostUpgradeData {
+	/// The migration ran, inner vec contains pre_upgrade data.
+	MigrationExecuted(Vec<u8>),
+	/// This migration is a noop, do not run post_upgrade checks.
+	Noop,
+}
+
+/// Implementation of the `OnRuntimeUpgrade` trait for `VersionedRuntimeUpgrade`.
+///
+/// Its main function is to perform the runtime upgrade in `on_runtime_upgrade` only if the on-chain
+/// version of the pallets storage matches `From`, and after the upgrade set the on-chain storage to
+/// `To`. If the versions do not match, it writes a log notifying the developer that the migration
+/// is a noop.
+#[cfg(feature = "experimental")]
+impl<
+		const FROM: u16,
+		const TO: u16,
+		Inner: OnRuntimeUpgrade,
+		Pallet: GetStorageVersion<CurrentStorageVersion = StorageVersion> + PalletInfoAccess,
+		DbWeight: Get<RuntimeDbWeight>,
+	> OnRuntimeUpgrade for VersionedRuntimeUpgrade<FROM, TO, Inner, Pallet, DbWeight>
+{
+	/// Executes pre_upgrade if the migration will run, and wraps the pre_upgrade bytes in
+	/// [`VersionedPostUpgradeData`] before passing them to post_upgrade, so it knows whether the
+	/// migration ran or not.
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
+		use codec::Encode;
+		let on_chain_version = Pallet::on_chain_storage_version();
+		if on_chain_version == FROM {
+			Ok(VersionedPostUpgradeData::MigrationExecuted(Inner::pre_upgrade()?).encode())
+		} else {
+			Ok(VersionedPostUpgradeData::Noop.encode())
+		}
+	}
+
+	/// Executes the versioned runtime upgrade.
+	///
+	/// First checks if the pallets on-chain storage version matches the version of this upgrade. If
+	/// it matches, it calls `Inner::on_runtime_upgrade`, updates the on-chain version, and returns
+	/// the weight. If it does not match, it writes a log notifying the developer that the migration
+	/// is a noop.
+	fn on_runtime_upgrade() -> Weight {
+		let on_chain_version = Pallet::on_chain_storage_version();
+		if on_chain_version == FROM {
+			log::info!(
+				"Running {} VersionedOnRuntimeUpgrade: version {:?} to {:?}.",
+				Pallet::name(),
+				FROM,
+				TO
+			);
+
+			// Execute the migration
+			let weight = Inner::on_runtime_upgrade();
+
+			// Update the on-chain version
+			StorageVersion::new(TO).put::<Pallet>();
+
+			weight.saturating_add(DbWeight::get().reads_writes(1, 1))
+		} else {
+			log::warn!(
+				"{} VersionedOnRuntimeUpgrade for version {:?} skipped because current on-chain version is {:?}.",
+				Pallet::name(),
+				FROM,
+				on_chain_version
+			);
+			DbWeight::get().reads(1)
+		}
+	}
+
+	/// Executes `Inner::post_upgrade` if the migration just ran.
+	///
+	/// pre_upgrade passes [`VersionedPostUpgradeData::MigrationExecuted`] to post_upgrade if
+	/// the migration ran, and [`VersionedPostUpgradeData::Noop`] otherwise.
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(
+		versioned_post_upgrade_data_bytes: Vec<u8>,
+	) -> Result<(), sp_runtime::TryRuntimeError> {
+		use codec::DecodeAll;
+		match <VersionedPostUpgradeData>::decode_all(&mut &versioned_post_upgrade_data_bytes[..])
+			.map_err(|_| "VersionedRuntimeUpgrade post_upgrade failed to decode PreUpgradeData")?
+		{
+			VersionedPostUpgradeData::MigrationExecuted(inner_bytes) =>
+				Inner::post_upgrade(inner_bytes),
+			VersionedPostUpgradeData::Noop => Ok(()),
+		}
+	}
+}
+
 /// Can store the current pallet version in storage.
 pub trait StoreCurrentStorageVersion<T: GetStorageVersion + PalletInfoAccess> {
 	/// Write the current storage version to the storage.
@@ -185,6 +322,8 @@ impl<P: Get<&'static str>, DbWeight: Get<RuntimeDbWeight>> frame_support::traits
 
 	#[cfg(feature = "try-runtime")]
 	fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
+		use crate::storage::unhashed::contains_prefixed_key;
+
 		let hashed_prefix = twox_128(P::get().as_bytes());
 		match contains_prefixed_key(&hashed_prefix) {
 			true => log::info!("Found {} keys pre-removal 👀", P::get()),
@@ -198,6 +337,8 @@ impl<P: Get<&'static str>, DbWeight: Get<RuntimeDbWeight>> frame_support::traits
 
 	#[cfg(feature = "try-runtime")]
 	fn post_upgrade(_state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
+		use crate::storage::unhashed::contains_prefixed_key;
+
 		let hashed_prefix = twox_128(P::get().as_bytes());
 		match contains_prefixed_key(&hashed_prefix) {
 			true => {

substrate/frame/support/test/Cargo.toml

@@ -54,6 +54,7 @@ std = [
 	"sp-version/std",
 	"test-pallet/std",
 ]
+experimental = ["frame-support/experimental"]
 try-runtime = [
 	"frame-support/try-runtime",
 	"frame-system/try-runtime",

substrate/frame/support/test/tests/versioned_runtime_upgrade.rs

@@ -0,0 +1,227 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Tests for VersionedRuntimeUpgrade
+
+#![cfg(all(feature = "experimental", feature = "try-runtime"))]
+
+use frame_support::{
+	construct_runtime, derive_impl,
+	migrations::VersionedRuntimeUpgrade,
+	parameter_types,
+	traits::{GetStorageVersion, OnRuntimeUpgrade, StorageVersion},
+	weights::constants::RocksDbWeight,
+};
+use frame_system::Config;
+
+type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic<Test>;
+type Block = frame_system::mocking::MockBlock<Test>;
+
+#[frame_support::pallet]
+mod dummy_pallet {
+	use frame_support::pallet_prelude::*;
+
+	const STORAGE_VERSION: StorageVersion = StorageVersion::new(4);
+
+	#[pallet::pallet]
+	#[pallet::storage_version(STORAGE_VERSION)]
+	pub struct Pallet<T>(_);
+
+	#[pallet::config]
+	pub trait Config: frame_system::Config {}
+
+	#[pallet::storage]
+	pub type SomeStorage<T: Config> = StorageValue<_, u32, ValueQuery>;
+
+	#[pallet::genesis_config]
+	#[derive(Default)]
+	pub struct GenesisConfig {}
+
+	#[pallet::genesis_build]
+	impl<T: Config> GenesisBuild<T> for GenesisConfig {
+		fn build(&self) {}
+	}
+}
+
+impl dummy_pallet::Config for Test {}
+
+construct_runtime!(
+	pub enum Test where
+		Block = Block,
+		NodeBlock = Block,
+		UncheckedExtrinsic = UncheckedExtrinsic,
+	{
+		System: frame_system::{Pallet, Call, Config, Storage, Event<T>} = 0,
+		DummyPallet: dummy_pallet::{Pallet, Config, Storage} = 1,
+	}
+);
+
+#[derive_impl(frame_system::config_preludes::TestDefaultConfig as frame_system::DefaultConfig)]
+impl frame_system::Config for Test {
+	type BaseCallFilter = frame_support::traits::Everything;
+	type RuntimeOrigin = RuntimeOrigin;
+	type RuntimeCall = RuntimeCall;
+	type RuntimeEvent = RuntimeEvent;
+	type PalletInfo = PalletInfo;
+	type OnSetCode = ();
+}
+
+pub(crate) fn new_test_ext() -> sp_io::TestExternalities {
+	let storage = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
+	let mut ext: sp_io::TestExternalities = sp_io::TestExternalities::from(storage);
+	ext.execute_with(|| System::set_block_number(1));
+	ext
+}
+
+/// A dummy migration for testing the `VersionedRuntimeUpgrade` trait.
+/// Sets SomeStorage to S.
+struct SomeUnversionedMigration<T: Config, const S: u32>(sp_std::marker::PhantomData<T>);
+
+parameter_types! {
+	const UpgradeReads: u64 = 4;
+	const UpgradeWrites: u64 = 2;
+	const PreUpgradeReturnBytes: [u8; 4] = [0, 1, 2, 3];
+	static PreUpgradeCalled: bool = false;
+	static PostUpgradeCalled: bool = false;
+	static PostUpgradeCalledWith: Vec<u8> = Vec::new();
+}
+
+/// Implement `OnRuntimeUpgrade` for `SomeUnversionedMigration`.
+/// It sets SomeStorage to S, and returns a weight derived from UpgradeReads and UpgradeWrites.
+impl<T: dummy_pallet::Config, const S: u32> OnRuntimeUpgrade for SomeUnversionedMigration<T, S> {
+	fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
+		PreUpgradeCalled::set(true);
+		Ok(PreUpgradeReturnBytes::get().to_vec())
+	}
+
+	fn on_runtime_upgrade() -> frame_support::weights::Weight {
+		dummy_pallet::SomeStorage::<T>::put(S);
+		RocksDbWeight::get().reads_writes(UpgradeReads::get(), UpgradeWrites::get())
+	}
+
+	fn post_upgrade(state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
+		PostUpgradeCalled::set(true);
+		PostUpgradeCalledWith::set(state);
+		Ok(())
+	}
+}
+
+type VersionedMigrationV0ToV1 =
+	VersionedRuntimeUpgrade<0, 1, SomeUnversionedMigration<Test, 1>, DummyPallet, RocksDbWeight>;
+
+type VersionedMigrationV1ToV2 =
+	VersionedRuntimeUpgrade<1, 2, SomeUnversionedMigration<Test, 2>, DummyPallet, RocksDbWeight>;
+
+type VersionedMigrationV2ToV4 =
+	VersionedRuntimeUpgrade<2, 4, SomeUnversionedMigration<Test, 4>, DummyPallet, RocksDbWeight>;
+
+#[test]
+fn successful_upgrade_path() {
+	new_test_ext().execute_with(|| {
+		// on-chain storage version and value in storage start at zero
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(0));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 0);
+
+		// Execute the migration from version 0 to 1 and verify it was successful
+		VersionedMigrationV0ToV1::on_runtime_upgrade();
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(1));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 1);
+
+		// Execute the migration from version 1 to 2 and verify it was successful
+		VersionedMigrationV1ToV2::on_runtime_upgrade();
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(2));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 2);
+
+		// Execute the migration from version 2 to 4 and verify it was successful
+		VersionedMigrationV2ToV4::on_runtime_upgrade();
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(4));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 4);
+	});
+}
+
+#[test]
+fn future_version_upgrade_is_ignored() {
+	new_test_ext().execute_with(|| {
+		// Executing V1 to V2 on V0 should be a noop
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(0));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 0);
+		VersionedMigrationV1ToV2::on_runtime_upgrade();
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(0));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 0);
+	});
+}
+
+#[test]
+fn past_version_upgrade_is_ignored() {
+	new_test_ext().execute_with(|| {
+		// Upgrade to V2
+		VersionedMigrationV0ToV1::on_runtime_upgrade();
+		VersionedMigrationV1ToV2::on_runtime_upgrade();
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(2));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 2);
+
+		// Now, V0 to V1 and V1 to V2 should both be noops
+		dummy_pallet::SomeStorage::<Test>::put(1000);
+		VersionedMigrationV0ToV1::on_runtime_upgrade();
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(2));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 1000);
+		VersionedMigrationV1ToV2::on_runtime_upgrade();
+		assert_eq!(DummyPallet::on_chain_storage_version(), StorageVersion::new(2));
+		assert_eq!(dummy_pallet::SomeStorage::<Test>::get(), 1000);
+	});
+}
+
+#[test]
+fn weights_are_returned_correctly() {
+	new_test_ext().execute_with(|| {
+		// Successful upgrade requires 1 additional read and write
+		let weight = VersionedMigrationV0ToV1::on_runtime_upgrade();
+		assert_eq!(
+			weight,
+			RocksDbWeight::get().reads_writes(UpgradeReads::get() + 1, UpgradeWrites::get() + 1)
+		);
+
+		// Noop upgrade requires only 1 read
+		let weight = VersionedMigrationV0ToV1::on_runtime_upgrade();
+		assert_eq!(weight, RocksDbWeight::get().reads(1));
+	});
+}
+
+#[test]
+fn pre_and_post_checks_behave_correctly() {
+	new_test_ext().execute_with(|| {
+		// Check initial state
+		assert_eq!(PreUpgradeCalled::get(), false);
+		assert_eq!(PostUpgradeCalled::get(), false);
+		assert_eq!(PostUpgradeCalledWith::get(), Vec::<u8>::new());
+
+		// Check pre/post hooks are called correctly when upgrade occurs.
+		VersionedMigrationV0ToV1::try_on_runtime_upgrade(true).unwrap();
+		assert_eq!(PreUpgradeCalled::get(), true);
+		assert_eq!(PostUpgradeCalled::get(), true);
+		assert_eq!(PostUpgradeCalledWith::get(), PreUpgradeReturnBytes::get().to_vec());
+
+		// Reset hook tracking state.
+		PreUpgradeCalled::set(false);
+		PostUpgradeCalled::set(false);
+
+		// Check pre/post hooks are not called when an upgrade is skipped.
+		VersionedMigrationV0ToV1::try_on_runtime_upgrade(true).unwrap();
+		assert_eq!(PreUpgradeCalled::get(), false);
+		assert_eq!(PostUpgradeCalled::get(), false);
+	})
+}