MBM try-runtime support (#4251)

# MBM try-runtime support

This MR adds support to the try-runtime trait such that the
try-runtime-CLI will be able to support MBM testing
[here](paritytech/try-runtime-cli#90). It mainly
adds two feature-gated hooks to the `SteppedMigration` hook to
facilitate testing. These hooks are named `pre_upgrade` and
`post_upgrade` and have the same signature and implications as for
single-block migrations.

## Integration

To make use of this in your Multi-Block-Migration, just implement the
two new hooks and test pre- and post-conditions in them:

```rust
#[cfg(feature = "try-runtime")]
fn pre_upgrade() -> Result<Vec<u8>, frame_support::sp_runtime::TryRuntimeError> {
	// ...
}

#[cfg(feature = "try-runtime")]
fn post_upgrade(prev: Vec<u8>) -> Result<(), frame_support::sp_runtime::TryRuntimeError> {
    // ...
}
```

You may return an error or panic in these functions to indicate failure.
This will then show up in the try-runtime-CLI and can be used in CI for
testing.

Changes:
- Adds `try-runtime` gated methods `pre_upgrade` and `post_upgrade` on
`SteppedMigration`
- Adds `try-runtime` gated methods `nth_pre_upgrade` and
`nth_post_upgrade` on `SteppedMigrations`
- Modifies `pallet_migrations` implementation to run pre_upgrade and
post_upgrade steps at the appropriate times, and panic in the event of
migration failure.

---------

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Signed-off-by: georgepisaltu <george.pisaltu@parity.io>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: claravanstaden <claravanstaden64@gmail.com>
Co-authored-by: ggwpez <ggwpez@users.noreply.github.com>
Co-authored-by: georgepisaltu <george.pisaltu@parity.io>

.gitlab/pipeline/test.yml

@@ -110,6 +110,30 @@ test-linux-stable-codecov:
         codecovcli -v do-upload -f target/coverage/result/report-${CI_NODE_INDEX}.lcov --disable-search -t ${CODECOV_TOKEN} -r paritytech/polkadot-sdk --commit-sha ${CI_COMMIT_SHA} --fail-on-error --git-service github;
       fi
 
+# some tests do not run with `try-runtime` feature enabled
+# https://github.com/paritytech/polkadot-sdk/pull/4251#discussion_r1624282143
+test-linux-stable-no-try-runtime:
+  stage: test
+  extends:
+    - .docker-env
+    - .common-refs
+    - .run-immediately
+    - .pipeline-stopper-artifacts
+  variables:
+    RUST_TOOLCHAIN: stable
+    # Enable debug assertions since we are running optimized builds for testing
+    # but still want to have debug assertions.
+    RUSTFLAGS: "-Cdebug-assertions=y -Dwarnings"
+  script:
+    - >
+      time cargo nextest run \
+        --workspace \
+        --locked \
+        --release \
+        --no-fail-fast \
+        --cargo-quiet \
+        --features experimental,riscv,ci-only-tests
+
 test-doc:
   stage: test
   extends:

cumulus/parachains/runtimes/people/people-rococo/src/lib.rs

@@ -111,6 +111,7 @@ pub type UncheckedExtrinsic =
 /// Migrations to apply on runtime upgrade.
 pub type Migrations = (
 	pallet_collator_selection::migration::v2::MigrationToV2<Runtime>,
+	cumulus_pallet_xcmp_queue::migration::v5::MigrateV4ToV5<Runtime>,
 	// permanent
 	pallet_xcm::migration::MigrateToLatestXcmVersion<Runtime>,
 );

prdoc/pr_4251.prdoc

@@ -0,0 +1,79 @@
+title: MBM `try-runtime` support
+doc:
+- audience: Runtime Dev
+  description: |
+    # MBM try-runtime support
+
+    This MR adds support to the try-runtime
+    trait such that the try-runtime-CLI will be able to support MBM testing [here](https://github.com/paritytech/try-runtime-cli/pull/90).
+    It mainly adds two feature-gated hooks to the `SteppedMigration` hook to facilitate
+    testing. These hooks are named `pre_upgrade` and `post_upgrade` and have the
+    same signature and implications as for single-block migrations.
+
+    ## Integration
+
+    To make use of this in your Multi-Block-Migration, just implement the two new hooks and test pre- and post-conditions in them:
+
+    ```rust
+    #[cfg(feature = "try-runtime")]
+    fn pre_upgrade() -> Result<Vec<u8>, frame_support::sp_runtime::TryRuntimeError>
+    {
+      // ...
+    }
+
+    #[cfg(feature = "try-runtime")]
+    fn post_upgrade(prev: Vec<u8>) -> Result<(), frame_support::sp_runtime::TryRuntimeError> {
+        // ...
+    }
+    ```
+
+    You may return an error or panic in these functions to indicate failure.
+    This will then show up in the try-runtime-CLI and can be used in CI for testing.
+
+
+    Changes:
+    - Adds `try-runtime` gated methods `pre_upgrade` and `post_upgrade`
+      on `SteppedMigration`
+    - Adds `try-runtime` gated methods `nth_pre_upgrade`
+      and `nth_post_upgrade` on `SteppedMigrations`
+    - Modifies `pallet_migrations`
+      implementation to run pre_upgrade and post_upgrade steps at the appropriate times, and panic in the event of migration failure.
+crates:
+- name: asset-hub-rococo-runtime
+  bump: minor
+- name: asset-hub-westend-runtime
+  bump: minor
+- name: bridge-hub-rococo-runtime
+  bump: minor
+- name: bridge-hub-westend-runtime
+  bump: minor
+- name: collectives-westend-runtime
+  bump: minor
+- name: contracts-rococo-runtime
+  bump: minor
+- name: coretime-rococo-runtime
+  bump: minor
+- name: coretime-westend-runtime
+  bump: minor
+- name: people-rococo-runtime
+  bump: minor
+- name: people-westend-runtime
+  bump: minor
+- name: penpal-runtime
+  bump: minor
+- name: polkadot-parachain-bin
+  bump: minor
+- name: rococo-runtime
+  bump: minor
+- name: westend-runtime
+  bump: minor
+- name: frame-executive
+  bump: minor
+- name: pallet-migrations
+  bump: minor
+- name: frame-support
+  bump: minor
+- name: frame-system
+  bump: minor
+- name: frame-try-runtime
+  bump: minor

substrate/frame/examples/multi-block-migrations/src/migrations/v1/mod.rs

@@ -21,6 +21,8 @@
 //! [`v0::MyMap`](`crate::migrations::v1::v0::MyMap`) storage map, transforms them,
 //! and inserts them into the [`MyMap`](`crate::pallet::MyMap`) storage map.
 
+extern crate alloc;
+
 use super::PALLET_MIGRATIONS_ID;
 use crate::pallet::{Config, MyMap};
 use frame_support::{
@@ -29,6 +31,12 @@ use frame_support::{
 	weights::WeightMeter,
 };
 
+#[cfg(feature = "try-runtime")]
+use alloc::collections::btree_map::BTreeMap;
+
+#[cfg(feature = "try-runtime")]
+use alloc::vec::Vec;
+
 mod benchmarks;
 mod tests;
 pub mod weights;
@@ -115,4 +123,39 @@ impl<T: Config, W: weights::WeightInfo> SteppedMigration for LazyMigrationV1<T,
 		}
 		Ok(cursor)
 	}
+
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<Vec<u8>, frame_support::sp_runtime::TryRuntimeError> {
+		use codec::Encode;
+
+		// Return the state of the storage before the migration.
+		Ok(v0::MyMap::<T>::iter().collect::<BTreeMap<_, _>>().encode())
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(prev: Vec<u8>) -> Result<(), frame_support::sp_runtime::TryRuntimeError> {
+		use codec::Decode;
+
+		// Check the state of the storage after the migration.
+		let prev_map = BTreeMap::<u32, u32>::decode(&mut &prev[..])
+			.expect("Failed to decode the previous storage state");
+
+		// Check the len of prev and post are the same.
+		assert_eq!(
+			MyMap::<T>::iter().count(),
+			prev_map.len(),
+			"Migration failed: the number of items in the storage after the migration is not the same as before"
+		);
+
+		for (key, value) in prev_map {
+			let new_value =
+				MyMap::<T>::get(key).expect("Failed to get the value after the migration");
+			assert_eq!(
+				value as u64, new_value,
+				"Migration failed: the value after the migration is not the same as before"
+			);
+		}
+
+		Ok(())
+	}
 }

substrate/frame/migrations/Cargo.toml

@@ -12,6 +12,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
 codec = { features = ["derive"], workspace = true }
+cfg-if = { workspace = true }
 docify = { workspace = true }
 impl-trait-for-tuples = { workspace = true }
 log = { workspace = true, default-features = true }

substrate/frame/migrations/src/lib.rs

@@ -70,21 +70,26 @@
 //! points to the currently active migration and stores its inner cursor. The inner cursor can then
 //! be used by the migration to store its inner state and advance. Each time when the migration
 //! returns `Some(cursor)`, it signals the pallet that it is not done yet.
+//!
 //! The cursor is reset on each runtime upgrade. This ensures that it starts to execute at the
 //! first migration in the vector. The pallets cursor is only ever incremented or set to `Stuck`
 //! once it encounters an error (Goal 4). Once in the stuck state, the pallet will stay stuck until
 //! it is fixed through manual governance intervention.
+//!
 //! As soon as the cursor of the pallet becomes `Some(_)`; [`MultiStepMigrator::ongoing`] returns
 //! `true` (Goal 2). This can be used by upstream code to possibly pause transactions.
 //! In `on_initialize` the pallet will load the current migration and check whether it was already
 //! executed in the past by checking for membership of its ID in the [`Historic`] set. Historic
 //! migrations are skipped without causing an error. Each successfully executed migration is added
 //! to this set (Goal 5).
+//!
 //! This proceeds until no more migrations remain. At that point, the event `UpgradeCompleted` is
 //! emitted (Goal 1).
+//!
 //! The execution of each migration happens by calling [`SteppedMigration::transactional_step`].
 //! This function wraps the inner `step` function into a transactional layer to allow rollback in
 //! the error case (Goal 6).
+//!
 //! Weight limits must be checked by the migration itself. The pallet provides a [`WeightMeter`] for
 //! that purpose. The pallet may return [`SteppedMigrationError::InsufficientWeight`] at any point.
 //! In that scenario, one of two things will happen: if that migration was exclusively executed
@@ -156,11 +161,15 @@ use core::ops::ControlFlow;
 use frame_support::{
 	defensive, defensive_assert,
 	migrations::*,
+	pallet_prelude::*,
 	traits::Get,
 	weights::{Weight, WeightMeter},
 	BoundedVec,
 };
-use frame_system::{pallet_prelude::BlockNumberFor, Pallet as System};
+use frame_system::{
+	pallet_prelude::{BlockNumberFor, *},
+	Pallet as System,
+};
 use sp_runtime::Saturating;
 
 /// Points to the next migration to execute.
@@ -262,18 +271,32 @@ pub type IdentifierOf<T> = BoundedVec<u8, <T as Config>::IdentifierMaxLen>;
 pub type ActiveCursorOf<T> = ActiveCursor<RawCursorOf<T>, BlockNumberFor<T>>;
 
 /// Trait for a tuple of No-OP migrations with one element.
+#[impl_trait_for_tuples::impl_for_tuples(30)]
 pub trait MockedMigrations: SteppedMigrations {
 	/// The migration should fail after `n` steps.
 	fn set_fail_after(n: u32);
 	/// The migration should succeed after `n` steps.
 	fn set_success_after(n: u32);
 }
 
+#[cfg(feature = "try-runtime")]
+/// Wrapper for pre-upgrade bytes, allowing us to impl MEL on it.
+///
+/// For `try-runtime` testing only.
+#[derive(Debug, Clone, Eq, PartialEq, Encode, Decode, scale_info::TypeInfo, Default)]
+struct PreUpgradeBytesWrapper(pub Vec<u8>);
+
+/// Data stored by the pre-upgrade hook of the MBMs. Only used for `try-runtime` testing.
+///
+/// Define this outside of the pallet so it is not confused with actual storage.
+#[cfg(feature = "try-runtime")]
+#[frame_support::storage_alias]
+type PreUpgradeBytes<T: Config> =
+	StorageMap<Pallet<T>, Twox64Concat, IdentifierOf<T>, PreUpgradeBytesWrapper, ValueQuery>;
+
 #[frame_support::pallet]
 pub mod pallet {
 	use super::*;
-	use frame_support::pallet_prelude::*;
-	use frame_system::pallet_prelude::*;
 
 	#[pallet::pallet]
 	pub struct Pallet<T>(_);
@@ -701,6 +724,16 @@ impl<T: Config> Pallet<T> {
 		}
 
 		let max_steps = T::Migrations::nth_max_steps(cursor.index);
+
+		// If this is the first time running this migration, exec the pre-upgrade hook.
+		#[cfg(feature = "try-runtime")]
+		if !PreUpgradeBytes::<T>::contains_key(&bounded_id) {
+			let bytes = T::Migrations::nth_pre_upgrade(cursor.index)
+				.expect("Invalid cursor.index")
+				.expect("Pre-upgrade failed");
+			PreUpgradeBytes::<T>::insert(&bounded_id, PreUpgradeBytesWrapper(bytes));
+		}
+
 		let next_cursor = T::Migrations::nth_transactional_step(
 			cursor.index,
 			cursor.inner_cursor.clone().map(|c| c.into_inner()),
@@ -735,6 +768,16 @@ impl<T: Config> Pallet<T> {
 			},
 			Ok(None) => {
 				// A migration is done when it returns cursor `None`.
+
+				// Run post-upgrade checks.
+				#[cfg(feature = "try-runtime")]
+				T::Migrations::nth_post_upgrade(
+					cursor.index,
+					PreUpgradeBytes::<T>::get(&bounded_id).0,
+				)
+				.expect("Invalid cursor.index.")
+				.expect("Post-upgrade failed.");
+
 				Self::deposit_event(Event::MigrationCompleted { index: cursor.index, took });
 				Historic::<T>::insert(&bounded_id, ());
 				cursor.goto_next_migration(System::<T>::block_number());
@@ -759,14 +802,21 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// Fail the current runtime upgrade, caused by `migration`.
+	///
+	/// When the `try-runtime` feature is enabled, this function will panic.
+	// Allow unreachable code so it can compile without warnings when `try-runtime` is enabled.
 	fn upgrade_failed(migration: Option<u32>) {
 		use FailedMigrationHandling::*;
 		Self::deposit_event(Event::UpgradeFailed);
 
-		match T::FailedMigrationHandler::failed(migration) {
-			KeepStuck => Cursor::<T>::set(Some(MigrationCursor::Stuck)),
-			ForceUnstuck => Cursor::<T>::kill(),
-			Ignore => {},
+		if cfg!(feature = "try-runtime") {
+			panic!("Migration with index {:?} failed.", migration);
+		} else {
+			match T::FailedMigrationHandler::failed(migration) {
+				KeepStuck => Cursor::<T>::set(Some(MigrationCursor::Stuck)),
+				ForceUnstuck => Cursor::<T>::kill(),
+				Ignore => {},
+			}
 		}
 	}
 

substrate/frame/migrations/src/mock_helpers.rs

@@ -43,6 +43,12 @@ pub enum MockedMigrationKind {
 	/// Cause an [`SteppedMigrationError::InsufficientWeight`] error after its number of steps
 	/// elapsed.
 	HighWeightAfter(Weight),
+	/// PreUpgrade should fail.
+	#[cfg(feature = "try-runtime")]
+	PreUpgradeFail,
+	/// PostUpgrade should fail.
+	#[cfg(feature = "try-runtime")]
+	PostUpgradeFail,
 }
 use MockedMigrationKind::*; // C style
 
@@ -99,6 +105,8 @@ impl SteppedMigrations for MockedMigrations {
 				Err(SteppedMigrationError::Failed)
 			},
 			TimeoutAfter => unreachable!(),
+			#[cfg(feature = "try-runtime")]
+			PreUpgradeFail | PostUpgradeFail => Ok(None),
 		})
 	}
 
@@ -115,6 +123,31 @@ impl SteppedMigrations for MockedMigrations {
 		MIGRATIONS::get().get(n as usize).map(|(_, s)| Some(*s))
 	}
 
+	#[cfg(feature = "try-runtime")]
+	fn nth_pre_upgrade(n: u32) -> Option<Result<Vec<u8>, sp_runtime::TryRuntimeError>> {
+		let (kind, _) = MIGRATIONS::get()[n as usize];
+
+		if let PreUpgradeFail = kind {
+			return Some(Err("Some pre-upgrade error".into()))
+		}
+
+		Some(Ok(vec![]))
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn nth_post_upgrade(
+		n: u32,
+		_state: Vec<u8>,
+	) -> Option<Result<(), sp_runtime::TryRuntimeError>> {
+		let (kind, _) = MIGRATIONS::get()[n as usize];
+
+		if let PostUpgradeFail = kind {
+			return Some(Err("Some post-upgrade error".into()))
+		}
+
+		Some(Ok(()))
+	}
+
 	fn cursor_max_encoded_len() -> usize {
 		65_536
 	}