MBM try-runtime support

cumulus/parachains/runtimes/assets/asset-hub-rococo/src/lib.rs

@@ -1286,7 +1286,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

cumulus/parachains/runtimes/assets/asset-hub-westend/src/lib.rs

@@ -1377,7 +1377,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

cumulus/parachains/runtimes/bridge-hubs/bridge-hub-rococo/Cargo.toml

@@ -41,6 +41,8 @@ pallet-timestamp = { path = "../../../../../substrate/frame/timestamp", default-
 pallet-transaction-payment = { path = "../../../../../substrate/frame/transaction-payment", default-features = false }
 pallet-transaction-payment-rpc-runtime-api = { path = "../../../../../substrate/frame/transaction-payment/rpc/runtime-api", default-features = false }
 pallet-utility = { path = "../../../../../substrate/frame/utility", default-features = false }
+pallet-example-mbm = { path = "../../../../../substrate/frame/examples/multi-block-migrations", default-features = false }
+pallet-migrations = { path = "../../../../../substrate/frame/migrations", default-features = false }
 sp-api = { path = "../../../../../substrate/primitives/api", default-features = false }
 sp-block-builder = { path = "../../../../../substrate/primitives/block-builder", default-features = false }
 sp-consensus-aura = { path = "../../../../../substrate/primitives/consensus/aura", default-features = false }
@@ -121,6 +123,8 @@ snowbridge-runtime-common = { path = "../../../../../bridges/snowbridge/runtime/
 
 bridge-hub-common = { path = "../common", default-features = false }
 
+
+
 [dev-dependencies]
 static_assertions = "1.1"
 bridge-hub-test-utils = { path = "../test-utils" }
@@ -221,6 +225,8 @@ std = [
 	"xcm-builder/std",
 	"xcm-executor/std",
 	"xcm/std",
+	"pallet-example-mbm/std",
+	"pallet-migrations/std"
 ]
 
 runtime-benchmarks = [
@@ -262,6 +268,8 @@ runtime-benchmarks = [
 	"sp-runtime/runtime-benchmarks",
 	"xcm-builder/runtime-benchmarks",
 	"xcm-executor/runtime-benchmarks",
+	"pallet-example-mbm/runtime-benchmarks",
+	"pallet-migrations/runtime-benchmarks"
 ]
 
 try-runtime = [
@@ -296,6 +304,8 @@ try-runtime = [
 	"snowbridge-pallet-outbound-queue/try-runtime",
 	"snowbridge-pallet-system/try-runtime",
 	"sp-runtime/try-runtime",
+	"pallet-example-mbm/try-runtime",
+	"pallet-migrations/try-runtime"
 ]
 
 fast-runtime = []

cumulus/parachains/runtimes/bridge-hubs/bridge-hub-rococo/src/lib.rs

@@ -62,6 +62,7 @@ use frame_support::{
 	construct_runtime, derive_impl,
 	dispatch::DispatchClass,
 	genesis_builder_helper::{build_state, get_preset},
+	migrations::FreezeChainOnFailedMigration,
 	parameter_types,
 	traits::{ConstBool, ConstU32, ConstU64, ConstU8, TransformOrigin},
 	weights::{ConstantMultiplier, Weight},
@@ -272,6 +273,7 @@ impl frame_system::Config for Runtime {
 	/// The action to take on a Runtime Upgrade
 	type OnSetCode = cumulus_pallet_parachain_system::ParachainSetCode<Self>;
 	type MaxConsumers = frame_support::traits::ConstU32<16>;
+	type MultiBlockMigrator = PalletMigrations;
 }
 
 impl pallet_timestamp::Config for Runtime {
@@ -662,6 +664,8 @@ impl snowbridge_pallet_system::Config for Runtime {
 	type InboundDeliveryCost = EthereumInboundQueue;
 }
 
+impl pallet_example_mbm::Config for Runtime {}
+
 // Create the runtime by composing the FRAME pallets that were previously configured.
 construct_runtime!(
 	pub enum Runtime
@@ -727,9 +731,26 @@ construct_runtime!(
 		// Message Queue. Importantly, is registered last so that messages are processed after
 		// the `on_initialize` hooks of bridging pallets.
 		MessageQueue: pallet_message_queue = 175,
+
+		PalletMigrations: pallet_migrations = 200,
 	}
 );
 
+frame_support::parameter_types! {
+	pub MbmServiceWeight: Weight = Perbill::from_percent(80) * RuntimeBlockWeights::get().max_block;
+}
+
+impl pallet_migrations::Config for Runtime {
+	type RuntimeEvent = RuntimeEvent;
+	type MaxServiceWeight = MbmServiceWeight;
+	type WeightInfo = ();
+	type Migrations = ();
+	type CursorMaxLen = ConstU32<65_536>;
+	type IdentifierMaxLen = ConstU32<256>;
+	type MigrationStatusHandler = ();
+	type FailedMigrationHandler = FreezeChainOnFailedMigration;
+}
+
 /// Proper alias for bridge GRANDPA pallet used to bridge with the bulletin chain.
 pub type BridgeRococoBulletinGrandpa = BridgePolkadotBulletinGrandpa;
 /// Proper alias for bridge messages pallet used to bridge with the bulletin chain.
@@ -1033,7 +1054,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

cumulus/parachains/runtimes/bridge-hubs/bridge-hub-westend/src/lib.rs

@@ -733,7 +733,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

cumulus/parachains/runtimes/collectives/collectives-westend/src/lib.rs

@@ -927,7 +927,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

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

@@ -658,7 +658,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

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

@@ -652,7 +652,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

cumulus/parachains/runtimes/coretime/coretime-westend/src/lib.rs

@@ -644,7 +644,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

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

@@ -619,7 +619,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

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

@@ -619,7 +619,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

cumulus/parachains/runtimes/testing/penpal/src/lib.rs

@@ -837,7 +837,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, RuntimeBlockWeights::get().max_block)
 		}

cumulus/polkadot-parachain/src/fake_runtime_api/asset_hub_polkadot_aura.rs

@@ -152,7 +152,7 @@ sp_api::impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(_: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(_: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			unimplemented!()
 		}
 

cumulus/polkadot-parachain/src/fake_runtime_api/aura.rs

@@ -152,7 +152,7 @@ sp_api::impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(_: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(_: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			unimplemented!()
 		}
 

polkadot/runtime/rococo/src/lib.rs

@@ -2229,7 +2229,7 @@ sp_api::impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			log::info!("try-runtime::on_runtime_upgrade rococo.");
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, BlockWeights::get().max_block)
@@ -2500,7 +2500,7 @@ sp_api::impl_runtime_apis! {
 #[cfg(all(test, feature = "try-runtime"))]
 mod remote_tests {
 	use super::*;
-	use frame_try_runtime::{runtime_decl_for_try_runtime::TryRuntime, UpgradeCheckSelect};
+	use frame_try_runtime::{runtime_decl_for_try_runtime::TryRuntime, TryOnRuntimeUpgradeOpts};
 	use remote_externalities::{
 		Builder, Mode, OfflineConfig, OnlineConfig, SnapshotConfig, Transport,
 	};
@@ -2532,6 +2532,8 @@ mod remote_tests {
 			.build()
 			.await
 			.unwrap();
-		ext.execute_with(|| Runtime::on_runtime_upgrade(UpgradeCheckSelect::PreAndPost));
+		ext.execute_with(|| {
+			Runtime::on_runtime_upgrade(TryOnRuntimeUpgradeOpts::new().pre_and_post(true).build())
+		});
 	}
 }

polkadot/runtime/westend/src/lib.rs

@@ -2253,7 +2253,7 @@ sp_api::impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			log::info!("try-runtime::on_runtime_upgrade westend.");
 			let weight = Executive::try_runtime_upgrade(checks).unwrap();
 			(weight, BlockWeights::get().max_block)
@@ -2536,7 +2536,7 @@ sp_api::impl_runtime_apis! {
 #[cfg(all(test, feature = "try-runtime"))]
 mod remote_tests {
 	use super::*;
-	use frame_try_runtime::{runtime_decl_for_try_runtime::TryRuntime, UpgradeCheckSelect};
+	use frame_try_runtime::{runtime_decl_for_try_runtime::TryRuntime, TryOnRuntimeUpgradeOpts};
 	use remote_externalities::{
 		Builder, Mode, OfflineConfig, OnlineConfig, SnapshotConfig, Transport,
 	};
@@ -2568,7 +2568,9 @@ mod remote_tests {
 			.build()
 			.await
 			.unwrap();
-		ext.execute_with(|| Runtime::on_runtime_upgrade(UpgradeCheckSelect::PreAndPost));
+		ext.execute_with(|| {
+			Runtime::on_runtime_upgrade(TryOnRuntimeUpgradeOpts::new().pre_and_post(true).build())
+		});
 	}
 }
 

substrate/bin/node/runtime/src/lib.rs

@@ -3157,7 +3157,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::TryOnRuntimeUpgradeOpts) -> (Weight, Weight) {
 			// NOTE: intentional unwrap: we don't want to propagate the error backwards, and want to
 			// have a backtrace here. If any of the pre/post migration checks fail, we shall stop
 			// right here and right now.

substrate/frame/examples/multi-block-migrations/Cargo.toml

@@ -21,6 +21,7 @@ frame-benchmarking = { path = "../../benchmarking", default-features = false, op
 log = { version = "0.4.20", default-features = false }
 scale-info = { version = "2.10.0", default-features = false }
 sp-io = { path = "../../../primitives/io", default-features = false }
+sp-std = { path = "../../../primitives/std", default-features = false }
 
 [features]
 default = ["std"]
@@ -33,6 +34,7 @@ std = [
 	"pallet-migrations/std",
 	"scale-info/std",
 	"sp-io/std",
+	"sp-std/std"
 ]
 runtime-benchmarks = [
 	"frame-benchmarking/runtime-benchmarks",

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

@@ -29,6 +29,9 @@ use frame_support::{
 	weights::WeightMeter,
 };
 
+#[cfg(feature = "try-runtime")]
+use sp_std::collections::btree_map::BTreeMap;
+
 mod benchmarks;
 mod tests;
 pub mod weights;
@@ -115,4 +118,29 @@ 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");
+
+		for (key, value) in MyMap::<T>::iter() {
+			let prev_value =
+				prev_map.get(&key).expect("Key not found in the previous storage state");
+			assert_eq!(value as u32, *prev_value, "Migration failed for key {}", key);
+		}
+
+		Ok(())
+	}
 }