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>

substrate/primitives/externalities/src/lib.rs

@@ -25,7 +25,7 @@
 
 use std::any::{Any, TypeId};
 
-use sp_storage::ChildInfo;
+use sp_storage::{ChildInfo, TrackedStorageKey};
 
 pub use scope_limited::{set_and_run_with_externalities, with_externalities};
 pub use extensions::{Extension, Extensions, ExtensionStore};
@@ -248,12 +248,19 @@ 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!
+	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+	///
+	/// Gets the current DB tracking whitelist.
+	fn get_whitelist(&self) -> Vec<TrackedStorageKey>;
+
 	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 	/// 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>>);
+	fn set_whitelist(&mut self, new: Vec<TrackedStorageKey>);
 }
 
 /// Extension for the [`Externalities`] trait.

substrate/primitives/runtime-interface/Cargo.toml

@@ -21,6 +21,7 @@ sp-externalities = { version = "0.8.0-rc5", optional = true, path = "../external
 codec = { package = "parity-scale-codec", version = "1.3.1", default-features = false }
 static_assertions = "1.0.0"
 primitive-types = { version = "0.7.0", default-features = false }
+sp-storage = { version = "2.0.0-rc5", default-features = false, path = "../storage" }
 
 [dev-dependencies]
 sp-runtime-interface-test-wasm = { version = "2.0.0-rc5", path = "test-wasm" }

substrate/primitives/runtime-interface/src/impls.rs

@@ -537,3 +537,7 @@ impl PassBy for sp_wasm_interface::ValueType {
 impl PassBy for sp_wasm_interface::Value {
 	type PassBy = Codec<sp_wasm_interface::Value>;
 }
+
+impl PassBy for sp_storage::TrackedStorageKey {
+	type PassBy = Codec<Self>;
+}

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

@@ -19,7 +19,10 @@
 
 use hash_db::Hasher;
 use codec::{Decode, Encode};
-use sp_core::{traits::RuntimeCode, storage::{ChildInfo, well_known_keys}};
+use sp_core::{
+	traits::RuntimeCode,
+	storage::{ChildInfo, well_known_keys, TrackedStorageKey}
+};
 use crate::{
 	trie_backend::TrieBackend,
 	trie_backend_essence::TrieBackendStorage,
@@ -226,10 +229,13 @@ pub trait Backend<H: Hasher>: std::fmt::Debug {
 		unimplemented!()
 	}
 
-	/// Update the whitelist for tracking db reads/writes
-	fn set_whitelist(&self, _: Vec<Vec<u8>>) {
-		unimplemented!()
+	/// Get the whitelist for tracking db reads/writes
+	fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
+		Default::default()
 	}
+
+	/// Update the whitelist for tracking db reads/writes
+	fn set_whitelist(&self, _: Vec<TrackedStorageKey>) {}
 }
 
 impl<'a, T: Backend<H>, H: Hasher> Backend<H> for &'a T {

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

@@ -27,7 +27,7 @@ use sp_trie::trie_types::Layout;
 use sp_core::{
 	storage::{
 		well_known_keys::is_child_storage_key, Storage,
-		ChildInfo, StorageChild,
+		ChildInfo, StorageChild, TrackedStorageKey,
 	},
 	traits::Externalities, Blake2Hasher,
 };
@@ -325,7 +325,11 @@ impl Externalities for BasicExternalities {
 		unimplemented!("reset_read_write_count is not supported in Basic")
 	}
 
-	fn set_whitelist(&mut self, _: Vec<Vec<u8>>) {
+	fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
+		unimplemented!("get_whitelist is not supported in Basic")
+	}
+
+	fn set_whitelist(&mut self, _: Vec<TrackedStorageKey>) {
 		unimplemented!("set_whitelist is not supported in Basic")
 	}
 }

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

@@ -26,7 +26,7 @@ use crate::{
 use hash_db::Hasher;
 use sp_core::{
 	offchain::storage::OffchainOverlayedChanges,
-	storage::{well_known_keys::is_child_storage_key, ChildInfo},
+	storage::{well_known_keys::is_child_storage_key, ChildInfo, TrackedStorageKey},
 	traits::Externalities, hexdisplay::HexDisplay,
 };
 use sp_trie::{trie_types::Layout, empty_child_trie_root};
@@ -609,7 +609,11 @@ where
 		self.backend.reset_read_write_count()
 	}
 
-	fn set_whitelist(&mut self, new: Vec<Vec<u8>>) {
+	fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
+		self.backend.get_whitelist()
+	}
+
+	fn set_whitelist(&mut self, new: Vec<TrackedStorageKey>) {
 		self.backend.set_whitelist(new)
 	}
 }

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

@@ -24,7 +24,7 @@ use std::{
 use crate::{Backend, StorageKey, StorageValue};
 use hash_db::Hasher;
 use sp_core::{
-	storage::ChildInfo,
+	storage::{ChildInfo, TrackedStorageKey},
 	traits::Externalities, Blake2Hasher,
 };
 use codec::Encode;
@@ -194,7 +194,11 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		unimplemented!("reset_read_write_count is not supported in ReadOnlyExternalities")
 	}
 
-	fn set_whitelist(&mut self, _: Vec<Vec<u8>>) {
+	fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
+		unimplemented!("get_whitelist is not supported in ReadOnlyExternalities")
+	}
+
+	fn set_whitelist(&mut self, _: Vec<TrackedStorageKey>) {
 		unimplemented!("set_whitelist is not supported in ReadOnlyExternalities")
 	}
 }

substrate/primitives/storage/Cargo.toml

@@ -18,7 +18,8 @@ serde = { version = "1.0.101", optional = true, features = ["derive"] }
 impl-serde = { version = "0.2.3", optional = true }
 ref-cast = "1.0.0"
 sp-debug-derive = { version = "2.0.0-rc5", path = "../debug-derive" }
+codec = { package = "parity-scale-codec", version = "1.3.1", default-features = false, features = ["derive"] }
 
 [features]
 default = [ "std" ]
-std = [ "sp-std/std", "serde", "impl-serde" ]
+std = [ "sp-std/std", "serde", "impl-serde", "codec/std" ]

substrate/primitives/storage/src/lib.rs

@@ -25,6 +25,7 @@ use sp_debug_derive::RuntimeDebug;
 
 use sp_std::{vec::Vec, ops::{Deref, DerefMut}};
 use ref_cast::RefCast;
+use codec::{Encode, Decode};
 
 /// Storage key.
 #[derive(PartialEq, Eq, RuntimeDebug)]
@@ -34,6 +35,26 @@ pub struct StorageKey(
 	pub Vec<u8>,
 );
 
+/// Storage key with read/write tracking information.
+#[derive(PartialEq, Eq, RuntimeDebug, Clone, Encode, Decode)]
+#[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,
+}
+
+// Easily convert a key to a `TrackedStorageKey` that has been read and written to.
+impl From<Vec<u8>> for TrackedStorageKey {
+	fn from(key: Vec<u8>) -> Self {
+		Self {
+			key: key,
+			has_been_read: true,
+			has_been_written: true,
+		}
+	}
+}
+
 /// Storage key of a child trie, it contains the prefix to the key.
 #[derive(PartialEq, Eq, RuntimeDebug)]
 #[cfg_attr(feature = "std", derive(Serialize, Deserialize, Hash, PartialOrd, Ord, Clone))]