Add Whitelist to DB Tracking in Benchmarks Pipeline

bin/node/runtime/Cargo.toml

@@ -18,6 +18,7 @@ codec = { package = "parity-scale-codec", version = "1.3.0", default-features =
 integer-sqrt = { version = "0.1.2" }
 serde = { version = "1.0.102", optional = true }
 static_assertions = "1.1.0"
+hex-literal = "0.2.1"
 
 # primitives
 sp-authority-discovery = { version = "2.0.0-rc3", default-features = false, path = "../../../primitives/authority-discovery" }

bin/node/runtime/src/lib.rs

@@ -1074,8 +1074,26 @@ impl_runtime_apis! {
 			impl pallet_offences_benchmarking::Trait for Runtime {}
 			impl frame_system_benchmarking::Trait for Runtime {}
 
+			let whitelist: Vec<Vec<u8>> = vec![
+				// Block Number
+				// frame_system::Number::<Runtime>::hashed_key().to_vec(),
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef702a5c1b19ab7a04f536c519aca4983ac").to_vec(),
+				// Total Issuance
+				hex_literal::hex!("c2261276cc9d1f8598ea4b6a74b15c2f57c875e4cff74148e4628f264b974c80").to_vec(),
+				// Execution Phase
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7ff553b5a9862a516939d82b3d3d8661a").to_vec(),
+				// Event Count
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef70a98fdbe9ce6c55837576c60c7af3850").to_vec(),
+				// System Events
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef780d41e5e16056765bc8461851072c9d7").to_vec(),
+				// Caller 0 Account
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7b99d880ec681799c0cf30e8886371da946c154ffd9992e395af90b5b13cc6f295c77033fce8a9045824a6690bbf99c6db269502f0a8d1d2a008542d5690a0749").to_vec(),
+				// Treasury Account
+				hex_literal::hex!("26aa394eea5630e07c48ae0c9558cef7b99d880ec681799c0cf30e8886371da95ecffd7b6c0f78751baa9d281e0bfa3a6d6f646c70792f74727372790000000000000000000000000000000000000000").to_vec(),
+			];
+
 			let mut batches = Vec::<BenchmarkBatch>::new();
-			let params = (&pallet, &benchmark, &lowest_range_values, &highest_range_values, &steps, repeat);
+			let params = (&pallet, &benchmark, &lowest_range_values, &highest_range_values, &steps, repeat, &whitelist);
 
 			add_benchmark!(params, batches, b"balances", Balances);
 			add_benchmark!(params, batches, b"collective", Council);

client/db/src/bench.rs

@@ -95,6 +95,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>>>,
 }
 
 impl<B: BlockT> BenchmarkingState<B> {
@@ -114,8 +115,11 @@ impl<B: BlockT> BenchmarkingState<B> {
 			shared_cache: new_shared_cache(0, (1, 10)),
 			key_tracker: Default::default(),
 			read_write_tracker: Default::default(),
+			whitelist: Default::default(),
 		};
 
+		state.add_whitelist_to_tracker();
+
 		state.reopen()?;
 		let child_delta = genesis.children_default.iter().map(|(_storage_key, child_content)| (
 			&child_content.child_info,
@@ -148,8 +152,24 @@ impl<B: BlockT> BenchmarkingState<B> {
 		Ok(())
 	}
 
+	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_mut();
+
+		whitelist.iter().for_each(|key| {
+			key_tracker.insert(key.to_vec(), whitelisted);
+		});
+	}
+
 	fn wipe_tracker(&self) {
 		*self.key_tracker.borrow_mut() = HashMap::new();
+		self.add_whitelist_to_tracker();
 		*self.read_write_tracker.borrow_mut() = Default::default();
 	}
 
@@ -406,6 +426,10 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 		self.wipe_tracker()
 	}
 
+	fn set_whitelist(&self, new: Vec<Vec<u8>>) {
+		*self.whitelist.borrow_mut() = new;
+	}
+
 	fn register_overlay_stats(&mut self, stats: &sp_state_machine::StateMachineStats) {
 		self.state.borrow_mut().as_mut().map(|s| s.register_overlay_stats(stats));
 	}

frame/benchmarking/src/lib.rs

@@ -708,6 +708,7 @@ macro_rules! impl_benchmark {
 				highest_range_values: &[u32],
 				steps: &[u32],
 				repeat: u32,
+				whitelist: &[Vec<u8>]
 			) -> Result<Vec<$crate::BenchmarkResults>, &'static str> {
 				// Map the input to the selected benchmark.
 				let extrinsic = sp_std::str::from_utf8(extrinsic)
@@ -717,6 +718,9 @@ macro_rules! impl_benchmark {
 					_ => return Err("Could not find extrinsic."),
 				};
 
+				// Add whitelist to DB
+				$crate::benchmarking::set_whitelist(whitelist.to_vec());
+
 				// Warm up the DB
 				$crate::benchmarking::commit_db();
 				$crate::benchmarking::wipe_db();
@@ -831,6 +835,7 @@ macro_rules! impl_benchmark {
 				highest_range_values: &[u32],
 				steps: &[u32],
 				repeat: u32,
+				whitelist: &[Vec<u8>]
 			) -> Result<Vec<$crate::BenchmarkResults>, &'static str> {
 				// Map the input to the selected benchmark.
 				let extrinsic = sp_std::str::from_utf8(extrinsic)
@@ -840,6 +845,9 @@ macro_rules! impl_benchmark {
 					_ => return Err("Could not find extrinsic."),
 				};
 
+				// Add whitelist to DB
+				$crate::benchmarking::set_whitelist(whitelist.to_vec());
+
 				// Warm up the DB
 				$crate::benchmarking::commit_db();
 				$crate::benchmarking::wipe_db();
@@ -1074,7 +1082,7 @@ macro_rules! impl_benchmark_tests {
 #[macro_export]
 macro_rules! add_benchmark {
 	( $params:ident, $batches:ident, $name:literal, $( $location:tt )* ) => (
-		let (pallet, benchmark, lowest_range_values, highest_range_values, steps, repeat) = $params;
+		let (pallet, benchmark, lowest_range_values, highest_range_values, steps, repeat, whitelist) = $params;
 		if &pallet[..] == &$name[..] || &pallet[..] == &b"*"[..] {
 			if &pallet[..] == &b"*"[..] || &benchmark[..] == &b"*"[..] {
 				for benchmark in $( $location )*::benchmarks().into_iter() {
@@ -1085,6 +1093,7 @@ macro_rules! add_benchmark {
 							&highest_range_values[..],
 							&steps[..],
 							repeat,
+							whitelist,
 						)?,
 						pallet: $name.to_vec(),
 						benchmark: benchmark.to_vec(),
@@ -1098,6 +1107,7 @@ macro_rules! add_benchmark {
 						&highest_range_values[..],
 						&steps[..],
 						repeat,
+						whitelist,
 					)?,
 					pallet: $name.to_vec(),
 					benchmark: benchmark.clone(),

frame/benchmarking/src/utils.rs

@@ -102,6 +102,10 @@ pub trait Benchmarking {
 	fn reset_read_write_count(&mut self) {
 		self.reset_read_write_count()
 	}
+
+	fn set_whitelist(&mut self, new: Vec<Vec<u8>>) {
+		self.set_whitelist(new)
+	}
 }
 
 /// The pallet benchmarking trait.
@@ -125,6 +129,7 @@ pub trait Benchmarking<T> {
 		highest_range_values: &[u32],
 		steps: &[u32],
 		repeat: u32,
+		whitelist: &[Vec<u8>]
 	) -> Result<Vec<T>, &'static str>;
 }
 

primitives/externalities/src/lib.rs

@@ -224,6 +224,13 @@ pub trait Externalities: ExtensionStore {
 	///
 	/// Resets read/write count for the benchmarking process.
 	fn reset_read_write_count(&mut self);
+
+	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+	/// Benchmarking related functionality and shouldn't be used anywhere else!
+	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+	///
+	/// Adds new storage keys to the DB tracking whitelist.
+	fn set_whitelist(&mut self, new: Vec<Vec<u8>>);
 }
 
 /// Extension for the [`Externalities`] trait.

primitives/state-machine/src/backend.rs

@@ -225,6 +225,11 @@ pub trait Backend<H: Hasher>: std::fmt::Debug {
 	fn reset_read_write_count(&self) {
 		unimplemented!()
 	}
+
+	/// Update the whitelist for tracking db reads/writes
+	fn set_whitelist(&self, _: Vec<Vec<u8>>) {
+		unimplemented!()
+	}
 }
 
 impl<'a, T: Backend<H>, H: Hasher> Backend<H> for &'a T {

primitives/state-machine/src/basic.rs

@@ -318,6 +318,10 @@ impl Externalities for BasicExternalities {
 	fn reset_read_write_count(&mut self) {
 		unimplemented!("reset_read_write_count is not supported in Basic")
 	}
+
+	fn set_whitelist(&mut self, _: Vec<Vec<u8>>) {
+		unimplemented!("set_whitelist is not supported in Basic")
+	}
 }
 
 impl sp_externalities::ExtensionStore for BasicExternalities {

primitives/state-machine/src/ext.rs

@@ -582,6 +582,10 @@ where
 	fn reset_read_write_count(&mut self) {
 		self.backend.reset_read_write_count()
 	}
+
+	fn set_whitelist(&mut self, new: Vec<Vec<u8>>) {
+		self.backend.set_whitelist(new)
+	}
 }
 
 

primitives/state-machine/src/read_only.rs

@@ -181,6 +181,10 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 	fn reset_read_write_count(&mut self) {
 		unimplemented!("reset_read_write_count is not supported in ReadOnlyExternalities")
 	}
+
+	fn set_whitelist(&mut self, _: Vec<Vec<u8>>) {
+		unimplemented!("set_whitelist is not supported in ReadOnlyExternalities")
+	}
 }
 
 impl<'a, H: Hasher, B: 'a + Backend<H>> sp_externalities::ExtensionStore for ReadOnlyExternalities<'a, H, B> {