Add special tag to exclude runtime storage items from benchmarking (#12205)

* initial setup

* add WhitelistedStorageKeys trait

* add (A, B) tuple implementation for whitelisted_storage_keys()

* fix formatting

* implement WhitelistedStorageKeys for all tuple combinations

* impl_for_tuples up to 128 for WhitelistedStorageKeys

* refactor to #[benchmarking(cached)]

* tweak error message and mark BlockNumber as cached

* add benchmarking(cached) to the other default types

* add docs for benchmarking(cached)

* properly parse storage type declaration

* make storage_alias structs public so we can use them in this macro

* use BTreeMap since TrackedStorageKey missing Ord outside of std

* make WhitelistedStorageKeys accessible

* basic detection of benchmarking(cached) 💥

* proper parsing of #[benchmarking(cached)] from pallet parse macro

* store presence of #[benchmarking(cached)] macro on StorageDef

* will be used for later expansion

* compiling blank impl for WhitelistedStorageKeys

* move impl to expand_pallet_struct

* use frame_support::sp_std::vec::Vec properly

* successfully compiling with storage info loaded into a variable 💥

* plausible implementation for whitelisted_storage_keys()

* depends on the assumption that storage_info.encode() can be loaded
  into TrackedStorageKey::new(..)

* use Pallet::whitelisted_storage_keys() instead of hard-coded list

* AllPallets::whitelisted_storage_keys() properly working 💥

* collect storage names

* whitelisted_storage_keys() impl working 💥

* clean up

* fix compiler error

* just one compiler error

* fix doc compiler error

* use better import path

* fix comment

* whoops

* whoops again

* fix macro import issue

* cargo fmt

* mark example as ignore

* use keyword tokens instead of string parsing

* fix keyword-based parsing of benchmarking(cached)

* preliminary spec for check_whitelist()

* add additional test for benchmarking whitelist

* add TODO note

* remove irrelevant line from example

* use filter_map instead of filter and map

* simplify syntax

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

* clean up

* fix test

* fix tests

* use keyword parsing instead of string parsing

* use collect() instead of a for loop

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* fix compiler error

* clean up benchmarking(cached) marking code

* use cloned()

* refactor to not use panic! and remove need for pub types in storage_alias

* remove unneeded use

Co-authored-by: Bastian Köcher <info@kchr.de>

* remove unneeded visibility changes

* don't manually hard code hash for treasury account as hex

* proper Ord, PartialOrd, and Hash impls for TrackedStorageKey

* now based just on key, and available in no-std

* use BTreeSet instead of BTreeMap

* fix comments

* cargo fmt

* switch to pallet::whitelist and re-do it basti's way :D

* make PartialOrd for TrackedStorageKey consistent with Ord

* more correct implementation of hash-related traits for TrackedStorageKey

* fix integration test

* update TODO

* remove unused keyword

* remove more unused keywords

* use into_iter()

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

* Update frame/support/procedural/src/pallet/parse/mod.rs

Co-authored-by: Bastian Köcher <info@kchr.de>

* add docs for whitelisted

* fix comment

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>
Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
Co-authored-by: Bastian Köcher <info@kchr.de>

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

@@ -521,18 +521,8 @@ impl_runtime_apis! {
 			impl frame_system_benchmarking::Config for Runtime {}
 			impl baseline::Config for Runtime {}
 
-			let whitelist: Vec<TrackedStorageKey> = vec![
-				// Block Number
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef702a5c1b19ab7a04f536c519aca4983ac").to_vec().into(),
-				// Total Issuance
-				hex_literal::hex!("c2261276cc9d1f8598ea4b6a74b15c2f57c875e4cff74148e4628f264b974c80").to_vec().into(),
-				// Execution Phase
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7ff553b5a9862a516939d82b3d3d8661a").to_vec().into(),
-				// Event Count
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef70a98fdbe9ce6c55837576c60c7af3850").to_vec().into(),
-				// System Events
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef780d41e5e16056765bc8461851072c9d7").to_vec().into(),
-			];
+			use frame_support::traits::WhitelistedStorageKeys;
+			let whitelist: Vec<TrackedStorageKey> = AllPalletsWithSystem::whitelisted_storage_keys();
 
 			let mut batches = Vec::<BenchmarkBatch>::new();
 			let params = (&config, &whitelist);
@@ -563,3 +553,40 @@ impl_runtime_apis! {
 		}
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use frame_support::traits::WhitelistedStorageKeys;
+	use sp_core::hexdisplay::HexDisplay;
+	use std::collections::HashSet;
+
+	#[test]
+	fn check_whitelist() {
+		let whitelist: HashSet<String> = AllPalletsWithSystem::whitelisted_storage_keys()
+			.iter()
+			.map(|e| HexDisplay::from(&e.key).to_string())
+			.collect();
+
+		// Block Number
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef702a5c1b19ab7a04f536c519aca4983ac")
+		);
+		// Total Issuance
+		assert!(
+			whitelist.contains("c2261276cc9d1f8598ea4b6a74b15c2f57c875e4cff74148e4628f264b974c80")
+		);
+		// Execution Phase
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef7ff553b5a9862a516939d82b3d3d8661a")
+		);
+		// Event Count
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef70a98fdbe9ce6c55837576c60c7af3850")
+		);
+		// System Events
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef780d41e5e16056765bc8461851072c9d7")
+		);
+	}
+}

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

