TryDecodeEntireState check for storage types and pallets (#1805)

### This PR is a port of this [PR for
substrate](https://github.com/paritytech/substrate/pull/13013) by
@kianenigma

Add infrastructure needed to have a Pallet::decode_entire_state(), which
makes sure all "typed" storage items defined in the pallet are
decode-able.

This is not enforced in any way at the moment. Teams who wish to
integrate/use this in the try-runtime feature flag should add
frame_support::storage::migration::EnsureStateDecodes as the LAST ITEM
of the runtime's custom migrations, and pass it to frame-executive. This
will make it usable in try-runtime on-runtime-upgrade.

This now catches cases like
https://github.com/paritytech/polkadot-sdk/pull/1969:
```pre
ERROR runtime::executive] failed to decode the value at key: Failed to decode value at key: 0x94eadf0156a8ad5156507773d0471e4ab8ebad86f546c7e0b135a4212aace339. Storage info StorageInfo { pallet_name: Ok("ParaScheduler"), storage_name: Ok("AvailabilityCores"), prefix: Err(Utf8Error { valid_up_to: 0, error_len: Some(1) }), max_values: Some(1), max_size: None }. Raw value: Some("0x0c010101010101")
```

... or:

![image](https://github.com/paritytech/polkadot-sdk/assets/10380170/73052d4f-4da5-4b21-a8dd-b17004e5965e)

Closes #241

---------

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: Liam Aharon <liam.aharon@hotmail.com>

prdoc/pr_1805.prdoc

@@ -0,0 +1,19 @@
+title: Introduce state decoding check after runtime upgrades.
+
+doc:
+  - audience: Core Dev
+    description: |
+      Adds a check to the try-runtime logic that will verify that all pallet on-chain storage still decodes. This can help to spot missing migrations before they become a problem. The check is enabled as soon as the `--checks` option of the `try-runtime` CLI is not `None`.
+
+migrations:
+  db: []
+
+  runtime: []
+
+crates: 
+  - name: frame-support
+    semver: minor
+  - name: frame-support-procedural
+    semver: minor
+
+host_functions: []

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

@@ -313,6 +313,12 @@ 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`.
+#[allow(unused_parens)]
+type Migrations = ();
+
 /// Unchecked extrinsic type as expected by this runtime.
 pub type UncheckedExtrinsic =
 	generic::UncheckedExtrinsic<Address, RuntimeCall, Signature, SignedExtra>;
@@ -325,6 +331,7 @@ pub type Executive = frame_executive::Executive<
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPalletsWithSystem,
+	Migrations,
 >;
 
 #[cfg(feature = "runtime-benchmarks")]

substrate/frame/executive/src/lib.rs

@@ -139,9 +139,15 @@ use sp_runtime::{
 use sp_std::{marker::PhantomData, prelude::*};
 
 #[cfg(feature = "try-runtime")]
-use log;
-#[cfg(feature = "try-runtime")]
-use sp_runtime::TryRuntimeError;
+use ::{
+	frame_support::{
+		traits::{TryDecodeEntireStorage, TryDecodeEntireStorageError, TryState},
+		StorageNoopGuard,
+	},
+	frame_try_runtime::{TryStateSelect, UpgradeCheckSelect},
+	log,
+	sp_runtime::TryRuntimeError,
+};
 
 #[allow(dead_code)]
 const LOG_TARGET: &str = "runtime::executive";
@@ -229,7 +235,8 @@ impl<
 			+ OnIdle<BlockNumberFor<System>>
 			+ OnFinalize<BlockNumberFor<System>>
 			+ OffchainWorker<BlockNumberFor<System>>
-			+ frame_support::traits::TryState<BlockNumberFor<System>>,
+			+ TryState<BlockNumberFor<System>>
+			+ TryDecodeEntireStorage,
 		COnRuntimeUpgrade: OnRuntimeUpgrade,
 	> Executive<System, Block, Context, UnsignedValidator, AllPalletsWithSystem, COnRuntimeUpgrade>
 where
@@ -308,11 +315,15 @@ where
 		let _guard = frame_support::StorageNoopGuard::default();
 		<AllPalletsWithSystem as frame_support::traits::TryState<
 			BlockNumberFor<System>,
-		>>::try_state(*header.number(), select)
+		>>::try_state(*header.number(), select.clone())
 		.map_err(|e| {
 			log::error!(target: LOG_TARGET, "failure: {:?}", e);
 			e
 		})?;
+		if select.any() {
+			let res = AllPalletsWithSystem::try_decode_entire_state();
+			Self::log_decode_result(res)?;
+		}
 		drop(_guard);
 
 		// do some of the checks that would normally happen in `final_checks`, but perhaps skip
@@ -352,26 +363,61 @@ where
 	/// Execute all `OnRuntimeUpgrade` of this runtime.
 	///
 	/// The `checks` param determines whether to execute `pre/post_upgrade` and `try_state` hooks.
-	pub fn try_runtime_upgrade(
-		checks: frame_try_runtime::UpgradeCheckSelect,
-	) -> Result<Weight, TryRuntimeError> {
+	pub fn try_runtime_upgrade(checks: UpgradeCheckSelect) -> Result<Weight, TryRuntimeError> {
 		let weight =
 			<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::try_on_runtime_upgrade(
 				checks.pre_and_post(),
 			)?;
+		// Nothing should modify the state after the migrations ran:
+		let _guard = StorageNoopGuard::default();
 
+		// The state must be decodable:
+		if checks.any() {
+			let res = AllPalletsWithSystem::try_decode_entire_state();
+			Self::log_decode_result(res)?;
+		}
+
+		// Check all storage invariants:
 		if checks.try_state() {
-			let _guard = frame_support::StorageNoopGuard::default();
-			<AllPalletsWithSystem as frame_support::traits::TryState<
-				BlockNumberFor<System>,
-			>>::try_state(
+			AllPalletsWithSystem::try_state(
 				frame_system::Pallet::<System>::block_number(),
-				frame_try_runtime::TryStateSelect::All,
+				TryStateSelect::All,
 			)?;
 		}
 
 		Ok(weight)
 	}
+
+	/// Logs the result of trying to decode the entire state.
+	fn log_decode_result(
+		res: Result<usize, Vec<TryDecodeEntireStorageError>>,
+	) -> Result<(), TryRuntimeError> {
+		match res {
+			Ok(bytes) => {
+				log::debug!(
+					target: LOG_TARGET,
+					"decoded the entire state ({bytes} bytes)",
+				);
+
+				Ok(())
+			},
+			Err(errors) => {
+				log::error!(
+					target: LOG_TARGET,
+					"`try_decode_entire_state` failed with {} errors",
+					errors.len(),
+				);
+
+				for (i, err) in errors.iter().enumerate() {
+					// We log the short version to `error` and then the full debug info to `debug`:
+					log::error!(target: LOG_TARGET, "- {i}. error: {err}");
+					log::debug!(target: LOG_TARGET, "- {i}. error: {err:?}");
+				}
+
+				Err("`try_decode_entire_state` failed".into())
+			},
+		}
+	}
 }
 
 impl<

substrate/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)]

substrate/frame/support/Cargo.toml

@@ -13,6 +13,7 @@ readme = "README.md"
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
+array-bytes = { version = "6.1", default-features = false }
 serde = { version = "1.0.188", default-features = false, features = ["alloc", "derive"] }
 codec = { package = "parity-scale-codec", version = "3.6.1", default-features = false, features = ["derive", "max-encoded-len"] }
 scale-info = { version = "2.10.0", default-features = false, features = ["derive"] }
@@ -52,7 +53,6 @@ aquamarine = { version = "0.3.2" }
 assert_matches = "1.3.0"
 pretty_assertions = "1.2.1"
 frame-system = { path = "../system" }
-array-bytes = "6.1"
 
 [features]
 default = [ "std" ]

substrate/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

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

@@ -822,12 +822,69 @@ 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 mut 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<_>>();
+		storage_names.sort_by_cached_key(|ident| ident.to_string());
+
+		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, #frame_support::__private::sp_std::vec::Vec<#frame_support::traits::TryDecodeEntireStorageError>> {
+					let pallet_name = <<T as #frame_system::Config>::PalletInfo	as frame_support::traits::PalletInfo>
+						::name::<#pallet_ident<#type_use_gen>>()
+						.expect("Every active pallet has a name in the runtime; qed");
+
+					#frame_support::__private::log::debug!(target: "runtime::try-decode-state", "trying to decode pallet: {pallet_name}");
+
+					// NOTE: for now, we have to exclude storage items that are feature gated.
+					let mut errors = #frame_support::__private::sp_std::vec::Vec::new();
+					let mut decoded = 0usize;
+
+					#(
+						#frame_support::__private::log::debug!(target: "runtime::try-decode-state", "trying to decode storage: \
+						{pallet_name}::{}", stringify!(#storage_names));
+
+						match <#storage_names as #frame_support::traits::TryDecodeEntireStorage>::try_decode_entire_state() {
+							Ok(count) => {
+								decoded += count;
+							},
+							Err(err) => {
+								errors.extend(err);
+							},
+						}
+					)*
+
+					if errors.is_empty() {
+						Ok(decoded)
+					} else {
+						Err(errors)
+					}
+				}
+			}
+		)
+	};
+
 	quote::quote!(
 		impl<#type_impl_gen> #pallet_ident<#type_use_gen>
 			#completed_where_clause
@@ -853,5 +910,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		#( #getters )*
 		#( #prefix_structs )*
 		#( #on_empty_structs )*
+
+		#try_decode_entire_state
 	)
 }

substrate/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,

substrate/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;

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

@@ -119,7 +119,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`].
@@ -423,14 +427,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()
 	}
 }
