Additional benchmark-storage flags (paritytech#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

utils/frame/benchmarking-cli/src/lib.rs

@@ -67,7 +67,7 @@ pub struct BenchmarkCmd {
 	#[clap(long = "json")]
 	pub json_output: bool,
 
-	/// Write the raw results in JSON format into the give file.
+	/// Write the raw results in JSON format into the given file.
 	#[clap(long, conflicts_with = "json-output")]
 	pub json_file: Option<PathBuf>,
 

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

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

@@ -49,17 +49,6 @@ impl StorageCmd {
 		let mut rng = Self::setup_rng();
 		keys.shuffle(&mut rng);
 
-		// Run some rounds of the benchmark as warmup.
-		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")?;
-			}
-		}
-
 		// Interesting part here:
 		// Read all the keys in the database and measure the time it takes to access each.
 		info!("Reading {} keys", keys.len());

utils/frame/benchmarking-cli/src/storage/record.rs

@@ -22,7 +22,7 @@ use sc_service::Configuration;
 
 use log::info;
 use serde::Serialize;
-use std::{fmt, fs, result, str::FromStr, time::Duration};
+use std::{fmt, fs, path::PathBuf, result, str::FromStr, time::Duration};
 
 /// Raw output of a Storage benchmark.
 #[derive(Debug, Default, Clone, Serialize)]
@@ -95,12 +95,18 @@ impl BenchRecord {
 		Ok((time, size)) // The swap of time/size here is intentional.
 	}
 
-	/// Saves the raw results in a json file in the current directory.
+	/// Unless a path is specified, saves the raw results in a json file in the current directory.
 	/// Prefixes it with the DB name and suffixed with `path_suffix`.
-	pub fn save_json(&self, cfg: &Configuration, path_suffix: &str) -> Result<()> {
-		let path = format!("{}_{}.json", cfg.database, path_suffix).to_lowercase();
+	pub fn save_json(&self, cfg: &Configuration, out_path: &PathBuf, suffix: &str) -> Result<()> {
+		let mut path = PathBuf::from(out_path);
+		if path.is_dir() || path.as_os_str().is_empty() {
+			path.push(&format!("{}_{}", cfg.database, suffix).to_lowercase());
+			path.set_extension("json");
+		}
+
 		let json = serde_json::to_string_pretty(&self)
 			.map_err(|e| format!("Serializing as JSON: {:?}", e))?;
+
 		fs::write(&path, json)?;
 		info!("Raw data written to {:?}", fs::canonicalize(&path)?);
 		Ok(())

utils/frame/benchmarking-cli/src/storage/template.rs

@@ -32,6 +32,8 @@ static TEMPLATE: &str = include_str!("./weights.hbs");
 pub(crate) struct TemplateData {
 	/// Name of the database used.
 	db_name: String,
+	/// Block number that was used.
+	block_number: String,
 	/// Name of the runtime. Taken from the chain spec.
 	runtime_name: String,
 	/// Version of the benchmarking CLI used.
@@ -85,28 +87,44 @@ impl TemplateData {
 		Ok(())
 	}
 
-	/// Filles out the `weights.hbs` HBS template with its own data.
+	/// Sets the block id that was used.
+	pub fn set_block_number(&mut self, block_number: String) {
+		self.block_number = block_number
+	}
+
+	/// Fills out the `weights.hbs` or specified HBS template with its own data.
 	/// Writes the result to `path` which can be a directory or file.
-	pub fn write(&self, path: &str) -> Result<()> {
+	pub fn write(&self, path: &Option<PathBuf>, hbs_template: &Option<PathBuf>) -> Result<()> {
 		let mut handlebars = handlebars::Handlebars::new();
 		// Format large integers with underscore.
 		handlebars.register_helper("underscore", Box::new(crate::writer::UnderscoreHelper));
 		// Don't HTML escape any characters.
 		handlebars.register_escape_fn(|s| -> String { s.to_string() });
+		// Use custom template if provided.
+		let template = match hbs_template {
+			Some(template) if template.is_file() => fs::read_to_string(template)?,
+			Some(_) => return Err("Handlebars template is not a valid file!".into()),
+			None => TEMPLATE.to_string(),
+		};
 
 		let out_path = self.build_path(path);
 		let mut fd = fs::File::create(&out_path)?;
 		info!("Writing weights to {:?}", fs::canonicalize(&out_path)?);
+
 		handlebars
-			.render_template_to_write(&TEMPLATE, &self, &mut fd)
+			.render_template_to_write(&template, &self, &mut fd)
 			.map_err(|e| format!("HBS template write: {:?}", e).into())
 	}
 
 	/// Builds a path for the weight file.
-	fn build_path(&self, weight_out: &str) -> PathBuf {
-		let mut path = PathBuf::from(weight_out);
-		if path.is_dir() {
-			path.push(format!("{}_weights.rs", self.db_name.to_lowercase()));
+	fn build_path(&self, weight_out: &Option<PathBuf>) -> PathBuf {
+		let mut path = match weight_out {
+			Some(p) => PathBuf::from(p),
+			None => PathBuf::new(),
+		};
+
+		if path.is_dir() || path.as_os_str().is_empty() {
+			path.push(format!("{}_weights", self.db_name.to_lowercase()));
 			path.set_extension("rs");
 		}
 		path

utils/frame/benchmarking-cli/src/storage/weights.hbs

@@ -19,6 +19,7 @@
 //! DATE: {{date}}
 //!
 //! DATABASE: `{{db_name}}`, RUNTIME: `{{runtime_name}}`
+//! BLOCK-NUM: `{{block_number}}`
 //! SKIP-WRITE: `{{params.skip_write}}`, SKIP-READ: `{{params.skip_read}}`, WARMUPS: `{{params.warmups}}`
 //! STATE-VERSION: `V{{params.state_version}}`, STATE-CACHE-SIZE: `{{params.state_cache_size}}`
 //! WEIGHT-PATH: `{{params.weight_path}}`