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>
2026-04-28 07:27:55 +00:00 · 2022-09-16 17:45:44 -04:00
parent 86b122d7a7
commit c0e007b50f
12 changed files with 175 additions and 34 deletions
@@ -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
 	)
 }
@@ -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,
 		})
 	}
 }
@@ -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