Include StorageInfo in Benchmarking Pipeline (#9090)

* extend storageinfo

* extend_storage_info

* use vec

* add storage info to pipeline

* get read and written keys

* undo storageinfo move

* refactor keytracker

* return read / write count

* playing with key matching

* add basic `StorageInfo` constructor

* add whitelisted to returned info

* fix some test stuff

* pipe comments into benchmark data

* add_storage_comments

* add comments to template

* track only storage prefix

* Update frame/benchmarking/src/lib.rs

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_balances --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/balances/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* fix test

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_balances --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/balances/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* remove test logs

* add temp benchmark script

* Apply suggestions from code review

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

* remove keytracker and use trackedstoragekey

* add comment for unknown keys

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_timestamp --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/timestamp/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* remove duplicate comments with unknown keys

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_timestamp --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/timestamp/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_balances --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/balances/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* refactor bench tracker, and fix results

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_balances --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/balances/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* fix child tries in new tracker

* extra newline

* fix unused warning

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_timestamp --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/timestamp/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* fix master merge

* storage info usage refactor

* remove now unused

* fix refactor

* use a vec for prefix

* fix tests

* also update writer to use vec

* disable read and written keys for now

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=frame_system --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/system/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* Update frame/system/src/weights.rs

* fix test

* Delete weights.rs

* reset weights

Co-authored-by: Parity Bot <admin@parity.io>
Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

substrate/primitives/externalities/src/lib.rs

@@ -296,6 +296,13 @@ pub trait Externalities: ExtensionStore {
 	fn proof_size(&self) -> Option<u32> {
 		None
 	}
+
+	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+	/// Benchmarking related functionality and shouldn't be used anywhere else!
+	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+	///
+	/// Get all the keys that have been read or written to during the benchmark.
+	fn get_read_and_written_keys(&self) -> Vec<(Vec<u8>, u32, u32, bool)>;
 }
 
 /// Extension for the [`Externalities`] trait.

substrate/primitives/state-machine/src/backend.rs

@@ -267,6 +267,11 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 	fn proof_size(&self) -> Option<u32> {
 		unimplemented!()
 	}
+
+	/// Extend storage info for benchmarking db
+	fn get_read_and_written_keys(&self) -> Vec<(Vec<u8>, u32, u32, bool)> {
+		unimplemented!()
+	}
 }
 
 /// Trait that allows consolidate two transactions together.

substrate/primitives/state-machine/src/basic.rs

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

substrate/primitives/state-machine/src/ext.rs

@@ -749,6 +749,10 @@ where
 	fn proof_size(&self) -> Option<u32> {
 		self.backend.proof_size()
 	}
+
+	fn get_read_and_written_keys(&self) -> Vec<(Vec<u8>, u32, u32, bool)> {
+		self.backend.get_read_and_written_keys()
+	}
 }
 
 impl<'a, H, N, B> Ext<'a, H, N, B>

substrate/primitives/state-machine/src/read_only.rs

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

substrate/primitives/storage/src/lib.rs

@@ -45,17 +45,56 @@ impl AsRef<[u8]> for StorageKey {
 #[cfg_attr(feature = "std", derive(Hash, PartialOrd, Ord))]
 pub struct TrackedStorageKey {
 	pub key: Vec<u8>,
-	pub has_been_read: bool,
-	pub has_been_written: bool,
+	pub reads: u32,
+	pub writes: u32,
+	pub whitelisted: bool,
 }
 
-// Easily convert a key to a `TrackedStorageKey` that has been read and written to.
+impl TrackedStorageKey {
+	/// Create a default `TrackedStorageKey`
+	pub fn new(key: Vec<u8>) -> Self {
+		Self {
+			key,
+			reads: 0,
+			writes: 0,
+			whitelisted: false,
+		}
+	}
+	/// Check if this key has been "read", i.e. it exists in the memory overlay.
+	///
+	/// Can be true if the key has been read, has been written to, or has been
+	/// whitelisted.
+	pub fn has_been_read(&self) -> bool {
+		self.whitelisted || self.reads > 0u32 || self.has_been_written()
+	}
+	/// Check if this key has been "written", i.e. a new value will be committed to the database.
+	///
+	/// Can be true if the key has been written to, or has been whitelisted.
+	pub fn has_been_written(&self) -> bool {
+		self.whitelisted || self.writes > 0u32
+	}
+	/// Add a storage read to this key.
+	pub fn add_read(&mut self) {
+		self.reads += 1;
+	}
+	/// Add a storage write to this key.
+	pub fn add_write(&mut self) {
+		self.writes += 1;
+	}
+	/// Whitelist this key.
+	pub fn whitelist(&mut self) {
+		self.whitelisted = true;
+	}
+}
+
+// Easily convert a key to a `TrackedStorageKey` that has been whitelisted.
 impl From<Vec<u8>> for TrackedStorageKey {
 	fn from(key: Vec<u8>) -> Self {
 		Self {
 			key: key,
-			has_been_read: true,
-			has_been_written: true,
+			reads: 0,
+			writes: 0,
+			whitelisted: true,
 		}
 	}
 }

substrate/primitives/tasks/src/async_externalities.rs

@@ -190,6 +190,10 @@ impl Externalities for AsyncExternalities {
 	fn set_whitelist(&mut self, _: Vec<TrackedStorageKey>) {
 		unimplemented!("set_whitelist is not supported in AsyncExternalities")
 	}
+
+	fn get_read_and_written_keys(&self) -> Vec<(Vec<u8>, u32, u32, bool)> {
+		unimplemented!("get_read_and_written_keys is not supported in AsyncExternalities")
+	}
 }
 
 impl sp_externalities::ExtensionStore for AsyncExternalities {