benchmark-cli: add child tree support (#12021)

* benchmark-cli: add child tree support

* removed extra comments

* addressed pr comments

* clean up

* addressed pr comments

substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs

@@ -24,7 +24,7 @@ use sp_core::storage::StorageKey;
 use sp_database::{ColumnId, Database};
 use sp_runtime::traits::{Block as BlockT, HashFor};
 use sp_state_machine::Storage;
-use sp_storage::StateVersion;
+use sp_storage::{ChildInfo, ChildType, PrefixedStorageKey, StateVersion};
 
 use clap::{Args, Parser};
 use log::info;
@@ -99,6 +99,10 @@ pub struct StorageParams {
 	/// State cache size.
 	#[clap(long, default_value = "0")]
 	pub state_cache_size: usize,
+
+	/// Include child trees in benchmark.
+	#[clap(long)]
+	pub include_child_trees: bool,
 }
 
 impl StorageCmd {
@@ -155,6 +159,16 @@ impl StorageCmd {
 		}
 	}
 
+	/// Returns Some if child node and None if regular
+	pub(crate) fn is_child_key(&self, key: Vec<u8>) -> Option<ChildInfo> {
+		if let Some((ChildType::ParentKeyId, storage_key)) =
+			ChildType::from_prefixed_key(&PrefixedStorageKey::new(key))
+		{
+			return Some(ChildInfo::new_default(storage_key))
+		}
+		None
+	}
+
 	/// Run some rounds of the (read) benchmark as warmup.
 	/// See `frame_benchmarking_cli::storage::read::bench_read` for detailed comments.
 	fn bench_warmup<B, BA, C>(&self, client: &Arc<C>) -> Result<()>
@@ -171,7 +185,7 @@ impl StorageCmd {
 
 		for i in 0..self.params.warmups {
 			info!("Warmup round {}/{}", i + 1, self.params.warmups);
-			for key in keys.clone() {
+			for key in keys.as_slice() {
 				let _ = client
 					.storage(&block, &key)
 					.expect("Checked above to exist")

substrate/utils/frame/benchmarking-cli/src/storage/read.rs

@@ -50,16 +50,43 @@ impl StorageCmd {
 		let (mut rng, _) = new_rng(None);
 		keys.shuffle(&mut rng);
 
+		let mut child_nodes = Vec::new();
 		// Interesting part here:
 		// Read all the keys in the database and measure the time it takes to access each.
 		info!("Reading {} keys", keys.len());
-		for key in keys.clone() {
-			let start = Instant::now();
-			let v = client
-				.storage(&block, &key)
-				.expect("Checked above to exist")
-				.ok_or("Value unexpectedly empty")?;
-			record.append(v.0.len(), start.elapsed())?;
+		for key in keys.as_slice() {
+			match (self.params.include_child_trees, self.is_child_key(key.clone().0)) {
+				(true, Some(info)) => {
+					// child tree key
+					let child_keys = client.child_storage_keys(&block, &info, &empty_prefix)?;
+					for ck in child_keys {
+						child_nodes.push((ck.clone(), info.clone()));
+					}
+				},
+				_ => {
+					// regular key
+					let start = Instant::now();
+					let v = client
+						.storage(&block, &key)
+						.expect("Checked above to exist")
+						.ok_or("Value unexpectedly empty")?;
+					record.append(v.0.len(), start.elapsed())?;
+				},
+			}
+		}
+
+		if self.params.include_child_trees {
+			child_nodes.shuffle(&mut rng);
+
+			info!("Reading {} child keys", child_nodes.len());
+			for (key, info) in child_nodes.as_slice() {
+				let start = Instant::now();
+				let v = client
+					.child_storage(&block, info, key)
+					.expect("Checked above to exist")
+					.ok_or("Value unexpectedly empty")?;
+				record.append(v.0.len(), start.elapsed())?;
+			}
 		}
 		Ok(record)
 	}

substrate/utils/frame/benchmarking-cli/src/storage/write.rs

@@ -16,7 +16,7 @@
 // limitations under the License.
 
 use sc_cli::Result;
-use sc_client_api::UsageProvider;
+use sc_client_api::{Backend as ClientBackend, StorageProvider, UsageProvider};
 use sc_client_db::{DbHash, DbState};
 use sp_api::StateBackend;
 use sp_blockchain::HeaderBackend;
@@ -29,7 +29,12 @@ use sp_trie::PrefixedMemoryDB;
 
 use log::{info, trace};
 use rand::prelude::*;
-use std::{fmt::Debug, sync::Arc, time::Instant};
+use sp_storage::{ChildInfo, StateVersion};
+use std::{
+	fmt::Debug,
+	sync::Arc,
+	time::{Duration, Instant},
+};
 
 use super::cmd::StorageCmd;
 use crate::shared::{new_rng, BenchRecord};
@@ -37,7 +42,7 @@ use crate::shared::{new_rng, BenchRecord};
 impl StorageCmd {
 	/// Benchmarks the time it takes to write a single Storage item.
 	/// Uses the latest state that is available for the given client.
-	pub(crate) fn bench_write<Block, H, C>(
+	pub(crate) fn bench_write<Block, BA, H, C>(
 		&self,
 		client: Arc<C>,
 		(db, state_col): (Arc<dyn sp_database::Database<DbHash>>, ColumnId),
@@ -46,7 +51,8 @@ impl StorageCmd {
 	where
 		Block: BlockT<Header = H, Hash = DbHash> + Debug,
 		H: HeaderT<Hash = DbHash>,
-		C: UsageProvider<Block> + HeaderBackend<Block>,
+		BA: ClientBackend<Block>,
+		C: UsageProvider<Block> + HeaderBackend<Block> + StorageProvider<Block, BA>,
 	{
 		// Store the time that it took to write each value.
 		let mut record = BenchRecord::default();
@@ -61,50 +67,96 @@ impl StorageCmd {
 		let mut kvs = trie.pairs();
 		let (mut rng, _) = new_rng(None);
 		kvs.shuffle(&mut rng);
+		info!("Writing {} keys", kvs.len());
+
+		let mut child_nodes = Vec::new();
 
 		// Generate all random values first; Make sure there are no collisions with existing
 		// db entries, so we can rollback all additions without corrupting existing entries.
-		for (k, original_v) in kvs.iter_mut() {
-			'retry: loop {
-				let mut new_v = vec![0; original_v.len()];
-				// Create a random value to overwrite with.
-				// NOTE: We use a possibly higher entropy than the original value,
-				// could be improved but acts as an over-estimation which is fine for now.
-				rng.fill_bytes(&mut new_v[..]);
-				let new_kv = vec![(k.as_ref(), Some(new_v.as_ref()))];
-				let (_, mut stx) = trie.storage_root(new_kv.iter().cloned(), self.state_version());
-				for (mut k, (_, rc)) in stx.drain().into_iter() {
-					if rc > 0 {
-						db.sanitize_key(&mut k);
-						if db.get(state_col, &k).is_some() {
-							trace!("Benchmark-store key creation: Key collision detected, retry");
-							continue 'retry
+		for (k, original_v) in kvs {
+			match (self.params.include_child_trees, self.is_child_key(k.to_vec())) {
+				(true, Some(info)) => {
+					let child_keys =
+						client.child_storage_keys_iter(&block, info.clone(), None, None)?;
+					for ck in child_keys {
+						child_nodes.push((ck.clone(), info.clone()));
+					}
+				},
+				_ => {
+					// regular key
+					let mut new_v = vec![0; original_v.len()];
+					loop {
+						// Create a random value to overwrite with.
+						// NOTE: We use a possibly higher entropy than the original value,
+						// could be improved but acts as an over-estimation which is fine for now.
+						rng.fill_bytes(&mut new_v[..]);
+						if check_new_value::<Block>(
+							db.clone(),
+							&trie,
+							&k.to_vec(),
+							&new_v,
+							self.state_version(),
+							state_col,
+							None,
+						) {
+							break
 						}
 					}
-				}
-				*original_v = new_v;
-				break
+
+					// Write each value in one commit.
+					let (size, duration) = measure_write::<Block>(
+						db.clone(),
+						&trie,
+						k.to_vec(),
+						new_v.to_vec(),
+						self.state_version(),
+						state_col,
+						None,
+					)?;
+					record.append(size, duration)?;
+				},
 			}
 		}
 
-		info!("Writing {} keys", kvs.len());
-		// Write each value in one commit.
-		for (k, new_v) in kvs.iter() {
-			// Interesting part here:
-			let start = Instant::now();
-			// Create a TX that will modify the Trie in the DB and
-			// calculate the root hash of the Trie after the modification.
-			let replace = vec![(k.as_ref(), Some(new_v.as_ref()))];
-			let (_, stx) = trie.storage_root(replace.iter().cloned(), self.state_version());
-			// Only the keep the insertions, since we do not want to benchmark pruning.
-			let tx = convert_tx::<Block>(db.clone(), stx.clone(), false, state_col);
-			db.commit(tx).map_err(|e| format!("Writing to the Database: {}", e))?;
-			record.append(new_v.len(), start.elapsed())?;
+		if self.params.include_child_trees {
+			child_nodes.shuffle(&mut rng);
+			info!("Writing {} child keys", child_nodes.len());
 
-			// Now undo the changes by removing what was added.
-			let tx = convert_tx::<Block>(db.clone(), stx.clone(), true, state_col);
-			db.commit(tx).map_err(|e| format!("Writing to the Database: {}", e))?;
+			for (key, info) in child_nodes {
+				if let Some(original_v) = client
+					.child_storage(&block, &info.clone(), &key)
+					.expect("Checked above to exist")
+				{
+					let mut new_v = vec![0; original_v.0.len()];
+					loop {
+						rng.fill_bytes(&mut new_v[..]);
+						if check_new_value::<Block>(
+							db.clone(),
+							&trie,
+							&key.0,
+							&new_v,
+							self.state_version(),
+							state_col,
+							Some(&info),
+						) {
+							break
+						}
+					}
+
+					let (size, duration) = measure_write::<Block>(
+						db.clone(),
+						&trie,
+						key.0,
+						new_v.to_vec(),
+						self.state_version(),
+						state_col,
+						Some(&info),
+					)?;
+					record.append(size, duration)?;
+				}
+			}
 		}
+
 		Ok(record)
 	}
 }
@@ -134,3 +186,62 @@ fn convert_tx<B: BlockT>(
 	}
 	ret
 }
+
+/// Measures write benchmark
+/// if `child_info` exist then it means this is a child tree key
+fn measure_write<Block: BlockT>(
+	db: Arc<dyn sp_database::Database<DbHash>>,
+	trie: &DbState<Block>,
+	key: Vec<u8>,
+	new_v: Vec<u8>,
+	version: StateVersion,
+	col: ColumnId,
+	child_info: Option<&ChildInfo>,
+) -> Result<(usize, Duration)> {
+	let start = Instant::now();
+	// Create a TX that will modify the Trie in the DB and
+	// calculate the root hash of the Trie after the modification.
+	let replace = vec![(key.as_ref(), Some(new_v.as_ref()))];
+	let stx = match child_info {
+		Some(info) => trie.child_storage_root(info, replace.iter().cloned(), version).2,
+		None => trie.storage_root(replace.iter().cloned(), version).1,
+	};
+	// Only the keep the insertions, since we do not want to benchmark pruning.
+	let tx = convert_tx::<Block>(db.clone(), stx.clone(), false, col);
+	db.commit(tx).map_err(|e| format!("Writing to the Database: {}", e))?;
+	let result = (new_v.len(), start.elapsed());
+
+	// Now undo the changes by removing what was added.
+	let tx = convert_tx::<Block>(db.clone(), stx.clone(), true, col);
+	db.commit(tx).map_err(|e| format!("Writing to the Database: {}", e))?;
+	Ok(result)
+}
+
+/// Checks if a new value causes any collision in tree updates
+/// returns true if there is no collision
+/// if `child_info` exist then it means this is a child tree key
+fn check_new_value<Block: BlockT>(
+	db: Arc<dyn sp_database::Database<DbHash>>,
+	trie: &DbState<Block>,
+	key: &Vec<u8>,
+	new_v: &Vec<u8>,
+	version: StateVersion,
+	col: ColumnId,
+	child_info: Option<&ChildInfo>,
+) -> bool {
+	let new_kv = vec![(key.as_ref(), Some(new_v.as_ref()))];
+	let mut stx = match child_info {
+		Some(info) => trie.child_storage_root(info, new_kv.iter().cloned(), version).2,
+		None => trie.storage_root(new_kv.iter().cloned(), version).1,
+	};
+	for (mut k, (_, rc)) in stx.drain().into_iter() {
+		if rc > 0 {
+			db.sanitize_key(&mut k);
+			if db.get(col, &k).is_some() {
+				trace!("Benchmark-store key creation: Key collision detected, retry");
+				return false
+			}
+		}
+	}
+	true
+}