migrations: VersionedRuntimeUpgrade

frame/support/Cargo.toml

@@ -80,6 +80,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"]

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,138 @@ 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 VersionedRuntimeUpgrades `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.
+///
+/// Example:
+/// ```ignore
+/// // Migrations to pass to the Executive pallet.
+/// pub type Migrations = (
+/// 	// ...other migrations
+/// 	VersionedRuntimeUpgrade<
+/// 		4, // From on-chain version 4
+/// 		5, // To on-chain version 5
+/// 		parachains_configuration::migration::v5::MigrateToV5<Runtime>,
+/// 		Configuration,
+/// 		RocksDbWeight,
+/// 	>,
+/// 	VersionedRuntimeUpgrade<
+/// 		5, // From on-chain version 5
+/// 		7, // To on-chain version 7
+/// 		parachains_configuration::migration::v6::MigrateToV6<Runtime>,
+/// 		Configuration,
+/// 		RocksDbWeight,
+/// 	>,
+/// 	// ...other migrations
+/// );
+/// ```
+#[cfg(feature = "experimental")]
+pub struct VersionedRuntimeUpgrade<const FROM: u16, const TO: u16, Inner, Pallet, Weight> {
+	_marker: PhantomData<(Inner, Pallet, Weight)>,
+}
+
+/// 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-chian 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 an Option
+	/// 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(Some(Inner::pre_upgrade()?).encode())
+		} else {
+			Ok(None::<Vec<u8>>.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 post_upgrade if the migration just ran.
+	///
+	/// pre_upgrade passes Some(Vec<u8>) to post_upgrade if the migration ran, and None otherwise.
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(pre_upgrade_return_bytes: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
+		use codec::DecodeAll;
+		let maybe_inner_bytes = <Option<Vec<u8>>>::decode_all(&mut &pre_upgrade_return_bytes[..])
+			.map_err(|_| {
+			"VersionedRuntimeUpgrade post_upgrade failed to decode pre_upgrade bytes"
+		})?;
+		match maybe_inner_bytes {
+			Some(inner_bytes) => Inner::post_upgrade(inner_bytes),
+			None => 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 +312,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 +327,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 => {