Fix benchmark read/write key tracker for keys in child storages. (#6905)

* WIP: read child trie and write child trie

* add test

* refactor a bit + improve log

* better naming

* trigger CI

* Revert "trigger CI"

This reverts commit d0aadae.

client/db/src/bench.rs

@@ -30,7 +30,9 @@ use sp_core::{
 };
 use sp_runtime::traits::{Block as BlockT, HashFor};
 use sp_runtime::Storage;
-use sp_state_machine::{DBValue, backend::Backend as StateBackend, StorageCollection};
+use sp_state_machine::{
+	DBValue, backend::Backend as StateBackend, StorageCollection, ChildStorageCollection
+};
 use kvdb::{KeyValueDB, DBTransaction};
 use crate::storage_cache::{CachingState, SharedCache, new_shared_cache};
 
@@ -96,7 +98,11 @@ pub struct BenchmarkingState<B: BlockT> {
 	genesis: HashMap<Vec<u8>, (Vec<u8>, i32)>,
 	record: Cell<Vec<Vec<u8>>>,
 	shared_cache: SharedCache<B>, // shared cache is always empty
-	key_tracker: RefCell<HashMap<Vec<u8>, KeyTracker>>,
+	/// Key tracker for keys in the main trie.
+	main_key_tracker: RefCell<HashMap<Vec<u8>, KeyTracker>>,
+	/// Key tracker for keys in a child trie.
+	/// Child trie are identified by their storage key (i.e. `ChildInfo::storage_key()`)
+	child_key_tracker: RefCell<HashMap<Vec<u8>, HashMap<Vec<u8>, KeyTracker>>>,
 	read_write_tracker: RefCell<ReadWriteTracker>,
 	whitelist: RefCell<Vec<TrackedStorageKey>>,
 }
@@ -116,7 +122,8 @@ impl<B: BlockT> BenchmarkingState<B> {
 			genesis_root: Default::default(),
 			record: Default::default(),
 			shared_cache: new_shared_cache(0, (1, 10)),
-			key_tracker: Default::default(),
+			main_key_tracker: Default::default(),
+			child_key_tracker: Default::default(),
 			read_write_tracker: Default::default(),
 			whitelist: Default::default(),
 		};
@@ -134,7 +141,7 @@ impl<B: BlockT> BenchmarkingState<B> {
 		);
 		state.genesis = transaction.clone().drain();
 		state.genesis_root = root.clone();
-		state.commit(root, transaction, Vec::new())?;
+		state.commit(root, transaction, Vec::new(), Vec::new())?;
 		state.record.take();
 		Ok(state)
 	}
@@ -156,7 +163,7 @@ impl<B: BlockT> BenchmarkingState<B> {
 	}
 
 	fn add_whitelist_to_tracker(&self) {
-		let mut key_tracker = self.key_tracker.borrow_mut();
+		let mut main_key_tracker = self.main_key_tracker.borrow_mut();
 
 		let whitelist = self.whitelist.borrow();
 
@@ -165,32 +172,37 @@ impl<B: BlockT> BenchmarkingState<B> {
 				has_been_read: key.has_been_read,
 				has_been_written: key.has_been_written,
 			};
-			key_tracker.insert(key.key.clone(), whitelisted);
+			main_key_tracker.insert(key.key.clone(), whitelisted);
 		});
 	}
 
 	fn wipe_tracker(&self) {
-		*self.key_tracker.borrow_mut() = HashMap::new();
+		*self.main_key_tracker.borrow_mut() = HashMap::new();
 		self.add_whitelist_to_tracker();
 		*self.read_write_tracker.borrow_mut() = Default::default();
 	}
 
