Selectable on-runtime-upgrade checks

bin/node-template/runtime/src/lib.rs

@@ -539,7 +539,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: bool) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (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.

bin/node/runtime/src/lib.rs

@@ -2197,7 +2197,7 @@ impl_runtime_apis! {
 
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
-		fn on_runtime_upgrade(checks: bool) -> (Weight, Weight) {
+		fn on_runtime_upgrade(checks: frame_try_runtime::UpgradeCheckSelect) -> (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.

frame/executive/src/lib.rs

@@ -333,8 +333,10 @@ where
 	///
 	/// Runs the try-state code both before and after the migration function if `checks` is set to
 	/// `true`. Also, if set to `true`, it runs the `pre_upgrade` and `post_upgrade` hooks.
-	pub fn try_runtime_upgrade(checks: bool) -> Result<Weight, &'static str> {
-		if checks {
+	pub fn try_runtime_upgrade(
+		checks: frame_try_runtime::UpgradeCheckSelect,
+	) -> Result<Weight, &'static str> {
+		if checks.try_state() {
 			let _guard = frame_support::StorageNoopGuard::default();
 			<AllPalletsWithSystem as frame_support::traits::TryState<System::BlockNumber>>::try_state(
 				frame_system::Pallet::<System>::block_number(),
@@ -344,10 +346,10 @@ where
 
 		let weight =
 			<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::try_on_runtime_upgrade(
-				checks,
+				checks.pre_and_post(),
 			)?;
 
-		if checks {
+		if checks.try_state() {
 			let _guard = frame_support::StorageNoopGuard::default();
 			<AllPalletsWithSystem as frame_support::traits::TryState<System::BlockNumber>>::try_state(
 				frame_system::Pallet::<System>::block_number(),

frame/support/src/traits.rs

@@ -121,4 +121,4 @@ pub use messages::{
 #[cfg(feature = "try-runtime")]
 mod try_runtime;
 #[cfg(feature = "try-runtime")]
-pub use try_runtime::{Select as TryStateSelect, TryState};
+pub use try_runtime::{Select as TryStateSelect, TryState, UpgradeCheckSelect};

frame/support/src/traits/try_runtime.rs

@@ -21,7 +21,7 @@ use impl_trait_for_tuples::impl_for_tuples;
 use sp_arithmetic::traits::AtLeast32BitUnsigned;
 use sp_std::prelude::*;
 
-// Which state tests to execute.
+/// Which state tests to execute.
 #[derive(codec::Encode, codec::Decode, Clone)]
 pub enum Select {
 	/// None of them.
@@ -81,6 +81,31 @@ impl sp_std::str::FromStr for Select {
 	}
 }
 
+/// Select which checks should be run when trying a runtime upgrade upgrade.
+#[derive(codec::Encode, codec::Decode, Clone, Debug, Copy)]
+pub enum UpgradeCheckSelect {
+	/// Run no checks.
+	None,
+	/// Run the `try_state`, `pre_upgrade` and `post_upgrade` checks.
+	All,
+	/// Run the `pre_upgrade` and `post_upgrade` checks.
+	PreAndPost,
+	/// Run the `try_state` checks.
+	TryState,
+}
+
+impl UpgradeCheckSelect {
+	/// Whether the pre- and post-upgrade checks are selected.
+	pub fn pre_and_post(&self) -> bool {
+		matches!(self, Self::All | Self::PreAndPost)
+	}
+
+	/// Whether the try-state checks are selected.
+	pub fn try_state(&self) -> bool {
+		matches!(self, Self::All | Self::TryState)
+	}
+}
+
 /// Execute some checks to ensure the internal state of a pallet is consistent.
 ///
 /// Usually, these checks should check all of the invariants that are expected to be held on all of

frame/try-runtime/src/lib.rs

@@ -20,7 +20,7 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 #![cfg(feature = "try-runtime")]
 
-pub use frame_support::traits::TryStateSelect;
+pub use frame_support::traits::{TryStateSelect, UpgradeCheckSelect};
 use frame_support::weights::Weight;
 
 sp_api::decl_runtime_apis! {
@@ -37,7 +37,7 @@ sp_api::decl_runtime_apis! {
 		/// If `checks` is `true`, `pre_migrate` and `post_migrate` of each migration and
 		/// `try_state` of all pallets will be executed. Else, no. If checks are executed, the PoV
 		/// tracking is likely inaccurate.
-		fn on_runtime_upgrade(checks: bool) -> (Weight, Weight);
+		fn on_runtime_upgrade(checks: UpgradeCheckSelect) -> (Weight, Weight);
 
 		/// Execute the given block, but optionally disable state-root and signature checks.
 		///

utils/frame/try-runtime/cli/src/commands/on_runtime_upgrade.rs

@@ -29,12 +29,37 @@ pub struct OnRuntimeUpgradeCmd {
 	#[command(subcommand)]
 	pub state: State,
 
-	/// Execute `try_state`, `pre_upgrade` and `post_upgrade` checks as well.
+	/// Select which optional checks to perform.
 	///
-	/// This will perform more checks, but it will also makes the reported PoV/Weight be
-	/// inaccurate.
-	#[clap(long)]
-	pub checks: bool,
+	/// Performing any checks will potentially invalidate the measured PoV/Weight.
+	#[clap(long, value_enum, default_value_t = UpgradeCheckSelect::All)]
+	pub checks: UpgradeCheckSelect,
+}
+
+/// This is an adapter for [`frame_try_runtime::UpgradeCheckSelect`] since that does not implement
+/// `clap::ValueEnum`.
+#[derive(clap::ValueEnum, Debug, Clone, Copy)]
+#[value(rename_all = "kebab-case")]
+pub enum UpgradeCheckSelect {
+	/// Perform no checks.
+	None,
+	/// Perform all checks.
+	All,
+	/// Perform pre- and post-upgrade checks.
+	PreAndPost,
+	/// Perform the try-state checks.
+	TryState,
+}
+
+impl From<UpgradeCheckSelect> for frame_try_runtime::UpgradeCheckSelect {
+	fn from(x: UpgradeCheckSelect) -> Self {
+		match x {
+			UpgradeCheckSelect::None => Self::None,
+			UpgradeCheckSelect::All => Self::All,
+			UpgradeCheckSelect::PreAndPost => Self::PreAndPost,
+			UpgradeCheckSelect::TryState => Self::TryState,
+		}
+	}
 }
 
 pub(crate) async fn on_runtime_upgrade<Block, HostFns>(
@@ -52,12 +77,13 @@ where
 {
 	let executor = build_executor(&shared);
 	let ext = command.state.into_ext::<Block, HostFns>(&shared, &executor, None).await?;
+	let checks: frame_try_runtime::UpgradeCheckSelect = command.checks.into();
 
 	let (_, encoded_result) = state_machine_call_with_proof::<Block, HostFns>(
 		&ext,
 		&executor,
 		"TryRuntime_on_runtime_upgrade",
-		command.checks.encode().as_ref(),
+		checks.encode().as_ref(),
 		Default::default(), // we don't really need any extensions here.
 		shared.export_proof,
 	)?;