@@ -1207,7 +1211,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![]

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

@@ -114,8 +114,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`].
@@ -472,7 +474,7 @@ where
 	}
 
 	/// Return the count.
-	pub fn count() -> u32 {
+	pub fn count() -> Counter {
 		CounterFor::<Prefix>::get()
 	}
 }

substrate/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::{

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

@@ -23,7 +23,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 frame_support::storage::StorageDecodeNonDedupLength;
@@ -72,7 +72,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 pallet_prefix() -> &'static [u8] {

substrate/frame/support/src/traits.rs

@@ -126,4 +126,7 @@ pub use tx_pause::{TransactionPause, TransactionPauseError};
 #[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, TryDecodeEntireStorageError, TryState,
+	UpgradeCheckSelect,
+};

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

@@ -93,9 +93,7 @@ pub trait StorageInstance {
 }
 
 /// Metadata about storage from the runtime.
-#[derive(
-	codec::Encode, codec::Decode, RuntimeDebug, Eq, PartialEq, Clone, scale_info::TypeInfo,
-)]
+#[derive(Debug, codec::Encode, codec::Decode, Eq, PartialEq, Clone, scale_info::TypeInfo)]
 pub struct StorageInfo {
 	/// Encoded string of pallet name.
 	pub pallet_name: Vec<u8>,

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

@@ -0,0 +1,498 @@
+// This file is part of Substrate.
+
+// Copyright (C) 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.
+
+//! Types to check that the entire storage can be decoded.
+
+use super::StorageInstance;
+use crate::{
+	storage::types::{
+		CountedStorageMapInstance, CountedStorageNMapInstance, Counter, KeyGenerator,
+		QueryKindTrait,
+	},
+	traits::{PartialStorageInfoTrait, StorageInfo},
+	StorageHasher,
+};
+use codec::{Decode, DecodeAll, FullCodec};
+use impl_trait_for_tuples::impl_for_tuples;
+use sp_core::Get;
+use sp_std::prelude::*;
+
+/// Decode the entire data under the given storage type.
+///
+/// For values, this is trivial. For all kinds of maps, it should decode all the values associated
+/// with all keys existing in the map.
+///
+/// Tuple implementations are provided and simply decode each type in the tuple, summing up the
+/// decoded bytes if `Ok(_)`.
+pub trait TryDecodeEntireStorage {
+	/// Decode the entire data under the given storage, returning `Ok(bytes_decoded)` if success.
+	fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>>;
+}
+
+#[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(feature = "tuples-128", impl_for_tuples(128))]
+impl TryDecodeEntireStorage for Tuple {
+	fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
+		let mut errors = Vec::new();
+		let mut len = 0usize;
+
+		for_tuples!(#(
+			match Tuple::try_decode_entire_state() {
+				Ok(bytes) => len += bytes,
+				Err(errs) => errors.extend(errs),
+			}
+		)*);
+
+		if errors.is_empty() {
+			Ok(len)
+		} else {
+			Err(errors)
+		}
+	}
+}
+
+/// A value could not be decoded.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct TryDecodeEntireStorageError {
+	/// The key of the undecodable value.
+	pub key: Vec<u8>,
+	/// The raw value.
+	pub raw: Option<Vec<u8>>,
+	/// The storage info of the key.
+	pub info: StorageInfo,
+}
+
+impl core::fmt::Display for TryDecodeEntireStorageError {
+	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+		write!(
+			f,
+			"Failed to decode storage item `{}::{}`",
+			&sp_std::str::from_utf8(&self.info.pallet_name).unwrap_or("<invalid>"),
+			&sp_std::str::from_utf8(&self.info.storage_name).unwrap_or("<invalid>"),
+		)
+	}
+}
+
+/// Decode all the values based on the prefix of `info` to `V`.
+///
+/// Basically, it decodes and sums up all the values who's key start with `info.prefix`. For values,
+/// this would be the value itself. For all sorts of maps, this should be all map items in the
+/// absence of key collision.
+fn decode_storage_info<V: Decode>(
+	info: StorageInfo,
+) -> Result<usize, Vec<TryDecodeEntireStorageError>> {
+	let mut next_key = info.prefix.clone();
+	let mut decoded = 0;
+
+	let decode_key = |key: &[u8]| match sp_io::storage::get(key) {
+		None => Ok(0),
+		Some(bytes) => {
+			let len = bytes.len();
+			let _ = <V as DecodeAll>::decode_all(&mut bytes.as_ref()).map_err(|_| {
+				vec![TryDecodeEntireStorageError {
+					key: key.to_vec(),
+					raw: Some(bytes.to_vec()),
+					info: info.clone(),
+				}]
+			})?;
+
+			Ok::<usize, Vec<_>>(len)
+		},
+	};
+
+	decoded += decode_key(&next_key)?;
+	loop {
+		match sp_io::storage::next_key(&next_key) {
+			Some(key) if key.starts_with(&info.prefix) => {
+				decoded += decode_key(&key)?;
+				next_key = key;
+			},
+			_ => break,
+		}
+	}
+
+	Ok(decoded)
+}
+
+impl<Prefix, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
+	for crate::storage::types::StorageValue<Prefix, Value, QueryKind, OnEmpty>
+where
+	Prefix: StorageInstance,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+{
+	fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
+		let info = Self::partial_storage_info()
+			.first()
+			.cloned()
+			.expect("Value has only one storage info; qed");
+		decode_storage_info::<Value>(info)
+	}
+}
+
+impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
+	for crate::storage::types::StorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues>
+where
+	Prefix: StorageInstance,
+	Hasher: StorageHasher,
+	Key: FullCodec,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
+{
+	fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
+		let info = Self::partial_storage_info()
+			.first()
+			.cloned()
+			.expect("Map has only one storage info; qed");
+		decode_storage_info::<Value>(info)
+	}
+}
+
+impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
+	for crate::storage::types::CountedStorageMap<
+		Prefix,
+		Hasher,
+		Key,
+		Value,
+		QueryKind,
+		OnEmpty,
+		MaxValues,
+	> where
+	Prefix: CountedStorageMapInstance,
+	Hasher: StorageHasher,
+	Key: FullCodec,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
+{
+	fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
+		let (map_info, counter_info) = match &Self::partial_storage_info()[..] {
+			[a, b] => (a.clone(), b.clone()),
+			_ => panic!("Counted map has two storage info items; qed"),
+		};
+		let mut decoded = decode_storage_info::<Counter>(counter_info)?;
+		decoded += decode_storage_info::<Value>(map_info)?;
+		Ok(decoded)
+	}
+}
+
+impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty, MaxValues>
+	TryDecodeEntireStorage
+	for crate::storage::types::StorageDoubleMap<
+		Prefix,
+		Hasher1,
+		Key1,
+		Hasher2,
+		Key2,
+		Value,
+		QueryKind,
+		OnEmpty,
+		MaxValues,
+	> where
+	Prefix: StorageInstance,
+	Hasher1: StorageHasher,
+	Key1: FullCodec,
+	Hasher2: StorageHasher,
+	Key2: FullCodec,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
+{
+	fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
+		let info = Self::partial_storage_info()
+			.first()
+			.cloned()
+			.expect("Double-map has only one storage info; qed");
+		decode_storage_info::<Value>(info)
+	}
+}
+
+impl<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
+	for crate::storage::types::StorageNMap<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues>
+where
+	Prefix: StorageInstance,
+	Key: KeyGenerator,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
+{
+	fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
+		let info = Self::partial_storage_info()
+			.first()
+			.cloned()
+			.expect("N-map has only one storage info; qed");
+		decode_storage_info::<Value>(info)
+	}
+}
+
+impl<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
+	for crate::storage::types::CountedStorageNMap<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues>
+where
+	Prefix: CountedStorageNMapInstance,
+	Key: KeyGenerator,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
+{
+	fn try_decode_entire_state() -> Result<usize, Vec<TryDecodeEntireStorageError>> {
+		let (map_info, counter_info) = match &Self::partial_storage_info()[..] {
+			[a, b] => (a.clone(), b.clone()),
+			_ => panic!("Counted NMap has two storage info items; qed"),
+		};
+
+		let mut decoded = decode_storage_info::<Counter>(counter_info)?;
+		decoded += decode_storage_info::<Value>(map_info)?;
+		Ok(decoded)
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::{
+		storage::types::{self, CountedStorageMapInstance, CountedStorageNMapInstance, Key},
+		Blake2_128Concat,
+	};
+
+	type H = Blake2_128Concat;
+
+	macro_rules! build_prefix {
+		($name:ident) => {
+			struct $name;
+			impl StorageInstance for $name {
+				fn pallet_prefix() -> &'static str {
+					"test_pallet"
+				}
+				const STORAGE_PREFIX: &'static str = stringify!($name);
+			}
+		};
+	}
+
+	build_prefix!(ValuePrefix);
+	type Value = types::StorageValue<ValuePrefix, u32>;
+
+	build_prefix!(MapPrefix);
+	type Map = types::StorageMap<MapPrefix, H, u32, u32>;
+
+	build_prefix!(CMapCounterPrefix);
+	build_prefix!(CMapPrefix);
+	impl CountedStorageMapInstance for CMapPrefix {
+		type CounterPrefix = CMapCounterPrefix;
+	}
+	type CMap = types::CountedStorageMap<CMapPrefix, H, u8, u16>;
+
+	build_prefix!(DMapPrefix);
+	type DMap = types::StorageDoubleMap<DMapPrefix, H, u32, H, u32, u32>;
+
+	build_prefix!(NMapPrefix);
+	type NMap = types::StorageNMap<NMapPrefix, (Key<H, u8>, Key<H, u8>), u128>;
+
+	build_prefix!(CountedNMapCounterPrefix);
+	build_prefix!(CountedNMapPrefix);
+	impl CountedStorageNMapInstance for CountedNMapPrefix {
+		type CounterPrefix = CountedNMapCounterPrefix;
+	}
+	type CNMap = types::CountedStorageNMap<CountedNMapPrefix, (Key<H, u8>, Key<H, u8>), u128>;
+
+	#[test]
+	fn try_decode_entire_state_value_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(Value::try_decode_entire_state(), Ok(0));
+
+			Value::put(42);
+			assert_eq!(Value::try_decode_entire_state(), Ok(4));
+
+			Value::kill();
+			assert_eq!(Value::try_decode_entire_state(), Ok(0));
+
+			// two bytes, cannot be decoded into u32.
+			sp_io::storage::set(&Value::hashed_key(), &[0u8, 1]);
+			assert!(Value::try_decode_entire_state().is_err());
+		})
+	}
+
+	#[test]
+	fn try_decode_entire_state_map_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(Map::try_decode_entire_state(), Ok(0));
+
+			Map::insert(0, 42);
+			assert_eq!(Map::try_decode_entire_state(), Ok(4));
+
+			Map::insert(0, 42);
+			assert_eq!(Map::try_decode_entire_state(), Ok(4));
+
+			Map::insert(1, 42);
+			assert_eq!(Map::try_decode_entire_state(), Ok(8));
+
+			Map::remove(0);
+			assert_eq!(Map::try_decode_entire_state(), Ok(4));
+
+			// two bytes, cannot be decoded into u32.
+			sp_io::storage::set(&Map::hashed_key_for(2), &[0u8, 1]);
+			assert!(Map::try_decode_entire_state().is_err());
+		})
+	}
+
+	#[test]
+	fn try_decode_entire_state_counted_map_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			// counter is not even initialized;
+			assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
+
+			let counter = 4;
+			let value_size = std::mem::size_of::<u16>();
+
+			CMap::insert(0, 42);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
+
+			CMap::insert(0, 42);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
+
+			CMap::insert(1, 42);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
+
+			CMap::remove(0);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
+
+			// counter is cleared again.
+			let _ = CMap::clear(u32::MAX, None);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
+
+			// 1 bytes, cannot be decoded into u16.
+			sp_io::storage::set(&CMap::hashed_key_for(2), &[0u8]);
+			assert!(CMap::try_decode_entire_state().is_err());
+		})
+	}
+
+	#[test]
+	fn try_decode_entire_state_double_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(DMap::try_decode_entire_state(), Ok(0));
+
+			DMap::insert(0, 0, 42);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(4));
+
+			DMap::insert(0, 0, 42);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(4));
+
+			DMap::insert(0, 1, 42);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(8));
+
+			DMap::insert(1, 0, 42);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(12));
+
+			DMap::remove(0, 0);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(8));
+
+			// two bytes, cannot be decoded into u32.
+			sp_io::storage::set(&DMap::hashed_key_for(1, 1), &[0u8, 1]);
+			assert!(DMap::try_decode_entire_state().is_err());
+		})
+	}
+
+	#[test]
+	fn try_decode_entire_state_n_map_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(NMap::try_decode_entire_state(), Ok(0));
+
+			let value_size = std::mem::size_of::<u128>();
+
+			NMap::insert((0u8, 0), 42);
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size));
+
+			NMap::insert((0, 0), 42);
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size));
+
+			NMap::insert((0, 1), 42);
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 2));
+
+			NMap::insert((1, 0), 42);
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 3));
+
+			NMap::remove((0, 0));
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 2));
+
+			// two bytes, cannot be decoded into u128.
+			sp_io::storage::set(&NMap::hashed_key_for((1, 1)), &[0u8, 1]);
+			assert!(NMap::try_decode_entire_state().is_err());
+		})
+	}
+
+	#[test]
+	fn try_decode_entire_state_counted_n_map_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			sp_io::TestExternalities::new_empty().execute_with(|| {
+				assert_eq!(NMap::try_decode_entire_state(), Ok(0));
+
+				let value_size = std::mem::size_of::<u128>();
+				let counter = 4;
+
+				CNMap::insert((0u8, 0), 42);
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size + counter));
+
+				CNMap::insert((0, 0), 42);
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size + counter));
+
+				CNMap::insert((0, 1), 42);
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
+
+				CNMap::insert((1, 0), 42);
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 3 + counter));
+
+				CNMap::remove((0, 0));
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
+
+				// two bytes, cannot be decoded into u128.
+				sp_io::storage::set(&CNMap::hashed_key_for((1, 1)), &[0u8, 1]);
+				assert!(CNMap::try_decode_entire_state().is_err());
+			})
+		})
+	}
+
+	#[test]
+	fn extra_bytes_are_rejected() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(Map::try_decode_entire_state(), Ok(0));
+
+			// 6bytes, too many to fit in u32, should be rejected.
+			sp_io::storage::set(&Map::hashed_key_for(2), &[0u8, 1, 3, 4, 5, 6]);
+			assert!(Map::try_decode_entire_state().is_err());
+		})
+	}
+
+	#[test]
+	fn try_decode_entire_state_tuple_of_storage_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(0));
+
+			Value::put(42);
+			assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(4));
+
+			Map::insert(0, 42);
+			assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(8));
+		});
+	}
+}

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

@@ -17,6 +17,11 @@
 
 //! Try-runtime specific traits and types.
 
+pub mod decode_entire_state;
+pub use decode_entire_state::{TryDecodeEntireStorage, TryDecodeEntireStorageError};
+
+use super::StorageInstance;
+
 use impl_trait_for_tuples::impl_for_tuples;
 use sp_arithmetic::traits::AtLeast32BitUnsigned;
 use sp_runtime::TryRuntimeError;
@@ -37,6 +42,13 @@ pub enum Select {
 	Only(Vec<Vec<u8>>),
 }
 
+impl Select {
+	/// Whether to run any checks at all.
+	pub fn any(&self) -> bool {
+		!matches!(self, Select::None)
+	}
+}
+
 impl Default for Select {
 	fn default() -> Self {
 		Select::None
@@ -105,6 +117,11 @@ impl UpgradeCheckSelect {
 	pub fn try_state(&self) -> bool {
 		matches!(self, Self::All | Self::TryState)
 	}
+
+	/// Whether to run any checks at all.
+	pub fn any(&self) -> bool {
+		!matches!(self, Self::None)
+	}
 }
 
 #[cfg(feature = "std")]

substrate/frame/support/test/tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen.stderr

@@ -128,3 +128,62 @@ error[E0277]: the trait bound `Bar: WrapperTypeEncode` is not satisfied
    = note: required for `Bar` to implement `FullEncode`
    = note: required for `Bar` to implement `FullCodec`
    = note: required for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>` to implement `StorageEntryMetadataBuilder`
