CI friendly fast mode for try state checks

frame/bags-list/src/lib.rs

@@ -270,7 +270,10 @@ pub mod pallet {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn try_state(_: BlockNumberFor<T>) -> Result<(), TryRuntimeError> {
+		fn try_state(
+			_: BlockNumberFor<T>,
+			_: frame_support::traits::TryStateSelect,
+		) -> Result<(), TryRuntimeError> {
 			<Self as SortedListProvider<T::AccountId>>::try_state()
 		}
 	}

frame/collective/src/lib.rs

@@ -350,7 +350,10 @@ pub mod pallet {
 	#[pallet::hooks]
 	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I> {
 		#[cfg(feature = "try-runtime")]
-		fn try_state(_n: BlockNumberFor<T>) -> Result<(), TryRuntimeError> {
+		fn try_state(
+			_n: BlockNumberFor<T>,
+			_s: frame_support::traits::TryStateSelect,
+		) -> Result<(), TryRuntimeError> {
 			Self::do_try_state()
 		}
 	}

frame/election-provider-multi-phase/src/lib.rs

@@ -886,7 +886,10 @@ pub mod pallet {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn try_state(_n: T::BlockNumber) -> Result<(), TryRuntimeError> {
+		fn try_state(
+			_n: T::BlockNumber,
+			_s: frame_support::traits::TryStateSelect,
+		) -> Result<(), TryRuntimeError> {
 			Self::do_try_state()
 		}
 	}

frame/election-provider-multi-phase/src/mock.rs

@@ -623,7 +623,8 @@ impl ExtBuilder {
 		#[cfg(feature = "try-runtime")]
 		ext.execute_with(|| {
 			assert_ok!(<MultiPhase as frame_support::traits::Hooks<u64>>::try_state(
-				System::block_number()
+				System::block_number(),
+				frame_support::traits::TryStateSelect::All
 			));
 		});
 	}

frame/elections-phragmen/src/lib.rs

@@ -330,7 +330,10 @@ pub mod pallet {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn try_state(_n: T::BlockNumber) -> Result<(), TryRuntimeError> {
+		fn try_state(
+			_n: T::BlockNumber,
+			_s: frame_support::traits::TryStateSelect,
+		) -> Result<(), TryRuntimeError> {
 			Self::do_try_state()
 		}
 	}
@@ -1527,7 +1530,8 @@ mod tests {
 			#[cfg(feature = "try-runtime")]
 			ext.execute_with(|| {
 				assert_ok!(<Elections as frame_support::traits::Hooks<u64>>::try_state(
-					System::block_number()
+					System::block_number(),
+					frame_support::traits::TryStateSelect::All
 				));
 			});
 		}

frame/executive/src/lib.rs

@@ -346,7 +346,7 @@ where
 			let _guard = frame_support::StorageNoopGuard::default();
 			<AllPalletsWithSystem as frame_support::traits::TryState<System::BlockNumber>>::try_state(
 				frame_system::Pallet::<System>::block_number(),
-				frame_try_runtime::TryStateSelect::All,
+				if checks.fast_try_state() { frame_try_runtime::TryStateSelect::Fast } else { frame_try_runtime::TryStateSelect::All },
 			)?;
 		}
 
@@ -359,7 +359,7 @@ where
 			let _guard = frame_support::StorageNoopGuard::default();
 			<AllPalletsWithSystem as frame_support::traits::TryState<System::BlockNumber>>::try_state(
 				frame_system::Pallet::<System>::block_number(),
-				frame_try_runtime::TryStateSelect::All,
+				if checks.fast_try_state() { frame_try_runtime::TryStateSelect::Fast } else { frame_try_runtime::TryStateSelect::All },
 			)?;
 		}
 

frame/fast-unstake/src/lib.rs

@@ -154,6 +154,7 @@ pub mod pallet {
 		pallet_prelude::*,
 		traits::{Defensive, ReservableCurrency, StorageVersion},
 	};
+
 	use frame_system::pallet_prelude::*;
 	use sp_runtime::{traits::Zero, DispatchResult};
 	use sp_staking::{EraIndex, StakingInterface};
@@ -295,7 +296,10 @@ pub mod pallet {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn try_state(_n: T::BlockNumber) -> Result<(), TryRuntimeError> {
+		fn try_state(
+			_n: T::BlockNumber,
+			_: frame_support::traits::TryStateSelect,
+		) -> Result<(), TryRuntimeError> {
 			// ensure that the value of `ErasToCheckPerBlock` is less than
 			// `T::MaxErasToCheckPerBlock`.
 			ensure!(

frame/nomination-pools/src/lib.rs

@@ -2629,7 +2629,10 @@ pub mod pallet {
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		#[cfg(feature = "try-runtime")]
-		fn try_state(_n: BlockNumberFor<T>) -> Result<(), TryRuntimeError> {
+		fn try_state(
+			_n: BlockNumberFor<T>,
+			_: frame_support::traits::TryStateSelect,
+		) -> Result<(), TryRuntimeError> {
 			Self::do_try_state(u8::MAX)
 		}
 

frame/nomination-pools/src/migration.rs

@@ -109,7 +109,10 @@ pub mod v1 {
 				Pallet::<T>::on_chain_storage_version() == 1,
 				"The onchain version must be updated after the migration."
 			);
-			Pallet::<T>::try_state(frame_system::Pallet::<T>::block_number())?;
+			Pallet::<T>::try_state(
+				frame_system::Pallet::<T>::block_number(),
+				frame_support::traits::TryStateSelect::All,
+			)?;
 			Ok(())
 		}
 	}

frame/staking/src/mock.rs

@@ -550,7 +550,7 @@ impl ExtBuilder {
 		let mut ext = self.build();
 		ext.execute_with(test);
 		ext.execute_with(|| {
-			Staking::do_try_state(System::block_number()).unwrap();
+			Staking::do_try_state(System::block_number(), false).unwrap();
 		});
 	}
 }

frame/staking/src/pallet/impls.rs

@@ -1718,17 +1718,26 @@ impl<T: Config> StakingInterface for Pallet<T> {
 
 #[cfg(any(test, feature = "try-runtime"))]
 impl<T: Config> Pallet<T> {
-	pub(crate) fn do_try_state(_: BlockNumberFor<T>) -> Result<(), TryRuntimeError> {
+	pub(crate) fn do_try_state(
+		_: BlockNumberFor<T>,
+		fast_mode: bool,
+	) -> Result<(), TryRuntimeError> {
 		ensure!(
 			T::VoterList::iter()
 				.all(|x| <Nominators<T>>::contains_key(&x) || <Validators<T>>::contains_key(&x)),
 			"VoterList contains non-staker"
 		);
 
-		Self::check_nominators()?;
 		Self::check_exposures()?;
 		Self::check_ledgers()?;
-		Self::check_count()
+		Self::check_count()?;
+
+		if !fast_mode {
+			// slow running tests (should be avoided on CI)
+			Self::check_nominators()?;
+		}
+
+		Ok(())
 	}
 
 	fn check_count() -> Result<(), TryRuntimeError> {

frame/staking/src/pallet/mod.rs

@@ -805,8 +805,11 @@ pub mod pallet {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn try_state(n: BlockNumberFor<T>) -> Result<(), sp_runtime::TryRuntimeError> {
-			Self::do_try_state(n)
+		fn try_state(
+			n: BlockNumberFor<T>,
+			select: frame_support::traits::TryStateSelect,
+		) -> Result<(), sp_runtime::TryRuntimeError> {
+			Self::do_try_state(n, select == frame_support::traits::TryStateSelect::Fast)
 		}
 	}
 

frame/support/procedural/src/pallet/expand/hooks.rs

@@ -267,14 +267,14 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 		{
 			fn try_state(
 				n: <T as #frame_system::Config>::BlockNumber,
-				_s: #frame_support::traits::TryStateSelect
+				s: #frame_support::traits::TryStateSelect
 			) -> Result<(), #frame_support::sp_runtime::TryRuntimeError> {
 				#log_try_state
 				<
 					Self as #frame_support::traits::Hooks<
 						<T as #frame_system::Config>::BlockNumber
 					>
-				>::try_state(n)
+				>::try_state(n, s)
 			}
 		}
 	)

frame/support/src/traits/hooks.rs

@@ -17,6 +17,8 @@
 
 //! Traits for hooking tasks to events in a blockchain's lifecycle.
 
+#[cfg(feature = "try-runtime")]
+use super::try_runtime::Select as TryStateSelect;
 use crate::weights::Weight;
 use impl_trait_for_tuples::impl_for_tuples;
 use sp_runtime::traits::AtLeast32BitUnsigned;
@@ -306,9 +308,12 @@ pub trait Hooks<BlockNumber> {
 	/// It should focus on certain checks to ensure that the state is sensible. This is never
 	/// executed in a consensus code-path, therefore it can consume as much weight as it needs.
 	///
+	/// Takes the block number and `TryStateSelect`as a parameter. The `TryStateSelect` is used to
+	/// select which state tests to execute.
+	///
 	/// This hook should not alter any storage.
 	#[cfg(feature = "try-runtime")]
-	fn try_state(_n: BlockNumber) -> Result<(), TryRuntimeError> {
+	fn try_state(_n: BlockNumber, _s: TryStateSelect) -> Result<(), TryRuntimeError> {
 		Ok(())
 	}
 

frame/support/src/traits/try_runtime.rs

@@ -23,7 +23,7 @@ use sp_runtime::TryRuntimeError;
 use sp_std::prelude::*;
 
 /// Which state tests to execute.
-#[derive(codec::Encode, codec::Decode, Clone, scale_info::TypeInfo)]
+#[derive(codec::Encode, codec::Decode, Clone, Eq, PartialEq, scale_info::TypeInfo)]
 pub enum Select {
 	/// None of them.
 	None,
@@ -35,6 +35,10 @@ pub enum Select {
 	///
 	/// Pallet names are obtained from [`super::PalletInfoAccess`].
 	Only(Vec<Vec<u8>>),
+	/// Run only fast running tests for all pallets.
+	///
+	/// Optimal mode for CI. Avoids long running tests. Each pallet chooses what it considers fast.
+	Fast,
 }
 
 impl Default for Select {
@@ -56,6 +60,7 @@ impl sp_std::fmt::Debug for Select {
 			),
 			Select::All => write!(f, "All"),
 			Select::None => write!(f, "None"),
+			Select::Fast => write!(f, "Fast"),
 		}
 	}
 }
@@ -64,9 +69,10 @@ impl sp_std::fmt::Debug for Select {
 impl sp_std::str::FromStr for Select {
 	type Err = &'static str;
 	fn from_str(s: &str) -> Result<Self, Self::Err> {
-		match s {
-			"all" | "All" => Ok(Select::All),
-			"none" | "None" => Ok(Select::None),
+		match s.to_lowercase().as_ref() {
+			"fast" => Ok(Select::Fast),
+			"all" => Ok(Select::All),
+			"none" => Ok(Select::None),
 			_ =>
 				if s.starts_with("rr-") {
 					let count = s
@@ -89,21 +95,33 @@ pub enum UpgradeCheckSelect {
 	None,
 	/// Run the `try_state`, `pre_upgrade` and `post_upgrade` checks.
 	All,
+	/// Run all fast running `try_state` checks, `pre_upgrade` and `post_upgrade` checks.
+	///
+	/// Skips long running try_state checks.
+	FastAll,
 	/// Run the `pre_upgrade` and `post_upgrade` checks.
 	PreAndPost,
 	/// Run the `try_state` checks.
 	TryState,
+	/// Run all fast running `try_state` checks.
+	///
+	/// Skips long running try_state checks.
+	FastTryState,
 }
 
 impl UpgradeCheckSelect {
 	/// Whether the pre- and post-upgrade checks are selected.
 	pub fn pre_and_post(&self) -> bool {
-		matches!(self, Self::All | Self::PreAndPost)
+		matches!(self, Self::All | Self::FastAll | Self::PreAndPost)
 	}
 
 	/// Whether the try-state checks are selected.
 	pub fn try_state(&self) -> bool {
-		matches!(self, Self::All | Self::TryState)
+		matches!(self, Self::All | Self::TryState | Self::FastAll | Self::FastTryState)
+	}
+
+	pub fn fast_try_state(&self) -> bool {
+		matches!(self, Self::FastAll | Self::FastTryState)
 	}
 }
 
@@ -115,8 +133,10 @@ impl core::str::FromStr for UpgradeCheckSelect {
 		match s.to_lowercase().as_str() {
 			"none" => Ok(Self::None),
 			"all" => Ok(Self::All),
+			"fast-all" => Ok(Self::FastAll),
 			"pre-and-post" => Ok(Self::PreAndPost),
 			"try-state" => Ok(Self::TryState),
+			"fast-try-state" => Ok(Self::FastTryState),
 			_ => Err("Invalid CheckSelector"),
 		}
 	}
@@ -143,7 +163,7 @@ impl<BlockNumber: Clone + sp_std::fmt::Debug + AtLeast32BitUnsigned> TryState<Bl
 	fn try_state(n: BlockNumber, targets: Select) -> Result<(), TryRuntimeError> {
 		match targets {
 			Select::None => Ok(()),
-			Select::All => {
+			Select::All | Select::Fast => {
 				let mut result = Ok(());
 				for_tuples!( #( result = result.and(Tuple::try_state(n.clone(), targets.clone())); )* );
 				result

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

@@ -39,6 +39,7 @@ pub struct ExecuteBlockCmd {
 	///
 	/// Expected values:
 	/// - `all`
+	/// - `fast`
 	/// - `none`
 	/// - A comma separated list of pallets, as per pallet names in `construct_runtime!()` (e.g.
 	///   `Staking, System`).

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

@@ -48,6 +48,7 @@ pub struct FollowChainCmd {
 	///
 	/// Expected values:
 	/// - `all`
+	/// - `fast`
 	/// - `none`
 	/// - A comma separated list of pallets, as per pallet names in `construct_runtime!()` (e.g.
 	///   `Staking, System`).

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

@@ -34,8 +34,10 @@ pub struct OnRuntimeUpgradeCmd {
 	///
 	/// - `none`: Perform no checks (default when the arg is not present).
 	/// - `all`: Perform all checks (default when the arg is present).
+	/// - `fast-all`: Perform fast running state, pre- and post-upgrade checks.
 	/// - `pre-and-post`: Perform pre- and post-upgrade checks.
 	/// - `try-state`: Perform the try-state checks.
+	/// - `fast-try-state`: Perform fast running state checks.
 	///
 	/// Performing any checks will potentially invalidate the measured PoV/Weight.
 	// NOTE: The clap attributes make it backwards compatible with the previous `--checks` flag.

utils/frame/try-runtime/cli/src/lib.rs

@@ -268,6 +268,18 @@
 //!     --at 0xa1b16c1efd889a9f17375ec4dd5c1b4351a2be17fa069564fced10d23b9b3836
 //! ```
 //!
+//! * Run migrations with pre- and post- upgrade checks and fast running try-state checks.
+//!
+//! ```bash
+//! ./substrate-try-runtime \
+//!     try-runtime \
+//!     --runtime runtime-try-runtime.wasm \
+//!     -lruntime=debug \
+//!     on-runtime-upgrade \
+//!     --checks=fast-all \
+//!     live --uri ws://localhost:9999
+//! ```
+//!
 //! * Executing the same command with the [`Runtime::Existing`] will fail because the existing
 //!   runtime, stored onchain in `substrate` binary that we compiled earlier does not have
 //!   `try-runtime` feature!