TryDecodeEntireState check for storage types and pallets

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

@@ -311,6 +311,14 @@ pub type SignedExtra = (
 	pallet_transaction_payment::ChargeTransactionPayment<Runtime>,
 );
 
+/// All migrations of the runtime, aside from the ones declared in the pallets.
+///
+/// This can be a tuple of types, each implementing `OnRuntimeUpgrade`. Add other migration types
+/// before `EnsureStateDecodes` as needed -- this is only for testing, and
+// should come last.
+#[allow(unused_parens)]
+type Migrations = (frame_support::migration::EnsureStateDecodes<AllPalletsWithSystem>);
+
 /// Unchecked extrinsic type as expected by this runtime.
 pub type UncheckedExtrinsic =
 	generic::UncheckedExtrinsic<Address, RuntimeCall, Signature, SignedExtra>;
@@ -323,6 +331,7 @@ pub type Executive = frame_executive::Executive<
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPalletsWithSystem,
+	Migrations,
 >;
 
 #[cfg(feature = "runtime-benchmarks")]

bin/node/runtime/src/lib.rs

@@ -2001,6 +2001,8 @@ type Migrations = (
 	pallet_nomination_pools::migration::v2::MigrateToV2<Runtime>,
 	pallet_alliance::migration::Migration<Runtime>,
 	pallet_contracts::Migration<Runtime>,
+	// This should always be the last migration item.
+	frame_support::storage::migration::EnsureStateDecodes<AllPalletsWithSystem>,
 );
 
 type EventRecord = frame_system::EventRecord<

frame/glutton/src/lib.rs

@@ -21,9 +21,8 @@
 //!
 //! # Glutton Pallet
 //!
-//! Pallet that consumes `ref_time` and `proof_size` of a block. Based on the
-//! `Compute` and `Storage` parameters the pallet consumes the adequate amount
-//! of weight.
+//! Pallet that consumes `ref_time` and `proof_size` of a block. Based on the `Compute` and
+//! `Storage` parameters the pallet consumes the adequate amount of weight.
 
 #![deny(missing_docs)]
 #![cfg_attr(not(feature = "std"), no_std)]

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

@@ -127,11 +127,12 @@ pub fn expand_event(def: &mut Def) -> proc_macro2::TokenStream {
 		let trait_use_gen = &def.trait_use_generics(event.attr_span);
 		let type_impl_gen = &def.type_impl_generics(event.attr_span);
 		let type_use_gen = &def.type_use_generics(event.attr_span);
+		let pallet_ident = &def.pallet_struct.pallet;
 
 		let PalletEventDepositAttr { fn_vis, fn_span, .. } = deposit_event;
 
 		quote::quote_spanned!(*fn_span =>
-			impl<#type_impl_gen> Pallet<#type_use_gen> #completed_where_clause {
+			impl<#type_impl_gen> #pallet_ident<#type_use_gen> #completed_where_clause {
 				#fn_vis fn deposit_event(event: Event<#event_use_gen>) {
 					let event = <
 						<T as Config #trait_use_gen>::RuntimeEvent as

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

@@ -780,12 +780,43 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		)
 	});
 
+	// aggregated where clause of all storage types and the whole pallet.
 	let mut where_clauses = vec![&def.config.where_clause];
 	where_clauses.extend(def.storages.iter().map(|storage| &storage.where_clause));
 	let completed_where_clause = super::merge_where_clauses(&where_clauses);
 	let type_impl_gen = &def.type_impl_generics(proc_macro2::Span::call_site());
 	let type_use_gen = &def.type_use_generics(proc_macro2::Span::call_site());
 