-	fn add_read_key(&self, key: &[u8]) {
-		log::trace!(target: "benchmark", "Read: {}", HexDisplay::from(&key));
-
-		let mut key_tracker = self.key_tracker.borrow_mut();
+	// Childtrie is identified by its storage key (i.e. `ChildInfo::storage_key`)
+	fn add_read_key(&self, childtrie: Option<&[u8]>, key: &[u8]) {
 		let mut read_write_tracker = self.read_write_tracker.borrow_mut();
+		let mut child_key_tracker = self.child_key_tracker.borrow_mut();
+		let mut main_key_tracker = self.main_key_tracker.borrow_mut();
 
-		let maybe_tracker = key_tracker.get(key);
+		let key_tracker = if let Some(childtrie) = childtrie {
+			child_key_tracker.entry(childtrie.to_vec()).or_insert_with(|| HashMap::new())
+	 	} else {
+			&mut main_key_tracker
+		};
 
-		match maybe_tracker {
+		let read = match key_tracker.get(key) {
 			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();
+				true
 			},
 			Some(tracker) => {
 				if !tracker.has_been_read {
@@ -200,40 +212,71 @@ impl<B: BlockT> BenchmarkingState<B> {
 					};
 					key_tracker.insert(key.to_vec(), has_been_read);
 					read_write_tracker.add_read();
+					true
 				} else {
 					read_write_tracker.add_repeat_read();
+					false
 				}
 			}
+		};
+
+		if read {
+			if let Some(childtrie) = childtrie {
+				log::trace!(
+					target: "benchmark",
+					"Childtrie Read: {} {}", HexDisplay::from(&childtrie), HexDisplay::from(&key)
+				);
+			} else {
+				log::trace!(target: "benchmark", "Read: {}", HexDisplay::from(&key));
+			}
 		}
 	}
 
-	fn add_write_key(&self, key: &[u8]) {
-		log::trace!(target: "benchmark", "Write: {}", HexDisplay::from(&key));
-
-		let mut key_tracker = self.key_tracker.borrow_mut();
+	// Childtrie is identified by its storage key (i.e. `ChildInfo::storage_key`)
+	fn add_write_key(&self, childtrie: Option<&[u8]>, key: &[u8]) {
 		let mut read_write_tracker = self.read_write_tracker.borrow_mut();
+		let mut child_key_tracker = self.child_key_tracker.borrow_mut();
+		let mut main_key_tracker = self.main_key_tracker.borrow_mut();
 
-		let maybe_tracker = key_tracker.get(key);
+		let key_tracker = if let Some(childtrie) = childtrie {
+			child_key_tracker.entry(childtrie.to_vec()).or_insert_with(|| HashMap::new())
+	 	} else {
+			&mut main_key_tracker
+		};
 
 		// If we have written to the key, we also consider that we have read from it.
 		let has_been_written = KeyTracker {
 			has_been_read: true,
 			has_been_written: true,
 		};
 
-		match maybe_tracker {
+		let write = match key_tracker.get(key) {
 			None => {
 				key_tracker.insert(key.to_vec(), has_been_written);
 				read_write_tracker.add_write();
+				true
 			},
 			Some(tracker) => {
 				if !tracker.has_been_written {
 					key_tracker.insert(key.to_vec(), has_been_written);
 					read_write_tracker.add_write();
+					true
 				} else {
 					read_write_tracker.add_repeat_write();
+					false
 				}
 			}
+		};
+
+		if write {
+			if let Some(childtrie) = childtrie {
+				log::trace!(
+					target: "benchmark",
+					"Childtrie Write: {} {}", HexDisplay::from(&childtrie), HexDisplay::from(&key)
+				);
+			} else {
+				log::trace!(target: "benchmark", "Write: {}", HexDisplay::from(&key));
+			}
 		}
 	}
 }
@@ -248,12 +291,12 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 	type TrieBackendStorage = <DbState<B> as StateBackend<HashFor<B>>>::TrieBackendStorage;
 
 	fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
-		self.add_read_key(key);
+		self.add_read_key(None, key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.storage(key)
 	}
 
 	fn storage_hash(&self, key: &[u8]) -> Result<Option<B::Hash>, Self::Error> {
-		self.add_read_key(key);
+		self.add_read_key(None, key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.storage_hash(key)
 	}
 
@@ -262,12 +305,12 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 		child_info: &ChildInfo,
 		key: &[u8],
 	) -> Result<Option<Vec<u8>>, Self::Error> {
-		self.add_read_key(key);
+		self.add_read_key(Some(child_info.storage_key()), key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.child_storage(child_info, key)
 	}
 
 	fn exists_storage(&self, key: &[u8]) -> Result<bool, Self::Error> {
-		self.add_read_key(key);
+		self.add_read_key(None, key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.exists_storage(key)
 	}
 
@@ -276,12 +319,12 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 		child_info: &ChildInfo,
 		key: &[u8],
 	) -> Result<bool, Self::Error> {
-		self.add_read_key(key);
+		self.add_read_key(Some(child_info.storage_key()), key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.exists_child_storage(child_info, key)
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
-		self.add_read_key(key);
+		self.add_read_key(None, key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.next_storage_key(key)
 	}
 
@@ -290,7 +333,7 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 		child_info: &ChildInfo,
 		key: &[u8],
 	) -> Result<Option<Vec<u8>>, Self::Error> {
-		self.add_read_key(key);
+		self.add_read_key(Some(child_info.storage_key()), key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.next_child_storage_key(child_info, key)
 	}
 
@@ -367,9 +410,9 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 	fn commit(&self,
 		storage_root: <HashFor<B> as Hasher>::Out,
 		mut transaction: Self::Transaction,
-		storage_changes: StorageCollection,
-	) -> Result<(), Self::Error>
-	{
+		main_storage_changes: StorageCollection,
+		child_storage_changes: ChildStorageCollection,
+	) -> Result<(), Self::Error> {
 		if let Some(db) = self.db.take() {
 			let mut db_transaction = DBTransaction::new();
 			let changes = transaction.drain();
@@ -390,8 +433,13 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 			self.db.set(Some(db));
 
 			// Track DB Writes
-			storage_changes.iter().for_each(|(key, _)| {
-				self.add_write_key(key);
+			main_storage_changes.iter().for_each(|(key, _)| {
+				self.add_write_key(None, key);
+			});
+			child_storage_changes.iter().for_each(|(child_storage_key, storage_changes)| {
+				storage_changes.iter().for_each(|(key, _)| {
+					self.add_write_key(Some(child_storage_key), key);
+				})
 			});
 		} else {
 			return Err("Trying to commit to a closed db".into())
@@ -453,3 +501,43 @@ impl<Block: BlockT> std::fmt::Debug for BenchmarkingState<Block> {
 		write!(f, "Bench DB")
 	}
 }
+
+#[cfg(test)]
+mod test {
+	use crate::bench::BenchmarkingState;
+	use sp_state_machine::backend::Backend as _;
+
+	#[test]
+	fn read_to_main_and_child_tries() {
+		let bench_state = BenchmarkingState::<crate::tests::Block>::new(Default::default(), None)
+			.unwrap();
+
+		let child1 = sp_core::storage::ChildInfo::new_default(b"child1");
+		let child2 = sp_core::storage::ChildInfo::new_default(b"child2");
+
+		bench_state.storage(b"foo").unwrap();
+		bench_state.child_storage(&child1, b"foo").unwrap();
+		bench_state.child_storage(&child2, b"foo").unwrap();
+
+		bench_state.storage(b"bar").unwrap();
+		bench_state.child_storage(&child1, b"bar").unwrap();
+		bench_state.child_storage(&child2, b"bar").unwrap();
+
+		bench_state.commit(
+			Default::default(),
+			Default::default(),
+			vec![
+				("foo".as_bytes().to_vec(), None)
+			],
+			vec![
+				("child1".as_bytes().to_vec(), vec![("foo".as_bytes().to_vec(), None)])
+			]
+		).unwrap();
+
+		let rw_tracker = bench_state.read_write_tracker.borrow();
+		assert_eq!(rw_tracker.reads, 6);
+		assert_eq!(rw_tracker.repeat_reads, 0);
+		assert_eq!(rw_tracker.writes, 2);
+		assert_eq!(rw_tracker.repeat_writes, 0);
+	}
+}

primitives/state-machine/src/backend.rs

@@ -26,7 +26,7 @@ use sp_core::{
 use crate::{
 	trie_backend::TrieBackend,
 	trie_backend_essence::TrieBackendStorage,
-	UsageInfo, StorageKey, StorageValue, StorageCollection,
+	UsageInfo, StorageKey, StorageValue, StorageCollection, ChildStorageCollection,
 };
 
 /// A state backend is used to read state data and can have changes committed
@@ -215,7 +215,13 @@ pub trait Backend<H: Hasher>: std::fmt::Debug {
 	}
 
 	/// Commit given transaction to storage.
-	fn commit(&self, _: H::Out, _: Self::Transaction, _: StorageCollection) -> Result<(), Self::Error> {
+	fn commit(
+		&self,
+		_: H::Out,
+		_: Self::Transaction,
+		_: StorageCollection,
+		_: ChildStorageCollection,
+	) -> Result<(), Self::Error> {
 		unimplemented!()
 	}
 

primitives/state-machine/src/ext.rs

@@ -594,6 +594,7 @@ where
 			changes.transaction_storage_root,
 			changes.transaction,
 			changes.main_storage_changes,
+			changes.child_storage_changes,
 		).expect(EXT_NOT_ALLOWED_TO_FAIL);
 		self.mark_dirty();
 		self.overlay