@@ -2148,22 +2148,14 @@ impl_runtime_apis! {
 			impl baseline::Config for Runtime {}
 			impl pallet_nomination_pools_benchmarking::Config for Runtime {}
 
-			let whitelist: Vec<TrackedStorageKey> = vec![
-				// Block Number
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef702a5c1b19ab7a04f536c519aca4983ac").to_vec().into(),
-				// Total Issuance
-				hex_literal::hex!("c2261276cc9d1f8598ea4b6a74b15c2f57c875e4cff74148e4628f264b974c80").to_vec().into(),
-				// Execution Phase
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7ff553b5a9862a516939d82b3d3d8661a").to_vec().into(),
-				// Event Count
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef70a98fdbe9ce6c55837576c60c7af3850").to_vec().into(),
-				// System Events
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef780d41e5e16056765bc8461851072c9d7").to_vec().into(),
-				// System BlockWeight
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef734abf5cb34d6244378cddbf18e849d96").to_vec().into(),
-				// Treasury Account
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7b99d880ec681799c0cf30e8886371da95ecffd7b6c0f78751baa9d281e0bfa3a6d6f646c70792f74727372790000000000000000000000000000000000000000").to_vec().into(),
-			];
+			use frame_support::traits::WhitelistedStorageKeys;
+			let mut whitelist: Vec<TrackedStorageKey> = AllPalletsWithSystem::whitelisted_storage_keys();
+
+			// Treasury Account
+			// TODO: this is manual for now, someday we might be able to use a
+			// macro for this particular key
+			let treasury_key = frame_system::Account::<Runtime>::hashed_key_for(Treasury::account_id());
+			whitelist.push(treasury_key.to_vec().into());
 
 			let mut batches = Vec::<BenchmarkBatch>::new();
 			let params = (&config, &whitelist);
@@ -2177,8 +2169,44 @@ impl_runtime_apis! {
 mod tests {
 	use super::*;
 	use frame_election_provider_support::NposSolution;
+	use frame_support::traits::WhitelistedStorageKeys;
 	use frame_system::offchain::CreateSignedTransaction;
+	use sp_core::hexdisplay::HexDisplay;
 	use sp_runtime::UpperOf;
+	use std::collections::HashSet;
+
+	#[test]
+	fn check_whitelist() {
+		let whitelist: HashSet<String> = AllPalletsWithSystem::whitelisted_storage_keys()
+			.iter()
+			.map(|e| HexDisplay::from(&e.key).to_string())
+			.collect();
+
+		// Block Number
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef702a5c1b19ab7a04f536c519aca4983ac")
+		);
+		// Total Issuance
+		assert!(
+			whitelist.contains("c2261276cc9d1f8598ea4b6a74b15c2f57c875e4cff74148e4628f264b974c80")
+		);
+		// Execution Phase
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef7ff553b5a9862a516939d82b3d3d8661a")
+		);
+		// Event Count
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef70a98fdbe9ce6c55837576c60c7af3850")
+		);
+		// System Events
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef780d41e5e16056765bc8461851072c9d7")
+		);
+		// System BlockWeight
+		assert!(
+			whitelist.contains("26aa394eea5630e07c48ae0c9558cef734abf5cb34d6244378cddbf18e849d96")
+		);
+	}
 
 	#[test]
 	fn validate_transaction_submitter_bounds() {

substrate/frame/balances/src/lib.rs

@@ -493,6 +493,7 @@ pub mod pallet {
 	/// The total units issued in the system.
 	#[pallet::storage]
 	#[pallet::getter(fn total_issuance)]
+	#[pallet::whitelist_storage]
 	pub type TotalIssuance<T: Config<I>, I: 'static = ()> = StorageValue<_, T::Balance, ValueQuery>;
 
 	/// The Balances pallet example of storing the balance of an account.

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

@@ -166,6 +166,24 @@ pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 		quote::quote! { #frame_support::traits::StorageVersion::default() }
 	};
 
+	let whitelisted_storage_idents: Vec<syn::Ident> = def
+		.storages
+		.iter()
+		.filter_map(|s| s.whitelisted.then_some(s.ident.clone()))
+		.collect();
+
+	let whitelisted_storage_keys_impl = quote::quote![
+		use #frame_support::traits::{StorageInfoTrait, TrackedStorageKey, WhitelistedStorageKeys};
+		impl<#type_impl_gen> WhitelistedStorageKeys for #pallet_ident<#type_use_gen> #storages_where_clauses {
+			fn whitelisted_storage_keys() -> #frame_support::sp_std::vec::Vec<TrackedStorageKey> {
+				use #frame_support::sp_std::vec;
+				vec![#(
+					TrackedStorageKey::new(#whitelisted_storage_idents::<#type_use_gen>::hashed_key().to_vec())
+				),*]
+			}
+		}
+	];
+
 	quote::quote_spanned!(def.pallet_struct.attr_span =>
 		#pallet_error_metadata
 
@@ -253,5 +271,6 @@ pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 		}
 
 		#storage_info
