Additional benchmark-storage flags (#11004)

* Fix typos

* Enable overwriting handlebars template

* Optionally name json output or disable json altogether

* Don't write to json by default

* Include block id in handlebars output

* Include warmups for write benchmarks

* PR comments

* Drop unnecessary file extension

* Use more appropriate types

* Use more appropriate error message

* More use of more appropriate types

* Rework write benchmark warmups

* Run same benchmark for both read and write

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

@@ -20,6 +20,7 @@ use sc_client_api::{Backend as ClientBackend, StorageProvider, UsageProvider};
 use sc_client_db::DbHash;
 use sc_service::Configuration;
 use sp_blockchain::HeaderBackend;
+use sp_core::storage::StorageKey;
 use sp_database::{ColumnId, Database};
 use sp_runtime::traits::{Block as BlockT, HashFor};
 use sp_state_machine::Storage;
@@ -29,7 +30,8 @@ use clap::{Args, Parser};
 use log::info;
 use rand::prelude::*;
 use serde::Serialize;
-use std::{fmt::Debug, sync::Arc};
+use sp_runtime::generic::BlockId;
+use std::{fmt::Debug, path::PathBuf, sync::Arc};
 
 use super::{record::StatSelect, template::TemplateData};
 
@@ -58,8 +60,8 @@ pub struct StorageCmd {
 pub struct StorageParams {
 	/// Path to write the *weight* file to. Can be a file or directory.
 	/// For substrate this should be `frame/support/src/weights`.
-	#[clap(long, default_value = ".")]
-	pub weight_path: String,
+	#[clap(long)]
+	pub weight_path: Option<PathBuf>,
 
 	/// Select a specific metric to calculate the final weight output.
 	#[clap(long = "metric", default_value = "average")]
@@ -83,8 +85,19 @@ pub struct StorageParams {
 	#[clap(long)]
 	pub skip_write: bool,
 
+	/// Specify the Handlebars template to use for outputting benchmark results.
+	#[clap(long)]
+	pub template_path: Option<PathBuf>,
+
+	/// Path to write the raw 'read' results in JSON format to. Can be a file or directory.
+	#[clap(long)]
+	pub json_read_path: Option<PathBuf>,
+
+	/// Path to write the raw 'write' results in JSON format to. Can be a file or directory.
+	#[clap(long)]
+	pub json_write_path: Option<PathBuf>,
+
 	/// Rounds of warmups before measuring.
-	/// Only supported for `read` benchmarks.
 	#[clap(long, default_value = "1")]
 	pub warmups: u32,
 
@@ -115,23 +128,32 @@ impl StorageCmd {
 	{
 		let mut template = TemplateData::new(&cfg, &self.params);
 
+		let block_id = BlockId::<Block>::Number(client.usage_info().chain.best_number);
+		template.set_block_number(block_id.to_string());
+
 		if !self.params.skip_read {
+			self.bench_warmup(&client)?;
 			let record = self.bench_read(client.clone())?;
-			record.save_json(&cfg, "read")?;
+			if let Some(path) = &self.params.json_read_path {
+				record.save_json(&cfg, path, "read")?;
+			}
 			let stats = record.calculate_stats()?;
 			info!("Time summary [ns]:\n{:?}\nValue size summary:\n{:?}", stats.0, stats.1);
 			template.set_stats(Some(stats), None)?;
 		}
 
 		if !self.params.skip_write {
+			self.bench_warmup(&client)?;
 			let record = self.bench_write(client, db, storage)?;
-			record.save_json(&cfg, "write")?;
+			if let Some(path) = &self.params.json_write_path {
+				record.save_json(&cfg, path, "write")?;
+			}
 			let stats = record.calculate_stats()?;
 			info!("Time summary [ns]:\n{:?}\nValue size summary:\n{:?}", stats.0, stats.1);
 			template.set_stats(None, Some(stats))?;
 		}
 
-		template.write(&self.params.weight_path)
+		template.write(&self.params.weight_path, &self.params.template_path)
 	}
 
 	/// Returns the specified state version.
@@ -149,6 +171,33 @@ impl StorageCmd {
 		info!("Using seed {}", seed);
 		StdRng::seed_from_u64(seed)
 	}
+
+	/// 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<()>
+	where
+		C: UsageProvider<B> + StorageProvider<B, BA>,
+		B: BlockT + Debug,
+		BA: ClientBackend<B>,
+	{
+		let block = BlockId::Number(client.usage_info().chain.best_number);
+		let empty_prefix = StorageKey(Vec::new());
+		let mut keys = client.storage_keys(&block, &empty_prefix)?;
+		let mut rng = Self::setup_rng();
+		keys.shuffle(&mut rng);
+
+		for i in 0..self.params.warmups {
+			info!("Warmup round {}/{}", i + 1, self.params.warmups);
+			for key in keys.clone() {
+				let _ = client
+					.storage(&block, &key)
+					.expect("Checked above to exist")
+					.ok_or("Value unexpectedly empty");
+			}
+		}
+
+		Ok(())
+	}
 }
 
 // Boilerplate