migrations: VersionedRuntimeUpgrade

frame/support/src/migrations.rs

@@ -15,19 +15,126 @@
 // 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},
+	traits::{
+		GetStorageVersion, NoStorageVersionSet, OnRuntimeUpgrade, PalletInfoAccess, StorageVersion,
+	},
 	weights::{RuntimeDbWeight, Weight},
 };
 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;
 
+/// 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 4 type parameters:
+/// - `Version`: The version being upgraded *from*.
+/// - `Inner`: An implementation of `OnRuntimeUpgrade`.
+/// - `Pallet`: The Pallet being upgraded.
+/// - `Weight`: The runtime's RuntimeDbWeight implementation.
+///
+/// `Version` is used to check if the migration should be run. If it is run, the pallets on-chain
+/// version is incremented.
+///
+/// Example:
+/// ```ignore
+/// parameter_types! {
+/// 	pub const V4ToV5: StorageVersion = StorageVersion::new(4);
+/// 	pub const V5ToV6: StorageVersion = StorageVersion::new(5);
+/// }
+///
+/// // Migrations to pass to the Executive pallet.
+/// pub type Migrations = (
+/// 	// ...other migrations
+/// 	VersionedRuntimeUpgrade<
+/// 		V4ToV5,
+/// 		parachains_configuration::migration::v5::MigrateToV5<Runtime>,
+/// 		Configuration,
+/// 		RocksDbWeight,
+/// 	>,
+/// 	VersionedRuntimeUpgrade<
+/// 		V5ToV6,
+/// 		parachains_configuration::migration::v6::MigrateToV6<Runtime>,
+/// 		Configuration,
+/// 		RocksDbWeight,
+/// 	>,
+/// 	// ...other migrations
+/// );
+/// ```
+pub struct VersionedRuntimeUpgrade<Version, Inner, Pallet, Weight> {
+	_marker: PhantomData<(Version, Inner, Pallet, Weight)>,
+}
+
+/// Implementation of the `OnRuntimeUpgrade` trait for `VersionedRuntimeUpgrade`.
+///
+/// Primarily, it
+/// 1. Logs a warning in `pre_upgrade` if the storage version specified in the pallet appears to be
+/// misconfigured before returning `Inner::pre_upgrade`.
+///
+/// 2. Performs the actual runtime upgrade in `on_runtime_upgrade` only if the on-chain version of
+/// the pallets storage matches `Version`. If the versions do not match, it writes a log notifying
+/// the developer that the migration is a noop.
+impl<
+		Version: Get<StorageVersion>,
+		Inner: OnRuntimeUpgrade,
+		Pallet: GetStorageVersion<CurrentStorageVersion = StorageVersion> + PalletInfoAccess,
+		DbWeight: Get<RuntimeDbWeight>,
+	> OnRuntimeUpgrade for VersionedRuntimeUpgrade<Version, Inner, Pallet, DbWeight>
+{
+	/// VersionedRuntimeUpgrade pre-upgrade checks.
+	///
+	/// Logs a warning if the storage version specified in the pallet appears to be misconfigured,
+	/// before returning `Inner::pre_upgrade`.
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
+		let current_version = Pallet::current_storage_version();
+		if current_version < Version::get() + 1 {
+			log::warn!(
+				"Pallet '{}' appears to be misconfigured: its version is less than the on-chain version that will be set by this VersionedRuntimeUpgrade. Please ensure the pallet version is set correctly.",
+				Pallet::name()
+			);
+		}
+
+		Inner::pre_upgrade()
+	}
+
+	/// Actually executes the versioned runtime upgrade.
+	///
+	/// First checks if the pallet's on-chain storage version matches the version of this upgrade.
+	/// If it matches, it calls `Inner::on_runtime_upgrade`, increments 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 == Version::get() {
+			// Execute the migration
+			let weight = Inner::on_runtime_upgrade();
+
+			// Update the on-chain version
+			let next = Version::get() + 1;
+			next.put::<Pallet>();
+
+			weight.saturating_add(DbWeight::get().reads_writes(1, 1))
+		} else {
+			log::warn!(
+				"{} VersionedOnRuntimeUpgrade {:?} skipped because current on-chain version is {:?}.",
+				Pallet::name(),
+				Version::get(),
+				on_chain_version
+			);
+			DbWeight::get().reads(1)
+		}
+	}
+}
+
 /// Can store the current pallet version in storage.
 pub trait StoreCurrentStorageVersion<T: GetStorageVersion + PalletInfoAccess> {
 	/// Write the current storage version to the storage.
@@ -185,6 +292,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 +307,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 => {