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,46 @@ 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)]
+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)
+	}
+}
+
+#[cfg(feature = "std")]
+impl core::str::FromStr for UpgradeCheckSelect {
+	type Err = &'static str;
+
+	fn from_str(s: &str) -> Result<Self, Self::Err> {
+		match s.to_lowercase().as_str() {
+			"none" => Ok(Self::None),
+			"all" => Ok(Self::All),
+			"pre-and-post" => Ok(Self::PreAndPost),
+			"try-state" => Ok(Self::TryState),
+			_ => Err("Invalid CheckSelector"),
+		}
+	}
+}
+
 /// 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

@@ -16,6 +16,7 @@
 // limitations under the License.
 
 use crate::{build_executor, state_machine_call_with_proof, SharedParams, State, LOG_TARGET};
+use frame_try_runtime::UpgradeCheckSelect;
 use parity_scale_codec::{Decode, Encode};
 use sc_executor::sp_wasm_interface::HostFunctions;
 use sp_runtime::traits::{Block as BlockT, NumberFor};
@@ -29,12 +30,16 @@ 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,
+	/// - `all`: Perform all checks (default).
+	/// - `pre-and-post`: Perform pre- and post-upgrade checks.
+	/// - `try-state`: Perform the try-state checks.
+	/// - `none`: Perform no checks.
+	///
+	/// Performing any checks will potentially invalidate the measured PoV/Weight.
+	#[clap(long, default_value = "All", verbatim_doc_comment)]
+	pub checks: UpgradeCheckSelect,
 }
 
 pub(crate) async fn on_runtime_upgrade<Block, HostFns>(