+	let try_decode_entire_state = {
+		let storage_names = def
+			.storages
+			.iter()
+			.filter_map(|storage| {
+				if storage.cfg_attrs.is_empty() {
+					let ident = &storage.ident;
+					let gen = &def.type_use_generics(storage.attr_span);
+					Some(quote::quote_spanned!(storage.attr_span => #ident<#gen> ))
+				} else {
+					None
+				}
+			})
+			.collect::<Vec<_>>();
+
+		quote::quote!(
+			#[cfg(feature = "try-runtime")]
+			impl<#type_impl_gen> #frame_support::traits::TryDecodeEntireStorage
+			for #pallet_ident<#type_use_gen> #completed_where_clause
+			{
+				fn try_decode_entire_state() -> Result<usize, &'static str> {
+					// simply delegate impl to a tuple of all storage items we have.
+					//
+					// NOTE: for now, we have to exclude storage items that are feature gated.
+					<( #( #storage_names ),*) as frame_support::traits::TryDecodeEntireStorage>::try_decode_entire_state()
+				}
+			}
+		)
+	};
+
 	quote::quote!(
 		impl<#type_impl_gen> #pallet_ident<#type_use_gen>
 			#completed_where_clause
@@ -811,5 +842,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		#( #getters )*
 		#( #prefix_structs )*
 		#( #on_empty_structs )*
+
+		#try_decode_entire_state
 	)
 }

frame/support/procedural/src/pallet/parse/storage.rs

@@ -151,7 +151,7 @@ pub enum QueryKind {
 /// `type MyStorage = StorageValue<MyStorageP, u32>`
 /// The keys and values types are parsed in order to get metadata
 pub struct StorageDef {
-	/// The index of error item in pallet module.
+	/// The index of storage item in pallet module.
 	pub index: usize,
 	/// Visibility of the storage type.
 	pub vis: syn::Visibility,

frame/support/src/storage/generator/mod.rs

@@ -24,10 +24,10 @@
 //!
 //! This is internal api and is subject to change.
 
-mod double_map;
+pub(crate) mod double_map;
 pub(crate) mod map;
-mod nmap;
-mod value;
+pub(crate) mod nmap;
+pub(crate) mod value;
 
 pub use double_map::StorageDoubleMap;
 pub use map::StorageMap;

frame/support/src/storage/migration.rs

@@ -385,6 +385,39 @@ pub fn move_prefix(from_prefix: &[u8], to_prefix: &[u8]) {
 	}
 }
 
+/// A phony migration that does nothing, except executing `TryDecodeEntireStorage` on
+/// `post_upgrade`, which implies it is only available if `try-state` feature is used.
+///
+/// This can be used typically in the top level runtime, where `AllPallets` typically comes from
+/// `construct_runtime!`.
+pub struct EnsureStateDecodes<AllPallets>(sp_std::marker::PhantomData<AllPallets>);
+
+#[cfg(not(feature = "try-runtime"))]
+impl<AllPallets> crate::traits::OnRuntimeUpgrade for EnsureStateDecodes<AllPallets> {
+	fn on_runtime_upgrade() -> crate::weights::Weight {
+		Default::default()
+	}
+}
+
+#[cfg(feature = "try-runtime")]
+impl<AllPallets: crate::traits::TryDecodeEntireStorage> crate::traits::OnRuntimeUpgrade
+	for EnsureStateDecodes<AllPallets>
+{
+	fn on_runtime_upgrade() -> sp_weights::Weight {
+		Default::default()
+	}
+
+	fn post_upgrade(_: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
+		let decoded = AllPallets::try_decode_entire_state()?;
+		crate::log::info!(
+			target: crate::LOG_TARGET,
+			"decoded the entire state, total size = {} bytes",
+			decoded
+		);
+		Ok(())
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::{

frame/support/src/storage/types/counted_map.rs

@@ -74,7 +74,11 @@ impl<P: CountedStorageMapInstance, H, K, V, Q, O, M> MapWrapper
 	type Map = StorageMap<P, H, K, V, Q, O, M>;
 }
 
-type CounterFor<P> = StorageValue<<P as CountedStorageMapInstance>::CounterPrefix, u32, ValueQuery>;
+/// The numeric counter type.
+pub type Counter = u32;
+
+type CounterFor<P> =
+	StorageValue<<P as CountedStorageMapInstance>::CounterPrefix, Counter, ValueQuery>;
 
 /// On removal logic for updating counter while draining upon some prefix with
 /// [`crate::storage::PrefixIterator`].
@@ -378,14 +382,14 @@ where
 	/// can be very heavy, so use with caution.
 	///
 	/// Returns the number of items in the map which is used to set the counter.
-	pub fn initialize_counter() -> u32 {
-		let count = Self::iter_values().count() as u32;
+	pub fn initialize_counter() -> Counter {
+		let count = Self::iter_values().count() as Counter;
 		CounterFor::<Prefix>::set(count);
 		count
 	}
 
 	/// Return the count.
-	pub fn count() -> u32 {
+	pub fn count() -> Counter {
 		CounterFor::<Prefix>::get()
 	}
 }
@@ -1162,7 +1166,7 @@ mod test {
 				StorageEntryMetadataIR {
 					name: "counter_for_foo",
 					modifier: StorageEntryModifierIR::Default,
-					ty: StorageEntryTypeIR::Plain(scale_info::meta_type::<u32>()),
+					ty: StorageEntryTypeIR::Plain(scale_info::meta_type::<Counter>()),
 					default: vec![0, 0, 0, 0],
 					docs: if cfg!(feature = "no-metadata-docs") {
 						vec![]

frame/support/src/storage/types/counted_nmap.rs

@@ -71,8 +71,10 @@ impl<P: CountedStorageNMapInstance, K, V, Q, O, M> MapWrapper
 	type Map = StorageNMap<P, K, V, Q, O, M>;
 }
 
+type Counter = super::counted_map::Counter;
+
 type CounterFor<P> =
-	StorageValue<<P as CountedStorageNMapInstance>::CounterPrefix, u32, ValueQuery>;
+	StorageValue<<P as CountedStorageNMapInstance>::CounterPrefix, Counter, ValueQuery>;
 
 /// On removal logic for updating counter while draining upon some prefix with
 /// [`crate::storage::PrefixIterator`].
@@ -429,7 +431,7 @@ where
 	}
 
 	/// Return the count.
-	pub fn count() -> u32 {
+	pub fn count() -> Counter {
 		CounterFor::<Prefix>::get()
 	}
 }

frame/support/src/storage/types/mod.rs

@@ -30,7 +30,7 @@ mod map;
 mod nmap;
 mod value;
 
-pub use counted_map::{CountedStorageMap, CountedStorageMapInstance};
+pub use counted_map::{CountedStorageMap, CountedStorageMapInstance, Counter};
 pub use counted_nmap::{CountedStorageNMap, CountedStorageNMapInstance};
 pub use double_map::StorageDoubleMap;
 pub use key::{

frame/support/src/storage/types/value.rs

@@ -24,7 +24,7 @@ use crate::{
 		types::{OptionQuery, QueryKindTrait, StorageEntryMetadataBuilder},
 		StorageAppend, StorageDecodeLength, StorageTryAppend,
 	},
-	traits::{GetDefault, StorageInfo, StorageInstance},
+	traits::{Get, GetDefault, StorageInfo, StorageInstance},
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
 use sp_arithmetic::traits::SaturatedConversion;
@@ -46,7 +46,7 @@ where
 	Prefix: StorageInstance,
 	Value: FullCodec,
 	QueryKind: QueryKindTrait<Value, OnEmpty>,
-	OnEmpty: crate::traits::Get<QueryKind::Query> + 'static,
+	OnEmpty: Get<QueryKind::Query> + 'static,
 {
 	type Query = QueryKind::Query;
 	fn module_prefix() -> &'static [u8] {

frame/support/src/traits.rs

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