MBM try-runtime support

.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:

prdoc/pr_4251.prdoc

@@ -0,0 +1,59 @@
+title: MBM `try-runtime` support
+doc:
+- audience: Runtime Dev
+  description: "# MBM try-runtime support\n\nThis 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.\n\n## Integration\n\
+    \nTo make use of this in your Multi-Block-Migration, just implement the two new\
+    \ hooks and test pre- and post-conditions in them:\n\n```rust\n#[cfg(feature =\
+    \ \"try-runtime\")]\nfn pre_upgrade() -> Result<Vec<u8>, frame_support::sp_runtime::TryRuntimeError>\
+    \ {\n\t// ...\n}\n\n#[cfg(feature = \"try-runtime\")]\nfn post_upgrade(prev: Vec<u8>)\
+    \ -> Result<(), frame_support::sp_runtime::TryRuntimeError> {\n    // ...\n}\n\
+    ```\n\nYou 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.\n\
+    \nChanges:\n- Adds `try-runtime` gated methods `pre_upgrade` and `post_upgrade`\
+    \ on `SteppedMigration`\n- Adds `try-runtime` gated methods `nth_pre_upgrade`\
+    \ and `nth_post_upgrade` on `SteppedMigrations`\n- 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/Cargo.toml

@@ -21,6 +21,7 @@ frame-benchmarking = { optional = true, workspace = true }
 log = { workspace = true }
 scale-info = { workspace = true }
 sp-io = { workspace = true }
+sp-std = { workspace = true }
 
 [features]
 default = ["std"]

substrate/frame/examples/multi-block-migrations/src/lib.rs

@@ -69,6 +69,8 @@
 pub mod migrations;
 mod mock;
 
+extern crate alloc;
+
 pub use pallet::*;
 
 #[frame_support::pallet]

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

@@ -29,6 +29,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 +121,38 @@ 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 value = MyMap::<T>::get(key).expect("Failed to get the value after the migration");
+			assert_eq!(
+				value, value as u64,
+				"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>(_);
@@ -700,6 +723,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()),
@@ -734,6 +767,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());
@@ -758,14 +801,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
 	}