+
+error[E0277]: the trait bound `Bar: WrapperTypeDecode` is not satisfied
+  --> tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen.rs:18:1
+   |
+18 | #[frame_support::pallet]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^ the trait `WrapperTypeDecode` is not implemented for `Bar`
+   |
+   = help: the following other types implement trait `WrapperTypeDecode`:
+             Box<T>
+             frame_support::sp_runtime::sp_application_crypto::sp_core::Bytes
+             Rc<T>
+             Arc<T>
+   = note: required for `Bar` to implement `Decode`
+   = note: required for `Bar` to implement `FullCodec`
+   = note: required for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>` to implement `TryDecodeEntireStorage`
+   = note: this error originates in the attribute macro `frame_support::pallet` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error[E0277]: the trait bound `Bar: EncodeLike` is not satisfied
+  --> tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen.rs:18:1
+   |
+18 | #[frame_support::pallet]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^ the trait `EncodeLike` is not implemented for `Bar`
+   |
+   = help: the following other types implement trait `EncodeLike<T>`:
+             <bool as EncodeLike>
+             <i8 as EncodeLike>
+             <i16 as EncodeLike>
+             <i32 as EncodeLike>
+             <i64 as EncodeLike>
+             <i128 as EncodeLike>
+             <u8 as EncodeLike>
+             <u16 as EncodeLike>
+           and $N others
+   = note: required for `Bar` to implement `FullEncode`
+   = note: required for `Bar` to implement `FullCodec`
+   = note: required for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>` to implement `TryDecodeEntireStorage`
+   = note: this error originates in the attribute macro `frame_support::pallet` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error[E0277]: the trait bound `Bar: WrapperTypeEncode` is not satisfied
+  --> tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen.rs:18:1
+   |
+18 | #[frame_support::pallet]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^ the trait `WrapperTypeEncode` is not implemented for `Bar`
+   |
+   = help: the following other types implement trait `WrapperTypeEncode`:
+             Box<T>
+             bytes::bytes::Bytes
+             Cow<'a, T>
+             parity_scale_codec::Ref<'a, T, U>
+             frame_support::sp_runtime::sp_application_crypto::sp_core::Bytes
+             Rc<T>
+             Arc<T>
+             Vec<T>
+           and $N others
+   = note: required for `Bar` to implement `Encode`
+   = note: required for `Bar` to implement `FullEncode`
+   = note: required for `Bar` to implement `FullCodec`
+   = note: required for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>` to implement `TryDecodeEntireStorage`
+   = note: this error originates in the attribute macro `frame_support::pallet` (in Nightly builds, run with -Z macro-backtrace for more info)

substrate/frame/support/test/tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen_unnamed.stderr

@@ -128,3 +128,62 @@ error[E0277]: the trait bound `Bar: WrapperTypeEncode` is not satisfied
    = note: required for `Bar` to implement `FullEncode`
    = note: required for `Bar` to implement `FullCodec`
    = note: required for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>` to implement `StorageEntryMetadataBuilder`