+		#whitelisted_storage_keys_impl
 	)
 }

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

@@ -28,6 +28,7 @@ mod keyword {
 	syn::custom_keyword!(getter);
 	syn::custom_keyword!(storage_prefix);
 	syn::custom_keyword!(unbounded);
+	syn::custom_keyword!(whitelist_storage);
 	syn::custom_keyword!(OptionQuery);
 	syn::custom_keyword!(ResultQuery);
 	syn::custom_keyword!(ValueQuery);
@@ -37,16 +38,21 @@ mod keyword {
 /// * `#[pallet::getter(fn dummy)]`
 /// * `#[pallet::storage_prefix = "CustomName"]`
 /// * `#[pallet::unbounded]`
+/// * `#[pallet::whitelist_storage]
 pub enum PalletStorageAttr {
 	Getter(syn::Ident, proc_macro2::Span),
 	StorageName(syn::LitStr, proc_macro2::Span),
 	Unbounded(proc_macro2::Span),
+	WhitelistStorage(proc_macro2::Span),
 }
 
 impl PalletStorageAttr {
 	fn attr_span(&self) -> proc_macro2::Span {
 		match self {
-			Self::Getter(_, span) | Self::StorageName(_, span) | Self::Unbounded(span) => *span,
+			Self::Getter(_, span) |
+			Self::StorageName(_, span) |
+			Self::Unbounded(span) |
+			Self::WhitelistStorage(span) => *span,
 		}
 	}
 }
@@ -84,6 +90,9 @@ impl syn::parse::Parse for PalletStorageAttr {
 			content.parse::<keyword::unbounded>()?;
 
 			Ok(Self::Unbounded(attr_span))
+		} else if lookahead.peek(keyword::whitelist_storage) {
+			content.parse::<keyword::whitelist_storage>()?;
+			Ok(Self::WhitelistStorage(attr_span))
 		} else {
 			Err(lookahead.error())
 		}
@@ -94,6 +103,7 @@ struct PalletStorageAttrInfo {
 	getter: Option<syn::Ident>,
 	rename_as: Option<syn::LitStr>,
 	unbounded: bool,
+	whitelisted: bool,
 }
 
 impl PalletStorageAttrInfo {
@@ -101,12 +111,14 @@ impl PalletStorageAttrInfo {
 		let mut getter = None;
 		let mut rename_as = None;
 		let mut unbounded = false;
+		let mut whitelisted = false;
 		for attr in attrs {
 			match attr {
 				PalletStorageAttr::Getter(ident, ..) if getter.is_none() => getter = Some(ident),
 				PalletStorageAttr::StorageName(name, ..) if rename_as.is_none() =>
 					rename_as = Some(name),
 				PalletStorageAttr::Unbounded(..) if !unbounded => unbounded = true,
+				PalletStorageAttr::WhitelistStorage(..) if !whitelisted => whitelisted = true,
 				attr =>
 					return Err(syn::Error::new(
 						attr.attr_span(),
@@ -115,7 +127,7 @@ impl PalletStorageAttrInfo {
 			}
 		}
 
-		Ok(PalletStorageAttrInfo { getter, rename_as, unbounded })
+		Ok(PalletStorageAttrInfo { getter, rename_as, unbounded, whitelisted })
 	}
 }
 
@@ -171,6 +183,8 @@ pub struct StorageDef {
 	pub named_generics: Option<StorageGenerics>,
 	/// If the value stored in this storage is unbounded.
 	pub unbounded: bool,
+	/// Whether or not reads to this storage key will be ignored by benchmarking
+	pub whitelisted: bool,
 }
 
 /// The parsed generic from the
@@ -672,7 +686,7 @@ impl StorageDef {
 		};
 
 		let attrs: Vec<PalletStorageAttr> = helper::take_item_pallet_attrs(&mut item.attrs)?;
-		let PalletStorageAttrInfo { getter, rename_as, unbounded } =
+		let PalletStorageAttrInfo { getter, rename_as, unbounded, whitelisted } =
 			PalletStorageAttrInfo::from_attrs(attrs)?;
 
 		let cfg_attrs = helper::get_item_cfg_attrs(&item.attrs);
@@ -814,6 +828,7 @@ impl StorageDef {
 			cfg_attrs,
 			named_generics,
 			unbounded,
+			whitelisted,
 		})
 	}
 }

substrate/frame/support/procedural/src/storage/mod.rs

@@ -32,6 +32,7 @@ pub(crate) use instance_trait::INHERENT_INSTANCE_NAME;
 use frame_support_procedural_tools::{
 	generate_crate_access, generate_hidden_includes, syn_ext as ext,
 };
+
 use quote::quote;
 
 /// All information contained in input of decl_storage

substrate/frame/support/src/lib.rs

@@ -1858,6 +1858,21 @@ pub mod pallet_prelude {
 /// pub(super) type MyStorage<T> = StorageValue<Value = u32>;
 /// ```
 ///
+/// The optional attribute `#[pallet::whitelist_storage]` will declare the
+/// storage as whitelisted from benchmarking. Doing so will exclude reads of
+/// that value's storage key from counting towards weight calculations during
+/// benchmarking.
+///
+/// This attribute should only be attached to storages that are known to be
+/// read/used in every block. This will result in a more accurate benchmarking weight.
+///
+/// ### Example
+/// ```ignore
+/// #[pallet::storage]
+/// #[pallet::whitelist_storage]
+/// pub(super) type Number<T: Config> = StorageValue<_, T::BlockNumber, ValueQuery>;
+/// ```
+///
 /// All the `cfg` attributes are automatically copied to the items generated for the storage,
 /// i.e. the getter, storage prefix, and the metadata element etc.
 ///

substrate/frame/support/src/traits.rs

@@ -89,6 +89,7 @@ pub mod schedule;
 mod storage;
 pub use storage::{
 	Instance, PartialStorageInfoTrait, StorageInfo, StorageInfoTrait, StorageInstance,
+	TrackedStorageKey, WhitelistedStorageKeys,
 };
 
 mod dispatch;

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

@@ -17,7 +17,9 @@
 
 //! Traits for encoding data related to pallet's storage items.
 
+use crate::sp_std::collections::btree_set::BTreeSet;
 use impl_trait_for_tuples::impl_for_tuples;
+pub use sp_core::storage::TrackedStorageKey;
 use sp_std::prelude::*;
 
 /// An instance of a pallet in the storage.
@@ -90,3 +92,29 @@ impl StorageInfoTrait for Tuple {
 pub trait PartialStorageInfoTrait {
 	fn partial_storage_info() -> Vec<StorageInfo>;
 }
+
+/// Allows a pallet to specify storage keys to whitelist during benchmarking.
+/// This means those keys will be excluded from the benchmarking performance
+/// calculation.
+pub trait WhitelistedStorageKeys {
+	/// Returns a [`Vec<TrackedStorageKey>`] indicating the storage keys that
+	/// should be whitelisted during benchmarking. This means that those keys
+	/// will be excluded from the benchmarking performance calculation.
+	fn whitelisted_storage_keys() -> Vec<TrackedStorageKey>;
+}
+
+#[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 WhitelistedStorageKeys for Tuple {
+	fn whitelisted_storage_keys() -> Vec<TrackedStorageKey> {
+		// de-duplicate the storage keys
+		let mut combined_keys: BTreeSet<TrackedStorageKey> = BTreeSet::new();
+		for_tuples!( #(
+			for storage_key in Tuple::whitelisted_storage_keys() {
+				combined_keys.insert(storage_key);
+			}
+		 )* );
+		combined_keys.into_iter().collect::<Vec<_>>()
+	}
+}

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

@@ -1,4 +1,4 @@
-error: expected one of: `getter`, `storage_prefix`, `unbounded`
+error: expected one of: `getter`, `storage_prefix`, `unbounded`, `whitelist_storage`
   --> $DIR/storage_invalid_attribute.rs:16:12
    |
 16 |     #[pallet::generate_store(pub trait Store)]

substrate/frame/system/src/lib.rs

@@ -562,6 +562,7 @@ pub mod pallet {
 
 	/// The current weight for the block.
 	#[pallet::storage]
+	#[pallet::whitelist_storage]
 	#[pallet::getter(fn block_weight)]
 	pub(super) type BlockWeight<T: Config> = StorageValue<_, ConsumedWeight, ValueQuery>;
 
@@ -584,6 +585,7 @@ pub mod pallet {
 
 	/// The current block number being processed. Set by `execute_block`.
 	#[pallet::storage]
+	#[pallet::whitelist_storage]
 	#[pallet::getter(fn block_number)]
 	pub(super) type Number<T: Config> = StorageValue<_, T::BlockNumber, ValueQuery>;
 
@@ -606,12 +608,14 @@ pub mod pallet {
 	/// Events have a large in-memory size. Box the events to not go out-of-memory
 	/// just in case someone still reads them from within the runtime.
 	#[pallet::storage]
+	#[pallet::whitelist_storage]
 	#[pallet::unbounded]
 	pub(super) type Events<T: Config> =
 		StorageValue<_, Vec<Box<EventRecord<T::RuntimeEvent, T::Hash>>>, ValueQuery>;
 
 	/// The number of events in the `Events<T>` list.
 	#[pallet::storage]
+	#[pallet::whitelist_storage]
 	#[pallet::getter(fn event_count)]
 	pub(super) type EventCount<T: Config> = StorageValue<_, EventIndex, ValueQuery>;
 
@@ -647,6 +651,7 @@ pub mod pallet {
 
 	/// The execution phase of the block.
 	#[pallet::storage]
+	#[pallet::whitelist_storage]
 	pub(super) type ExecutionPhase<T: Config> = StorageValue<_, Phase>;
 
 	#[cfg_attr(feature = "std", derive(Default))]

substrate/primitives/storage/src/lib.rs

@@ -47,8 +47,9 @@ impl AsRef<[u8]> for StorageKey {
 }
 
 /// Storage key with read/write tracking information.
-#[derive(PartialEq, Eq, RuntimeDebug, Clone, Encode, Decode)]
-#[cfg_attr(feature = "std", derive(Hash, PartialOrd, Ord))]
+#[derive(
+	PartialEq, Eq, Ord, PartialOrd, sp_std::hash::Hash, RuntimeDebug, Clone, Encode, Decode,
+)]
 pub struct TrackedStorageKey {
 	pub key: Vec<u8>,
 	pub reads: u32,