Add special tag to exclude runtime storage items from benchmarking

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

@@ -514,18 +514,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);
@@ -556,3 +546,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")
+		);
+	}
+}

bin/node/runtime/src/lib.rs

@@ -2142,22 +2142,15 @@ 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 (manually whitelist this for now since no
+			// accessible storage definition we can attach the
+			// benchmarking(cached) attr macro to)
+			// TODO: figure out a workaround for this that doesn't involve
+			// manually hard-coding
+			whitelist.push(hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7b99d880ec681799c0cf30e8886371da95ecffd7b6c0f78751baa9d281e0bfa3a6d6f646c70792f74727372790000000000000000000000000000000000000000").to_vec().into());
 
 			let mut batches = Vec::<BenchmarkBatch>::new();
 			let params = (&config, &whitelist);
@@ -2171,8 +2164,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() {

frame/balances/src/lib.rs

@@ -167,7 +167,7 @@ use codec::{Codec, Decode, Encode, MaxEncodedLen};
 #[cfg(feature = "std")]
 use frame_support::traits::GenesisBuild;
 use frame_support::{
-	ensure,
+	benchmarking, ensure,
 	pallet_prelude::DispatchResult,
 	traits::{
 		tokens::{fungible, BalanceStatus as Status, DepositConsequence, WithdrawConsequence},
@@ -492,6 +492,7 @@ pub mod pallet {
 	/// The total units issued in the system.
 	#[pallet::storage]
 	#[pallet::getter(fn total_issuance)]
+	#[benchmarking(cached)]
 	pub type TotalIssuance<T: Config<I>, I: 'static = ()> = StorageValue<_, T::Balance, ValueQuery>;
 
 	/// The Balances pallet example of storing the balance of an account.

frame/support/procedural/src/lib.rs

@@ -38,6 +38,7 @@ mod tt_macro;
 use proc_macro::TokenStream;
 use std::{cell::RefCell, str::FromStr};
 pub(crate) use storage::INHERENT_INSTANCE_NAME;
+use syn::{parse_macro_input, Ident};
 
 thread_local! {
 	/// A global counter, can be used to generate a relatively unique identifier.
@@ -583,3 +584,19 @@ pub fn storage_alias(_: TokenStream, input: TokenStream) -> TokenStream {
 		.unwrap_or_else(|r| r.into_compile_error())
 		.into()
 }
+
+#[proc_macro_attribute]
+pub fn benchmarking(attr: TokenStream, item: TokenStream) -> TokenStream {
+	match parse_macro_input!(attr as Ident).to_string().as_str() {
+		"cached" => benchmarking_cached(item),
+		_ => panic!("only benchmarking(cached) is supported at this time"),
+	}
+}
+
+fn benchmarking_cached(item: TokenStream) -> TokenStream {
+	// re-use storage_alias's Input parser since it is accessible
+	// and valid for all storage item declarations
+	syn::parse2::<crate::storage_alias::Input>(item.clone().into())
+		.expect("benchmarking(cached) can only be attached to a valid storage type declaration");
+	item
+}

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.benchmarking_cached.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::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
 	)
 }

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

@@ -36,7 +36,9 @@ pub mod type_value;
 pub mod validate_unsigned;
 
 use frame_support_procedural_tools::generate_crate_access_2018;
-use syn::spanned::Spanned;
+use proc_macro2::Group;
+use quote::ToTokens;
+use syn::{spanned::Spanned, AttrStyle};
 
 /// Parsed definition of a pallet.
 pub struct Def {
@@ -123,8 +125,31 @@ impl Def {
 					origin = Some(origin::OriginDef::try_from(index, item)?),
 				Some(PalletAttr::Inherent(_)) if inherent.is_none() =>
 					inherent = Some(inherent::InherentDef::try_from(index, item)?),
-				Some(PalletAttr::Storage(span)) =>
-					storages.push(storage::StorageDef::try_from(span, index, item)?),
+				Some(PalletAttr::Storage(span)) => {
+					let mut storage_def = storage::StorageDef::try_from(span, index, item)?;
+					// check for #[benchmarking(cached)] calls
+					if let syn::Item::Type(typ) = item {
+						for attr in typ.attrs.as_slice() {
+							if attr.style == AttrStyle::Outer {
+								if let Some(seg) = attr.path.segments.last() {
+									if let Ok(_) =
+										syn::parse2::<keyword::benchmarking>(seg.to_token_stream())
+									{
+										if let Ok(group) = syn::parse2::<Group>(attr.tokens.clone())
+										{
+											if let Ok(_) =
+												syn::parse2::<keyword::cached>(group.stream())
+											{
+												storage_def.benchmarking_cached = true;
+											}
+										}
+									}
+								}
+							}
+						}
+					}
+					storages.push(storage_def);
+				},
 				Some(PalletAttr::ValidateUnsigned(_)) if validate_unsigned.is_none() => {
 					let v = validate_unsigned::ValidateUnsignedDef::try_from(index, item)?;
 					validate_unsigned = Some(v);
@@ -383,6 +408,8 @@ mod keyword {
 	syn::custom_keyword!(generate_store);
 	syn::custom_keyword!(Store);
 	syn::custom_keyword!(extra_constants);
+	syn::custom_keyword!(cached);
+	syn::custom_keyword!(benchmarking);
 }
 
 /// Parse attributes for item in pallet module

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

@@ -171,6 +171,7 @@ pub struct StorageDef {
 	pub named_generics: Option<StorageGenerics>,
 	/// If the value stored in this storage is unbounded.
 	pub unbounded: bool,
+	pub benchmarking_cached: bool,
 }
 
 /// The parsed generic from the
@@ -814,6 +815,7 @@ impl StorageDef {
 			cfg_attrs,
 			named_generics,
 			unbounded,
+			benchmarking_cached: false,
 		})
 	}
 }

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

frame/support/procedural/src/storage_alias.rs

@@ -29,7 +29,7 @@ use syn::{
 };
 
 /// Represents a path that only consists of [`Ident`] separated by `::`.
-struct SimplePath {
+pub struct SimplePath {
 	leading_colon: Option<Token![::]>,
 	segments: Punctuated<Ident, Token![::]>,
 }
@@ -65,7 +65,7 @@ impl ToTokens for SimplePath {
 
 /// Represents generics which only support [`Ident`] separated by commas as you would pass it to a
 /// type.
-struct TypeGenerics {
+pub struct TypeGenerics {
 	lt_token: Token![<],
 	params: Punctuated<Ident, token::Comma>,
 	gt_token: Token![>],
@@ -97,9 +97,9 @@ impl ToTokens for TypeGenerics {
 }
 
 /// Represents generics which only support [`TypeParam`] separated by commas.
-struct SimpleGenerics {
+pub struct SimpleGenerics {
 	lt_token: Token![<],
-	params: Punctuated<TypeParam, token::Comma>,
+	pub params: Punctuated<TypeParam, token::Comma>,
 	gt_token: Token![>],
 }
 
@@ -141,7 +141,7 @@ mod storage_types {
 }
 
 /// The supported storage types
-enum StorageType {
+pub enum StorageType {
 	Value {
 		_kw: storage_types::StorageValue,
 		_lt_token: Token![<],
@@ -406,15 +406,15 @@ impl Parse for StorageType {
 }
 
 /// The input expected by this macro.
-struct Input {
-	attributes: Vec<Attribute>,
-	visibility: Visibility,
+pub struct Input {
+	pub attributes: Vec<Attribute>,
+	pub visibility: Visibility,
 	_type: Token![type],
-	storage_name: Ident,
-	storage_generics: Option<SimpleGenerics>,
-	where_clause: Option<WhereClause>,
+	pub storage_name: Ident,
+	pub storage_generics: Option<SimpleGenerics>,
+	pub where_clause: Option<WhereClause>,
 	_equal: Token![=],
-	storage_type: StorageType,
+	pub storage_type: StorageType,
 	_semicolon: Token![;],
 }
 

frame/support/src/lib.rs

@@ -210,6 +210,27 @@ impl TypeId for PalletId {
 /// ```
 pub use frame_support_procedural::storage_alias;
 
+/// A series of attribute macros related to benchmarking. Right now just
+/// [`benchmarking(cached)`](frame_support_procedural::benchmarking(cached)) is provided.
+///
+/// # benchmarking(cached)
+///
+/// This is an attribute macro that can be attached to a [`StorageValue`]. Doing
+/// so will exclude reads of that value's storage key from counting towards
+/// weight calculations during benchmarking.
+///
+/// This attribute should be attached to [`StorageValue`]s that are known to be
+/// read/used in every block. This will result in a more efficient read pattern
+/// for the value and a more accurate benchmarking weight
+///
+/// ## Example
+/// ```ignore
+/// #[pallet::storage]
+/// #[benchmarking(cached)]
+/// pub(super) type Number<T: Config> = StorageValue<_, T::BlockNumber, ValueQuery>;
+/// ```
+pub use frame_support_procedural::benchmarking;
+
 /// Create new implementations of the [`Get`](crate::traits::Get) trait.
 ///
 /// The so-called parameter type can be created in four different ways:

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;