feat: add runtime-wide try-state hooks

polkadot/xcm/xcm-simulator/fuzzer/src/fuzz.rs

@@ -24,7 +24,7 @@ use sp_runtime::{traits::AccountIdConversion, BuildStorage};
 use xcm_simulator::{decl_test_network, decl_test_parachain, decl_test_relay_chain, TestExt};
 
 #[cfg(feature = "try-runtime")]
-use frame_support::traits::{TryState, TryStateSelect::All};
+use frame_support::traits::TryStateLogic;
 use frame_support::{assert_ok, traits::IntegrityTest};
 use xcm::{latest::prelude::*, MAX_XCM_DECODE_DEPTH};
 

substrate/frame/delegated-staking/src/mock.rs

@@ -242,9 +242,8 @@ impl ExtBuilder {
 		ext.execute_with(test);
 		ext.execute_with(|| {
 			#[cfg(feature = "try-runtime")]
-			<AllPalletsWithSystem as frame_support::traits::TryState<u64>>::try_state(
+			<AllPalletsWithSystem as frame_support::traits::TryStateLogic<u64>>::try_state(
 				frame_system::Pallet::<Runtime>::block_number(),
-				frame_support::traits::TryStateSelect::All,
 			)
 			.unwrap();
 			#[cfg(not(feature = "try-runtime"))]

substrate/frame/executive/README.md

@@ -53,7 +53,7 @@ generic parameter. The custom logic will be called before the on runtime upgrade
 struct CustomOnRuntimeUpgrade;
 impl frame_support::traits::OnRuntimeUpgrade for CustomOnRuntimeUpgrade {
     fn on_runtime_upgrade() -> frame_support::weights::Weight {
-        // Do whatever you want.
+        // Do whatever you want besides modifying storage.
         frame_support::weights::Weight::zero()
     }
 }

substrate/frame/executive/src/lib.rs

@@ -86,11 +86,18 @@
 //! done by setting an optional generic parameter. The custom logic will be called before
 //! the on runtime upgrade logic of all modules is called.
 //!
+//! ### Custom `TryState` logic
+//!
+//! You can add custom logic that should be called in your runtime on a try-state test. This is
+//! done by setting an optional generic parameter after the `OnRuntimeUpgrade` type. The custom
+//! logic will be called before the `try-state` logic of all modules is called.
+//!
 //! ```
 //! # use sp_runtime::generic;
 //! # use frame_executive as executive;
 //! # pub struct UncheckedExtrinsic {};
 //! # pub struct Header {};
+//! # pub type BlockNumber = u32;
 //! # type Context = frame_system::ChainContext<Runtime>;
 //! # pub type Block = generic::Block<Header, UncheckedExtrinsic>;
 //! # pub type Balances = u64;
@@ -115,7 +122,23 @@
 //!     }
 //! }
 //!
-//! pub type Executive = executive::Executive<Runtime, Block, Context, Runtime, AllPalletsWithSystem, CustomOnRuntimeUpgrade>;
+//! struct CustomTryState;
+//!
+//! #[cfg(feature = "try-runtime")]
+//! impl frame_support::traits::IdentifiableTryStateLogic<BlockNumber> for CustomTryState {
+//!
+//! 	fn try_state(_: BlockNumber) -> Result<(), sp_runtime::TryRuntimeError> {
+//!         // Do whatever you want.
+//!         Ok(())
+//! 	}
+//!
+//! 	fn matches_id(id: &[u8]) -> bool {
+//! 		// Used to decide whether to include this or not when `Select::Only` is specified.
+//!         id == b"my_custom_try_state"
+//! 	}
+//! }
+//!
+//! pub type Executive = executive::Executive<Runtime, Block, Context, Runtime, AllPalletsWithSystem, CustomOnRuntimeUpgrade, CustomTryState>;
 //! ```
 
 #[cfg(doc)]
@@ -179,7 +202,10 @@ use sp_std::{marker::PhantomData, prelude::*};
 #[cfg(feature = "try-runtime")]
 use ::{
 	frame_support::{
-		traits::{TryDecodeEntireStorage, TryDecodeEntireStorageError, TryState},
+		traits::{
+			IdentifiableTryStateLogic, TryDecodeEntireStorage, TryDecodeEntireStorageError,
+			TryState,
+		},
 		StorageNoopGuard,
 	},
 	frame_try_runtime::{TryStateSelect, UpgradeCheckSelect},
@@ -205,13 +231,16 @@ pub type OriginOf<E, C> = <CallOf<E, C> as Dispatchable>::RuntimeOrigin;
 ///   used to call hooks e.g. `on_initialize`.
 /// - `OnRuntimeUpgrade`: Custom logic that should be called after a runtime upgrade. Modules are
 ///   already called by `AllPalletsWithSystem`. It will be called before all modules will be called.
+/// - `TryState`: Custom logic that should be called when running `try-state` tests. Modules are
+///   already called by `AllPalletsWithSystem`. It will be called before all modules will be called.
 pub struct Executive<
 	System,
 	Block,
 	Context,
 	UnsignedValidator,
 	AllPalletsWithSystem,
 	OnRuntimeUpgrade = (),
+	TryState = (),
 >(
 	PhantomData<(
 		System,
@@ -220,6 +249,7 @@ pub struct Executive<
 		UnsignedValidator,
 		AllPalletsWithSystem,
 		OnRuntimeUpgrade,
+		TryState,
 	)>,
 );
 
@@ -239,9 +269,17 @@ impl<
 			+ OffchainWorker<BlockNumberFor<System>>
 			+ OnPoll<BlockNumberFor<System>>,
 		COnRuntimeUpgrade: OnRuntimeUpgrade,
+		CTryState,
 	> ExecuteBlock<Block>
-	for Executive<System, Block, Context, UnsignedValidator, AllPalletsWithSystem, COnRuntimeUpgrade>
-where
+	for Executive<
+		System,
+		Block,
+		Context,
+		UnsignedValidator,
+		AllPalletsWithSystem,
+		COnRuntimeUpgrade,
+		CTryState,
+	> where
 	Block::Extrinsic: Checkable<Context> + Codec,
 	CheckedOf<Block::Extrinsic, Context>: Applyable + GetDispatchInfo,
 	CallOf<Block::Extrinsic, Context>:
@@ -277,11 +315,20 @@ impl<
 			+ OnFinalize<BlockNumberFor<System>>
 			+ OffchainWorker<BlockNumberFor<System>>
 			+ OnPoll<BlockNumberFor<System>>
-			+ TryState<BlockNumberFor<System>>
+			+ IdentifiableTryStateLogic<BlockNumberFor<System>>
 			+ TryDecodeEntireStorage,
 		COnRuntimeUpgrade: OnRuntimeUpgrade,
-	> Executive<System, Block, Context, UnsignedValidator, AllPalletsWithSystem, COnRuntimeUpgrade>
-where
+		CTryState: IdentifiableTryStateLogic<BlockNumberFor<System>>,
+	>
+	Executive<
+		System,
+		Block,
+		Context,
+		UnsignedValidator,
+		AllPalletsWithSystem,
+		COnRuntimeUpgrade,
+		CTryState,
+	> where
 	Block::Extrinsic: Checkable<Context> + Codec,
 	CheckedOf<Block::Extrinsic, Context>: Applyable + GetDispatchInfo,
 	CallOf<Block::Extrinsic, Context>:
@@ -377,9 +424,10 @@ where
 
 		// run the try-state checks of all pallets, ensuring they don't alter any state.
 		let _guard = frame_support::StorageNoopGuard::default();
-		<AllPalletsWithSystem as frame_support::traits::TryState<
-			BlockNumberFor<System>,
-		>>::try_state(*header.number(), select.clone())
+		<(CTryState, AllPalletsWithSystem) as TryState<BlockNumberFor<System>>>::try_state(
+			*header.number(),
+			select.clone(),
+		)
 		.map_err(|e| {
 			log::error!(target: LOG_TARGET, "failure: {:?}", e);
 			e
@@ -456,7 +504,7 @@ where
 
 		// Check all storage invariants:
 		if checks.try_state() {
-			AllPalletsWithSystem::try_state(
+			<(CTryState, AllPalletsWithSystem) as TryState<BlockNumberFor<System>>>::try_state(
 				frame_system::Pallet::<System>::block_number(),
 				TryStateSelect::All,
 			)?;
@@ -514,8 +562,17 @@ impl<
 			+ OffchainWorker<BlockNumberFor<System>>
 			+ OnPoll<BlockNumberFor<System>>,
 		COnRuntimeUpgrade: OnRuntimeUpgrade,
-	> Executive<System, Block, Context, UnsignedValidator, AllPalletsWithSystem, COnRuntimeUpgrade>
-where
+		CTryState,
+	>
+	Executive<
+		System,
+		Block,
+		Context,
+		UnsignedValidator,
+		AllPalletsWithSystem,
+		COnRuntimeUpgrade,
+		CTryState,
+	> where
 	Block::Extrinsic: Checkable<Context> + Codec,
 	CheckedOf<Block::Extrinsic, Context>: Applyable + GetDispatchInfo,
 	CallOf<Block::Extrinsic, Context>:

substrate/frame/executive/src/tests.rs

@@ -44,6 +44,43 @@ use pallet_balances::Call as BalancesCall;
 
 const TEST_KEY: &[u8] = b":test:key:";
 
+#[cfg(feature = "try-runtime")]
+mod try_runtime {
+	use frame_support::{
+		parameter_types,
+		traits::{IdentifiableTryStateLogic, TryStateLogic},
+	};
+	use sp_runtime::traits::AtLeast32BitUnsigned;
+
+	// Will contain `true` when the custom runtime logic is called.
+	parameter_types! {
+		pub(super) static CanaryFlag: bool = false;
+	}
+
+	const CUSTOM_TRY_STATE_ID: &[u8] = b"custom_try_state";
+
+	pub(super) struct CustomTryState;
+
+	impl<BlockNumber> TryStateLogic<BlockNumber> for CustomTryState
+	where
+		BlockNumber: Clone + sp_std::fmt::Debug + AtLeast32BitUnsigned,
+	{
+		fn try_state(_: BlockNumber) -> Result<(), sp_runtime::TryRuntimeError> {
+			CanaryFlag::set(true);
+			Ok(())
+		}
+	}
+
+	impl<BlockNumber> IdentifiableTryStateLogic<BlockNumber> for CustomTryState
+	where
+		BlockNumber: Clone + sp_std::fmt::Debug + AtLeast32BitUnsigned,
+	{
+		fn matches_id(id: &[u8]) -> bool {
+			id == CUSTOM_TRY_STATE_ID
+		}
+	}
+}
+
 #[frame_support::pallet(dev_mode)]
 mod custom {
 	use super::*;
@@ -78,6 +115,16 @@ mod custom {
 		fn offchain_worker(n: BlockNumberFor<T>) {
 			assert_eq!(BlockNumberFor::<T>::from(1u32), n);
 		}
+
+		// Verify that `CustomTryState` has been called before.
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: BlockNumberFor<T>) -> Result<(), sp_runtime::TryRuntimeError> {
+			assert!(
+				try_runtime::CanaryFlag::get(),
+				"Custom `try-runtime` did not run before pallet `try-runtime`."
+			);
+			Ok(())
+		}
 	}
 
 	#[pallet::call]
@@ -397,13 +444,19 @@ impl OnRuntimeUpgrade for CustomOnRuntimeUpgrade {
 	}
 }
 
+#[cfg(not(feature = "try-runtime"))]
+type CustomOnTryState = ();
+#[cfg(feature = "try-runtime")]
+type CustomOnTryState = try_runtime::CustomTryState;
+
 type Executive = super::Executive<
 	Runtime,
 	Block<TestXt>,
 	ChainContext<Runtime>,
 	Runtime,
 	AllPalletsWithSystem,
 	CustomOnRuntimeUpgrade,
+	CustomOnTryState,
 >;
 
 parameter_types! {
@@ -523,9 +576,12 @@ fn new_test_ext(balance_factor: Balance) -> sp_io::TestExternalities {
 	ext.execute_with(|| {
 		SystemCallbacksCalled::set(0);
 	});
+
+	#[cfg(feature = "try-runtime")]
+	try_runtime::CanaryFlag::set(false);
+
 	ext
 }
-
 fn new_test_ext_v0(balance_factor: Balance) -> sp_io::TestExternalities {
 	let mut t = frame_system::GenesisConfig::<Runtime>::default().build_storage().unwrap();
 	pallet_balances::GenesisConfig::<Runtime> { balances: vec![(1, 111 * balance_factor)] }

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

@@ -18,6 +18,8 @@
 use crate::pallet::Def;
 
 /// * implement the individual traits using the Hooks trait
+/// * implement the `TryStateLogic` and `IdentifiableTryStateLogic` traits, that are strictly
+///   dependent on the `TryState` hook
 pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 	let (where_clause, span, has_runtime_upgrade) = match def.hooks.as_ref() {
 		Some(hooks) => {
@@ -336,5 +338,27 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 				})
 			}
 		}
+
+		// Implement `TryStateLogic<BlockNumber>` for `Pallet`
+		#[cfg(feature = "try-runtime")]
+		impl<#type_impl_gen> #frame_support::traits::TryStateLogic<#frame_system::pallet_prelude::BlockNumberFor<T>>
+			for #pallet_ident<#type_use_gen>
+			#where_clause
+		{
+			fn try_state(n: frame_system::pallet_prelude::BlockNumberFor<T>) -> Result<(), #frame_support::sp_runtime::TryRuntimeError> {
+				<Self as #frame_support::traits::TryState<#frame_system::pallet_prelude::BlockNumberFor::<T>>>::try_state(n, #frame_support::traits::TryStateSelect::All)
+			}
+		}
+
+		// Implement `IdentifiableTryStateLogic<BlockNumber>` for `Pallet`
+		#[cfg(feature = "try-runtime")]
+		impl<#type_impl_gen> #frame_support::traits::IdentifiableTryStateLogic<frame_system::pallet_prelude::BlockNumberFor<T>>
+			for #pallet_ident<#type_use_gen>
+			#where_clause
+		{
+			fn matches_id(id: &[u8]) -> bool {
+				<Self as #frame_support::pallet_prelude::PalletInfoAccess>::name().as_bytes() == id
+			}
+		}
 	)
 }

substrate/frame/support/src/traits.rs

@@ -134,6 +134,6 @@ pub use tasks::Task;
 mod try_runtime;
 #[cfg(feature = "try-runtime")]
 pub use try_runtime::{
-	Select as TryStateSelect, TryDecodeEntireStorage, TryDecodeEntireStorageError, TryState,
-	UpgradeCheckSelect,
+	IdentifiableTryStateLogic, Select as TryStateSelect, TryDecodeEntireStorage,
+	TryDecodeEntireStorageError, TryState, TryStateLogic, UpgradeCheckSelect,
 };