Support Multiple Instances with Benchmarks (#7669)

* Support multiple instances with benchmarks

* fix tests

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <[email protected]>

* docs

* fix output

* Update lib.rs

Co-authored-by: Bastian Köcher <[email protected]>

frame/benchmarking/src/lib.rs

@@ -1052,10 +1052,29 @@ macro_rules! impl_benchmark_test {
 /// ```
 ///
 /// At the end of `dispatch_benchmark`, you should return this batches object.
+///
+/// In the case where you have multiple instances of a pallet that you need to separately benchmark,
+/// the name of your module struct will be used as a suffix to your outputted weight file. For
+/// example:
+///
+/// ```ignore
+/// add_benchmark!(params, batches, pallet_balances, Balances); // pallet_balances.rs
+/// add_benchmark!(params, batches, pallet_collective, Council); // pallet_collective_council.rs
+/// add_benchmark!(params, batches, pallet_collective, TechnicalCommittee); // pallet_collective_technical_committee.rs
+/// ```
+///
+/// You can manipulate this suffixed string by using a type alias if needed. For example:
+///
+/// ```ignore
+/// type Council2 = TechnicalCommittee;
+/// add_benchmark!(params, batches, pallet_collective, Council2); // pallet_collective_council_2.rs
+/// ```
+
 #[macro_export]
 macro_rules! add_benchmark {
 	( $params:ident, $batches:ident, $name:ident, $( $location:tt )* ) => (
 		let name_string = stringify!($name).as_bytes();
+		let instance_string = stringify!( $( $location )* ).as_bytes();
 		let (config, whitelist) = $params;
 		let $crate::BenchmarkConfig {
 			pallet,
@@ -1071,6 +1090,9 @@ macro_rules! add_benchmark {
 			if &pallet[..] == &b"*"[..] || &benchmark[..] == &b"*"[..] {
 				for benchmark in $( $location )*::benchmarks(*extra).into_iter() {
 					$batches.push($crate::BenchmarkBatch {
+						pallet: name_string.to_vec(),
+						instance: instance_string.to_vec(),
+						benchmark: benchmark.to_vec(),
 						results: $( $location )*::run_benchmark(
 							benchmark,
 							&lowest_range_values[..],
@@ -1080,12 +1102,13 @@ macro_rules! add_benchmark {
 							whitelist,
 							*verify,
 						)?,
-						pallet: name_string.to_vec(),
-						benchmark: benchmark.to_vec(),
 					});
 				}
 			} else {
 				$batches.push($crate::BenchmarkBatch {
+					pallet: name_string.to_vec(),
+					instance: instance_string.to_vec(),
+					benchmark: benchmark.clone(),
 					results: $( $location )*::run_benchmark(
 						&benchmark[..],
 						&lowest_range_values[..],
@@ -1095,8 +1118,6 @@ macro_rules! add_benchmark {
 						whitelist,
 						*verify,
 					)?,
-					pallet: name_string.to_vec(),
-					benchmark: benchmark.clone(),
 				});
 			}
 		}

frame/benchmarking/src/utils.rs

@@ -43,6 +43,8 @@ impl std::fmt::Display for BenchmarkParameter {
 pub struct BenchmarkBatch {
 	/// The pallet containing this benchmark.
 	pub pallet: Vec<u8>,
+	/// The instance of this pallet being benchmarked.
+	pub instance: Vec<u8>,
 	/// The extrinsic (or benchmark name) of this benchmark.
 	pub benchmark: Vec<u8>,
 	/// The results from this benchmark.

utils/frame/benchmarking-cli/Cargo.toml

@@ -23,11 +23,12 @@ sp-externalities = { version = "0.8.0", path = "../../../primitives/externalitie
 sp-keystore = { version = "0.8.0", path = "../../../primitives/keystore" }
 sp-runtime = { version = "2.0.0", path = "../../../primitives/runtime" }
 sp-state-machine = { version = "0.8.0", path = "../../../primitives/state-machine" }
-structopt = "0.3.8"
 codec = { version = "1.3.1", package = "parity-scale-codec" }
+structopt = "0.3.8"
 chrono = "0.4"
 serde = "1.0.116"
 handlebars = "3.5.0"
+Inflector = "0.11.4"
 
 [features]
 default = ["db"]

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

@@ -22,6 +22,7 @@ use std::fs;
 use std::path::PathBuf;
 
 use serde::Serialize;
+use inflector::Inflector;
 
 use crate::BenchmarkCmd;
 use frame_benchmarking::{BenchmarkBatch, BenchmarkSelector, Analysis, RegressionModel};
@@ -37,6 +38,7 @@ struct TemplateData {
 	date: String,
 	version: String,
 	pallet: String,
+	instance: String,
 	header: String,
 	cmd: CmdData,
 	benchmarks: Vec<BenchmarkData>,
@@ -102,7 +104,7 @@ fn io_error(s: &str) -> std::io::Error {
 // p1 -> [b1, b2, b3]
 // p2 -> [b1, b2]
 // ```
-fn map_results(batches: &[BenchmarkBatch]) -> Result<HashMap<String, Vec<BenchmarkData>>, std::io::Error> {
+fn map_results(batches: &[BenchmarkBatch]) -> Result<HashMap<(String, String), Vec<BenchmarkData>>, std::io::Error> {
 	// Skip if batches is empty.
 	if batches.is_empty() { return Err(io_error("empty batches")) }
 
@@ -115,20 +117,22 @@ fn map_results(batches: &[BenchmarkBatch]) -> Result<HashMap<String, Vec<Benchma
 		if batch.results.is_empty() { continue }
 
 		let pallet_string = String::from_utf8(batch.pallet.clone()).unwrap();
+		let instance_string = String::from_utf8(batch.instance.clone()).unwrap();
 		let benchmark_data = get_benchmark_data(batch);
 		pallet_benchmarks.push(benchmark_data);
 
 		// Check if this is the end of the iterator
 		if let Some(next) = batches_iter.peek() {
 			// Next pallet is different than current pallet, save and create new data.
 			let next_pallet = String::from_utf8(next.pallet.clone()).unwrap();
-			if next_pallet != pallet_string {
-				all_benchmarks.insert(pallet_string, pallet_benchmarks.clone());
+			let next_instance = String::from_utf8(next.instance.clone()).unwrap();
+			if next_pallet != pallet_string || next_instance != instance_string {
+				all_benchmarks.insert((pallet_string, instance_string), pallet_benchmarks.clone());
 				pallet_benchmarks = Vec::new();
 			}
 		} else {
 			// This is the end of the iterator, so push the final data.
-			all_benchmarks.insert(pallet_string, pallet_benchmarks.clone());
+			all_benchmarks.insert((pallet_string, instance_string), pallet_benchmarks.clone());
 		}
 	}
 	Ok(all_benchmarks)
@@ -272,23 +276,30 @@ pub fn write_results(
 
 	// Organize results by pallet into a JSON map
 	let all_results = map_results(batches)?;
-	for (pallet, results) in all_results.into_iter() {
+	for ((pallet, instance), results) in all_results.iter() {
 		let mut file_path = path.clone();
 		// If a user only specified a directory...
 		if file_path.is_dir() {
-			// Create new file: "path/to/pallet_name.rs".
-			file_path.push(&pallet);
+			// Check if there might be multiple instances benchmarked.
+			if all_results.keys().any(|(p, i)| p == pallet && i != instance) {
+				// Create new file: "path/to/pallet_name_instance_name.rs".
+				file_path.push(pallet.clone() + "_" + &instance.to_snake_case());
+			} else {
+				// Create new file: "path/to/pallet_name.rs".
+				file_path.push(pallet.clone());
+			}
 			file_path.set_extension("rs");
 		}
 
 		let hbs_data = TemplateData {
 			args: args.clone(),
 			date: date.clone(),
 			version: VERSION.to_string(),
-			pallet: pallet,
+			pallet: pallet.to_string(),
+			instance: instance.to_string(),
 			header: header_text.clone(),
 			cmd: cmd_data.clone(),
-			benchmarks: results,
+			benchmarks: results.clone(),
 		};
 
 		let mut output_file = fs::File::create(file_path)?;
@@ -393,6 +404,7 @@ mod test {
 
 		return BenchmarkBatch {
 			pallet: [pallet.to_vec(), b"_pallet".to_vec()].concat(),
+			instance: b"instance".to_vec(),
 			benchmark: [benchmark.to_vec(), b"_benchmark".to_vec()].concat(),
 			results,
 		}
@@ -445,15 +457,21 @@ mod test {
 			test_data(b"second", b"first", BenchmarkParameter::c, 3, 4),
 		]).unwrap();
 
-		let first_benchmark = &mapped_results.get("first_pallet").unwrap()[0];
+		let first_benchmark = &mapped_results.get(
+			&("first_pallet".to_string(), "instance".to_string())
+		).unwrap()[0];
 		assert_eq!(first_benchmark.name, "first_benchmark");
 		check_data(first_benchmark, "a", 10, 3);
 
-		let second_benchmark = &mapped_results.get("first_pallet").unwrap()[1];
+		let second_benchmark = &mapped_results.get(
+			&("first_pallet".to_string(), "instance".to_string())
+		).unwrap()[1];
 		assert_eq!(second_benchmark.name, "second_benchmark");
 		check_data(second_benchmark, "b", 9, 2);
 
-		let second_pallet_benchmark = &mapped_results.get("second_pallet").unwrap()[0];
+		let second_pallet_benchmark = &mapped_results.get(
+			&("second_pallet".to_string(), "instance".to_string())
+		).unwrap()[0];
 		assert_eq!(second_pallet_benchmark.name, "first_benchmark");
 		check_data(second_pallet_benchmark, "c", 3, 4);
 	}