Runtime State Test + Integration with try-runtime (#10174)

* add missing version to dependencies

* Huh

* add features more

* more fixing

* last touches

* it all finally works

* remove some feature gates

* remove unused

* fix old macro

* make it work again

* fmt

* remove unused import

* ".git/.scripts/fmt.sh" 1

* Cleanup more

* fix and rename everything

* a few clippy fixes

* Add try-runtime feature

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

* small fixes

* fmt

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

* fix build

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

Co-authored-by: David <dvdplm@gmail.com>

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

Co-authored-by: David <dvdplm@gmail.com>

* address all review comments

* fix typos

* revert spec change

* last touches

* update docs

* fmt

* remove some debug_assertions

* fmt

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: command-bot <>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: David <dvdplm@gmail.com>

substrate/frame/support/src/dispatch.rs

@@ -1549,6 +1549,35 @@ macro_rules! decl_module {
 		{}
 	};
 
+	(@impl_try_state_default
+		{ $system:ident }
+		$module:ident<$trait_instance:ident: $trait_name:ident$(<I>, $instance:ident: $instantiable:path)?>;
+		{ $( $other_where_bounds:tt )* }
+	) => {
+		#[cfg(feature = "try-runtime")]
+		impl<$trait_instance: $system::Config + $trait_name$(<I>, $instance: $instantiable)?>
+			$crate::traits::TryState<<$trait_instance as $system::Config>::BlockNumber>
+			for $module<$trait_instance$(, $instance)?> where $( $other_where_bounds )*
+		{
+			fn try_state(
+				_: <$trait_instance as $system::Config>::BlockNumber,
+				_: $crate::traits::TryStateSelect,
+			) -> Result<(), &'static str> {
+				let pallet_name = <<
+					$trait_instance
+					as
+					$system::Config
+				>::PalletInfo as $crate::traits::PalletInfo>::name::<Self>().unwrap_or("<unknown pallet name>");
+				$crate::log::debug!(
+					target: $crate::LOG_TARGET,
+					"⚠️ pallet {} cannot have try-state because it is using decl_module!",
+					pallet_name,
+				);
+				Ok(())
+			}
+		}
+	};
+
 	(@impl_on_runtime_upgrade
 		{ $system:ident }
 		$module:ident<$trait_instance:ident: $trait_name:ident$(<I>, $instance:ident: $instantiable:path)?>;
@@ -2026,6 +2055,13 @@ macro_rules! decl_module {
 			$( $on_initialize )*
 		}
 
+		$crate::decl_module! {
+			@impl_try_state_default
+			{ $system }
+			$mod_type<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?>;
+			{ $( $other_where_bounds )* }
+		}
+
 		$crate::decl_module! {
 			@impl_on_runtime_upgrade
 			{ $system }

substrate/frame/support/src/traits.rs

@@ -84,8 +84,6 @@ pub use hooks::{
 	Hooks, IntegrityTest, OnFinalize, OnGenesis, OnIdle, OnInitialize, OnRuntimeUpgrade,
 	OnTimestampSet,
 };
-#[cfg(feature = "try-runtime")]
-pub use hooks::{OnRuntimeUpgradeHelpersExt, ON_RUNTIME_UPGRADE_PREFIX};
 
 pub mod schedule;
 mod storage;
@@ -106,3 +104,8 @@ pub use voting::{
 	ClassCountOf, CurrencyToVote, PollStatus, Polling, SaturatingCurrencyToVote,
 	U128CurrencyToVote, VoteTally,
 };
+
+#[cfg(feature = "try-runtime")]
+mod try_runtime;
+#[cfg(feature = "try-runtime")]
+pub use try_runtime::{OnRuntimeUpgradeHelpersExt, Select as TryStateSelect, TryState};

substrate/frame/support/src/traits/hooks.rs

@@ -20,6 +20,7 @@
 use crate::weights::Weight;
 use impl_trait_for_tuples::impl_for_tuples;
 use sp_runtime::traits::AtLeast32BitUnsigned;
+use sp_std::prelude::*;
 
 /// The block initialization trait.
 ///
@@ -93,9 +94,9 @@ impl<BlockNumber: Copy + AtLeast32BitUnsigned> OnIdle<BlockNumber> for Tuple {
 		let start_index = start_index.try_into().ok().expect(
 			"`start_index % len` always fits into `usize`, because `len` can be in maximum `usize::MAX`; qed"
 		);
-		for on_idle in on_idle_functions.iter().cycle().skip(start_index).take(len) {
+		for on_idle_fn in on_idle_functions.iter().cycle().skip(start_index).take(len) {
 			let adjusted_remaining_weight = remaining_weight.saturating_sub(weight);
-			weight = weight.saturating_add(on_idle(n, adjusted_remaining_weight));
+			weight = weight.saturating_add(on_idle_fn(n, adjusted_remaining_weight));
 		}
 		weight
 	}
@@ -114,47 +115,6 @@ pub trait OnGenesis {
 	fn on_genesis() {}
 }
 
-/// Prefix to be used (optionally) for implementing [`OnRuntimeUpgradeHelpersExt::storage_key`].
-#[cfg(feature = "try-runtime")]
-pub const ON_RUNTIME_UPGRADE_PREFIX: &[u8] = b"__ON_RUNTIME_UPGRADE__";
-
-/// Some helper functions for [`OnRuntimeUpgrade`] during `try-runtime` testing.
-#[cfg(feature = "try-runtime")]
-pub trait OnRuntimeUpgradeHelpersExt {
-	/// Generate a storage key unique to this runtime upgrade.
-	///
-	/// This can be used to communicate data from pre-upgrade to post-upgrade state and check
-	/// them. See [`Self::set_temp_storage`] and [`Self::get_temp_storage`].
-	#[cfg(feature = "try-runtime")]
-	fn storage_key(ident: &str) -> [u8; 32] {
-		crate::storage::storage_prefix(ON_RUNTIME_UPGRADE_PREFIX, ident.as_bytes())
-	}
-
-	/// Get temporary storage data written by [`Self::set_temp_storage`].
-	///
-	/// Returns `None` if either the data is unavailable or un-decodable.
-	///
-	/// A `at` storage identifier must be provided to indicate where the storage is being read from.
-	#[cfg(feature = "try-runtime")]
-	fn get_temp_storage<T: codec::Decode>(at: &str) -> Option<T> {
-		sp_io::storage::get(&Self::storage_key(at))
-			.and_then(|bytes| codec::Decode::decode(&mut &*bytes).ok())
-	}
-
-	/// Write some temporary data to a specific storage that can be read (potentially in
-	/// post-upgrade hook) via [`Self::get_temp_storage`].
-	///
-	/// A `at` storage identifier must be provided to indicate where the storage is being written
-	/// to.
-	#[cfg(feature = "try-runtime")]
-	fn set_temp_storage<T: codec::Encode>(data: T, at: &str) {
-		sp_io::storage::set(&Self::storage_key(at), &data.encode());
-	}
-}
-
-#[cfg(feature = "try-runtime")]
-impl<U: OnRuntimeUpgrade> OnRuntimeUpgradeHelpersExt for U {}
-
 /// The runtime upgrade trait.
 ///
 /// Implementing this lets you express what should happen when the runtime upgrades,
@@ -272,6 +232,15 @@ pub trait Hooks<BlockNumber> {
 		Weight::new()
 	}
 
+	/// Execute the sanity checks of this pallet, per block.
+	///
+	/// 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.
+	#[cfg(feature = "try-runtime")]
+	fn try_state(_n: BlockNumber) -> Result<(), &'static str> {
+		Ok(())
+	}
+
 	/// Execute some pre-checks prior to a runtime upgrade.
 	///
 	/// This hook is never meant to be executed on-chain but is meant to be used by testing tools.

substrate/frame/support/src/traits/try_runtime.rs

@@ -0,0 +1,174 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Try-runtime specific traits and types.
+
+use super::*;
+use impl_trait_for_tuples::impl_for_tuples;
+use sp_arithmetic::traits::AtLeast32BitUnsigned;
+use sp_std::prelude::*;
+
+/// Prefix to be used (optionally) for implementing [`OnRuntimeUpgradeHelpersExt::storage_key`].
+const ON_RUNTIME_UPGRADE_PREFIX: &[u8] = b"__ON_RUNTIME_UPGRADE__";
+
+/// Some helper functions for [`OnRuntimeUpgrade`] during `try-runtime` testing.
+pub trait OnRuntimeUpgradeHelpersExt {
+	/// Generate a storage key unique to this runtime upgrade.
+	///
+	/// This can be used to communicate data from pre-upgrade to post-upgrade state and check
+	/// them. See [`Self::set_temp_storage`] and [`Self::get_temp_storage`].
+	fn storage_key(ident: &str) -> [u8; 32] {
+		crate::storage::storage_prefix(ON_RUNTIME_UPGRADE_PREFIX, ident.as_bytes())
+	}
+
+	/// Get temporary storage data written by [`Self::set_temp_storage`].
+	///
+	/// Returns `None` if either the data is unavailable or un-decodable.
+	///
+	/// A `at` storage identifier must be provided to indicate where the storage is being read from.
+	fn get_temp_storage<T: codec::Decode>(at: &str) -> Option<T> {
+		sp_io::storage::get(&Self::storage_key(at))
+			.and_then(|bytes| codec::Decode::decode(&mut &*bytes).ok())
+	}
+
+	/// Write some temporary data to a specific storage that can be read (potentially in
+	/// post-upgrade hook) via [`Self::get_temp_storage`].
+	///
+	/// A `at` storage identifier must be provided to indicate where the storage is being written
+	/// to.
+	fn set_temp_storage<T: codec::Encode>(data: T, at: &str) {
+		sp_io::storage::set(&Self::storage_key(at), &data.encode());
+	}
+}
+
+impl<U: OnRuntimeUpgrade> OnRuntimeUpgradeHelpersExt for U {}
+
+// Which state tests to execute.
+#[derive(codec::Encode, codec::Decode, Clone)]
+pub enum Select {
+	/// None of them.
+	None,
+	/// All of them.
+	All,
+	/// Run a fixed number of them in a round robin manner.
+	RoundRobin(u32),
+	/// Run only pallets who's name matches the given list.
+	///
+	/// Pallet names are obtained from [`PalletInfoAccess`].
+	Only(Vec<Vec<u8>>),
+}
+
+impl Default for Select {
+	fn default() -> Self {
+		Select::None
+	}
+}
+
+impl sp_std::fmt::Debug for Select {
+	fn fmt(&self, f: &mut sp_std::fmt::Formatter<'_>) -> sp_std::fmt::Result {
+		match self {
+			Select::RoundRobin(x) => write!(f, "RoundRobin({})", x),
+			Select::Only(x) => write!(
+				f,
+				"Only({:?})",
+				x.iter()
+					.map(|x| sp_std::str::from_utf8(x).unwrap_or("<invalid?>"))
+					.collect::<Vec<_>>(),
+			),
+			Select::All => write!(f, "All"),
+			Select::None => write!(f, "None"),
+		}
+	}
+}
+
+#[cfg(feature = "std")]
+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),
+			_ =>
+				if s.starts_with("rr-") {
+					let count = s
+						.split_once('-')
+						.and_then(|(_, count)| count.parse::<u32>().ok())
+						.ok_or("failed to parse count")?;
+					Ok(Select::RoundRobin(count))
+				} else {
+					let pallets = s.split(',').map(|x| x.as_bytes().to_vec()).collect::<Vec<_>>();
+					Ok(Select::Only(pallets))
+				},
+		}
+	}
+}
+
+/// 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
+/// the storage items of your pallet.
+pub trait TryState<BlockNumber> {
+	/// Execute the state checks.
+	fn try_state(_: BlockNumber, _: Select) -> Result<(), &'static str>;
+}
+
+#[cfg_attr(all(not(feature = "tuples-96"), not(feature = "tuples-128")), impl_for_tuples(64))]
+#[cfg_attr(all(feature = "tuples-96", not(feature = "tuples-128")), impl_for_tuples(96))]
+#[cfg_attr(all(feature = "tuples-128"), impl_for_tuples(128))]
+impl<BlockNumber: Clone + sp_std::fmt::Debug + AtLeast32BitUnsigned> TryState<BlockNumber>
+	for Tuple
+{
+	for_tuples!( where #( Tuple: crate::traits::PalletInfoAccess )* );
+	fn try_state(n: BlockNumber, targets: Select) -> Result<(), &'static str> {
+		match targets {
+			Select::None => Ok(()),
+			Select::All => {
+				let mut result = Ok(());
+				for_tuples!( #( result = result.and(Tuple::try_state(n.clone(), targets.clone())); )* );
+				result
+			},
+			Select::RoundRobin(len) => {
+				let functions: &[fn(BlockNumber, Select) -> Result<(), &'static str>] =
+					&[for_tuples!(#( Tuple::try_state ),*)];
+				let skip = n.clone() % (functions.len() as u32).into();
+				let skip: u32 =
+					skip.try_into().unwrap_or_else(|_| sp_runtime::traits::Bounded::max_value());
+				let mut result = Ok(());
+				for try_state_fn in functions.iter().cycle().skip(skip as usize).take(len as usize)
+				{
+					result = result.and(try_state_fn(n.clone(), targets.clone()));
+				}
+				result
+			},
+			Select::Only(ref pallet_names) => {
+				let try_state_fns: &[(
+					&'static str,
+					fn(BlockNumber, Select) -> Result<(), &'static str>,
+				)] = &[for_tuples!(
+					#( (<Tuple as crate::traits::PalletInfoAccess>::name(), Tuple::try_state) ),*
+				)];
+				let mut result = Ok(());
+				for (name, try_state_fn) in try_state_fns {
+					if pallet_names.iter().any(|n| n == name.as_bytes()) {
+						result = result.and(try_state_fn(n.clone(), targets.clone()));
+					}
+				}
+				result
+			},
+		}
+	}
+}