+
+error[E0277]: the trait bound `Bar: WrapperTypeDecode` is not satisfied
+  --> tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen_unnamed.rs:18:1
+   |
+18 | #[frame_support::pallet]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^ the trait `WrapperTypeDecode` is not implemented for `Bar`
+   |
+   = help: the following other types implement trait `WrapperTypeDecode`:
+             Box<T>
+             frame_support::sp_runtime::sp_application_crypto::sp_core::Bytes
+             Rc<T>
+             Arc<T>
+   = note: required for `Bar` to implement `Decode`
+   = note: required for `Bar` to implement `FullCodec`
+   = note: required for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>` to implement `TryDecodeEntireStorage`
+   = note: this error originates in the attribute macro `frame_support::pallet` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error[E0277]: the trait bound `Bar: EncodeLike` is not satisfied
+  --> tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen_unnamed.rs:18:1
+   |
+18 | #[frame_support::pallet]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^ the trait `EncodeLike` is not implemented for `Bar`
+   |
+   = help: the following other types implement trait `EncodeLike<T>`:
+             <bool as EncodeLike>
+             <i8 as EncodeLike>
+             <i16 as EncodeLike>
+             <i32 as EncodeLike>
+             <i64 as EncodeLike>
+             <i128 as EncodeLike>
+             <u8 as EncodeLike>
+             <u16 as EncodeLike>
+           and $N others
+   = note: required for `Bar` to implement `FullEncode`
+   = note: required for `Bar` to implement `FullCodec`
+   = note: required for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>` to implement `TryDecodeEntireStorage`
+   = note: this error originates in the attribute macro `frame_support::pallet` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error[E0277]: the trait bound `Bar: WrapperTypeEncode` is not satisfied
+  --> tests/pallet_ui/storage_ensure_span_are_ok_on_wrong_gen_unnamed.rs:18:1
+   |
+18 | #[frame_support::pallet]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^ the trait `WrapperTypeEncode` is not implemented for `Bar`
+   |
+   = help: the following other types implement trait `WrapperTypeEncode`:
+             Box<T>
+             bytes::bytes::Bytes
+             Cow<'a, T>
+             parity_scale_codec::Ref<'a, T, U>
+             frame_support::sp_runtime::sp_application_crypto::sp_core::Bytes
+             Rc<T>
+             Arc<T>
+             Vec<T>
+           and $N others
+   = note: required for `Bar` to implement `Encode`
+   = note: required for `Bar` to implement `FullEncode`
+   = note: required for `Bar` to implement `FullCodec`
+   = note: required for `frame_support::pallet_prelude::StorageValue<_GeneratedPrefixForStorageFoo<T>, Bar>` to implement `TryDecodeEntireStorage`
+   = note: this error originates in the attribute macro `frame_support::pallet` (in Nightly builds, run with -Z macro-backtrace for more info)