Dynamic Benchmarking DB Whitelist (#6815)

* Add `get_whitelist` api

* add whitelisted caller

* Whitelist caller

* remove caller 0

* initial piping of origin (not actual value yet)

* remove attempt to pass origin around

* Add whitelist for `DidUpdate` storage on `pallet_timestamp`

* fix traits

* only add to whitelist if !contains

* PassBy not implemented error

* Whitelist read/writes explicitly per key

* update docs

* reduce trait constraint

* copy pasta

* Apply suggestions from code review

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>
Co-authored-by: Alexander Popiak <alexander.popiak@parity.io>

* rename functions @apopiak

* missed some renaming

* enable doc tests

* Update docs

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>
Co-authored-by: Alexander Popiak <alexander.popiak@parity.io>

bin/node/runtime/src/lib.rs

@@ -1130,7 +1130,7 @@ impl_runtime_apis! {
 			repeat: u32,
 			extra: bool,
 		) -> Result<Vec<frame_benchmarking::BenchmarkBatch>, sp_runtime::RuntimeString> {
-			use frame_benchmarking::{Benchmarking, BenchmarkBatch, add_benchmark};
+			use frame_benchmarking::{Benchmarking, BenchmarkBatch, add_benchmark, TrackedStorageKey};
 			// Trying to add benchmarks directly to the Session Pallet caused cyclic dependency issues.
 			// To get around that, we separated the Session benchmarks into its own crate, which is why
 			// we need these two lines below.
@@ -1142,21 +1142,19 @@ impl_runtime_apis! {
 			impl pallet_offences_benchmarking::Trait for Runtime {}
 			impl frame_system_benchmarking::Trait for Runtime {}
 
-			let whitelist: Vec<Vec<u8>> = vec![
+			let whitelist: Vec<TrackedStorageKey> = vec![
 				// Block Number
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef702a5c1b19ab7a04f536c519aca4983ac").to_vec(),
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef702a5c1b19ab7a04f536c519aca4983ac").to_vec().into(),
 				// Total Issuance
-				hex_literal::hex!("c2261276cc9d1f8598ea4b6a74b15c2f57c875e4cff74148e4628f264b974c80").to_vec(),
+				hex_literal::hex!("c2261276cc9d1f8598ea4b6a74b15c2f57c875e4cff74148e4628f264b974c80").to_vec().into(),
 				// Execution Phase
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7ff553b5a9862a516939d82b3d3d8661a").to_vec(),
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7ff553b5a9862a516939d82b3d3d8661a").to_vec().into(),
 				// Event Count
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef70a98fdbe9ce6c55837576c60c7af3850").to_vec(),
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef70a98fdbe9ce6c55837576c60c7af3850").to_vec().into(),
 				// System Events
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef780d41e5e16056765bc8461851072c9d7").to_vec(),
-				// Caller 0 Account
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7b99d880ec681799c0cf30e8886371da946c154ffd9992e395af90b5b13cc6f295c77033fce8a9045824a6690bbf99c6db269502f0a8d1d2a008542d5690a0749").to_vec(),
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef780d41e5e16056765bc8461851072c9d7").to_vec().into(),
 				// Treasury Account
-				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7b99d880ec681799c0cf30e8886371da95ecffd7b6c0f78751baa9d281e0bfa3a6d6f646c70792f74727372790000000000000000000000000000000000000000").to_vec(),
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7b99d880ec681799c0cf30e8886371da95ecffd7b6c0f78751baa9d281e0bfa3a6d6f646c70792f74727372790000000000000000000000000000000000000000").to_vec().into(),
 			];
 
 			let mut batches = Vec::<BenchmarkBatch>::new();

client/db/src/bench.rs

@@ -24,7 +24,10 @@ use std::collections::HashMap;
 
 use hash_db::{Prefix, Hasher};
 use sp_trie::{MemoryDB, prefixed_key};
-use sp_core::{storage::ChildInfo, hexdisplay::HexDisplay};
+use sp_core::{
+	storage::{ChildInfo, TrackedStorageKey},
+	hexdisplay::HexDisplay
+};
 use sp_runtime::traits::{Block as BlockT, HashFor};
 use sp_runtime::Storage;
 use sp_state_machine::{DBValue, backend::Backend as StateBackend, StorageCollection};
@@ -95,7 +98,7 @@ pub struct BenchmarkingState<B: BlockT> {
 	shared_cache: SharedCache<B>, // shared cache is always empty
 	key_tracker: RefCell<HashMap<Vec<u8>, KeyTracker>>,
 	read_write_tracker: RefCell<ReadWriteTracker>,
-	whitelist: RefCell<Vec<Vec<u8>>>,
+	whitelist: RefCell<Vec<TrackedStorageKey>>,
 }
 
 impl<B: BlockT> BenchmarkingState<B> {
@@ -155,15 +158,14 @@ impl<B: BlockT> BenchmarkingState<B> {
 	fn add_whitelist_to_tracker(&self) {
 		let mut key_tracker = self.key_tracker.borrow_mut();
 
-		let whitelisted = KeyTracker {
-			has_been_read: true,
-			has_been_written: true,
-		};
-
 		let whitelist = self.whitelist.borrow();
 
 		whitelist.iter().for_each(|key| {
-			key_tracker.insert(key.to_vec(), whitelisted);
+			let whitelisted = KeyTracker {
+				has_been_read: key.has_been_read,
+				has_been_written: key.has_been_written,
+			};
+			key_tracker.insert(key.key.clone(), whitelisted);
 		});
 	}
 
@@ -181,18 +183,21 @@ impl<B: BlockT> BenchmarkingState<B> {
 
 		let maybe_tracker = key_tracker.get(key);
 
-		let has_been_read = KeyTracker {
-			has_been_read: true,
-			has_been_written: false,
-		};
-
 		match maybe_tracker {
 			None => {
+				let has_been_read = KeyTracker {
+					has_been_read: true,
+					has_been_written: false,
+				};
 				key_tracker.insert(key.to_vec(), has_been_read);
 				read_write_tracker.add_read();
 			},
 			Some(tracker) => {
 				if !tracker.has_been_read {
+					let has_been_read = KeyTracker {
+						has_been_read: true,
+						has_been_written: tracker.has_been_written,
+					};
 					key_tracker.insert(key.to_vec(), has_been_read);
 					read_write_tracker.add_read();
 				} else {
@@ -426,7 +431,11 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 		self.wipe_tracker()
 	}
 
-	fn set_whitelist(&self, new: Vec<Vec<u8>>) {
+	fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
+		self.whitelist.borrow().to_vec()
+	}
+
+	fn set_whitelist(&self, new: Vec<TrackedStorageKey>) {
 		*self.whitelist.borrow_mut() = new;
 	}
 

frame/balances/src/benchmarking.rs

@@ -22,7 +22,7 @@
 use super::*;
 
 use frame_system::RawOrigin;
-use frame_benchmarking::{benchmarks, account};
+use frame_benchmarking::{benchmarks, account, whitelisted_caller};
 use sp_runtime::traits::Bounded;
 
 use crate::Module as Balances;
@@ -40,7 +40,7 @@ benchmarks! {
 	// * Transfer will create the recipient account.
 	transfer {
 		let existential_deposit = T::ExistentialDeposit::get();
-		let caller = account("caller", 0, SEED);
+		let caller = whitelisted_caller();
 
 		// Give some multiple of the existential deposit + creation fee + transfer fee
 		let balance = existential_deposit.saturating_mul(ED_MULTIPLIER.into());
@@ -60,7 +60,7 @@ benchmarks! {
 	// * Both accounts exist and will continue to exist.
 	#[extra]
 	transfer_best_case {
-		let caller = account("caller", 0, SEED);
+		let caller = whitelisted_caller();
 		let recipient: T::AccountId = account("recipient", 0, SEED);
 		let recipient_lookup: <T::Lookup as StaticLookup>::Source = T::Lookup::unlookup(recipient.clone());
 
@@ -80,7 +80,7 @@ benchmarks! {
 	// Benchmark `transfer_keep_alive` with the worst possible condition:
 	// * The recipient account is created.
 	transfer_keep_alive {
-		let caller = account("caller", 0, SEED);
+		let caller = whitelisted_caller();
 		let recipient: T::AccountId = account("recipient", 0, SEED);
 		let recipient_lookup: <T::Lookup as StaticLookup>::Source = T::Lookup::unlookup(recipient.clone());
 

frame/benchmarking/Cargo.toml

@@ -20,9 +20,13 @@ sp-runtime-interface = { version = "2.0.0-rc5", path = "../../primitives/runtime
 sp-runtime = { version = "2.0.0-rc5", path = "../../primitives/runtime", default-features = false }
 sp-std = { version = "2.0.0-rc5", path = "../../primitives/std", default-features = false }
 sp-io = { version = "2.0.0-rc5", path = "../../primitives/io", default-features = false }
+sp-storage = { version = "2.0.0-rc5", path = "../../primitives/storage", default-features = false }
 frame-support = { version = "2.0.0-rc5", default-features = false, path = "../support" }
 frame-system = { version = "2.0.0-rc5", default-features = false, path = "../system" }
 
+[dev-dependencies]
+hex-literal = "0.2.1"
+
 [features]
 default = [ "std" ]
 std = [