Add PoV Tracking to Benchmarking Pipeline (#8559)

* Added a function to estimate proof size for benchmarking

* integrate proof_size into benchmarking pipeline

* Update client/db/src/bench.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update client/db/src/bench.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* fix tests

* one more test

* Update bench.rs

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

Co-authored-by: Alexander Popiak <alexander.popiak@parity.io>

* Update utils/frame/benchmarking-cli/src/command.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

Co-authored-by: arkpar <arkady.paronyan@gmail.com>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
Co-authored-by: Alexander Popiak <alexander.popiak@parity.io>
Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

client/db/src/bench.rs

@@ -23,17 +23,18 @@ use std::cell::{Cell, RefCell};
 use std::collections::HashMap;
 
 use hash_db::{Prefix, Hasher};
-use sp_trie::{MemoryDB, prefixed_key};
+use sp_trie::{MemoryDB, prefixed_key, StorageProof};
 use sp_core::{
 	storage::{ChildInfo, TrackedStorageKey},
 	hexdisplay::HexDisplay
 };
 use sp_runtime::traits::{Block as BlockT, HashFor};
 use sp_runtime::Storage;
 use sp_state_machine::{
-	DBValue, backend::Backend as StateBackend, StorageCollection, ChildStorageCollection
+	DBValue, backend::Backend as StateBackend, StorageCollection, ChildStorageCollection, ProofRecorder,
 };
 use kvdb::{KeyValueDB, DBTransaction};
+use codec::Encode;
 use crate::storage_cache::{CachingState, SharedCache, new_shared_cache};
 
 type DbState<B> = sp_state_machine::TrieBackend<
@@ -44,14 +45,25 @@ type State<B> = CachingState<DbState<B>, B>;
 
 struct StorageDb<Block: BlockT> {
 	db: Arc<dyn KeyValueDB>,
+	proof_recorder: Option<ProofRecorder<HashFor<Block>>>,
 	_block: std::marker::PhantomData<Block>,
 }
 
 impl<Block: BlockT> sp_state_machine::Storage<HashFor<Block>> for StorageDb<Block> {
 	fn get(&self, key: &Block::Hash, prefix: Prefix) -> Result<Option<DBValue>, String> {
-		let key = prefixed_key::<HashFor<Block>>(key, prefix);
-		self.db.get(0, &key)
-			.map_err(|e| format!("Database backend error: {:?}", e))
+		let prefixed_key = prefixed_key::<HashFor<Block>>(key, prefix);
+		if let Some(recorder) = &self.proof_recorder {
+			if let Some(v) = recorder.read().get(&key) {
+				return Ok(v.clone());
+			}
+			let backend_value = self.db.get(0, &prefixed_key)
+				.map_err(|e| format!("Database backend error: {:?}", e))?;
+			recorder.write().insert(key.clone(), backend_value.clone());
+			Ok(backend_value)
+		} else {
+			self.db.get(0, &prefixed_key)
+				.map_err(|e| format!("Database backend error: {:?}", e))
+		}
 	}
 }
 
@@ -105,11 +117,12 @@ pub struct BenchmarkingState<B: BlockT> {
 	child_key_tracker: RefCell<HashMap<Vec<u8>, HashMap<Vec<u8>, KeyTracker>>>,
 	read_write_tracker: RefCell<ReadWriteTracker>,
 	whitelist: RefCell<Vec<TrackedStorageKey>>,
+	proof_recorder: Option<ProofRecorder<HashFor<B>>>,
 }
 
 impl<B: BlockT> BenchmarkingState<B> {
 	/// Create a new instance that creates a database in a temporary dir.
-	pub fn new(genesis: Storage, _cache_size_mb: Option<usize>) -> Result<Self, String> {
+	pub fn new(genesis: Storage, _cache_size_mb: Option<usize>, record_proof: bool) -> Result<Self, String> {
 		let mut root = B::Hash::default();
 		let mut mdb = MemoryDB::<HashFor<B>>::default();
 		sp_state_machine::TrieDBMut::<HashFor<B>>::new(&mut mdb, &mut root);
@@ -126,6 +139,7 @@ impl<B: BlockT> BenchmarkingState<B> {
 			child_key_tracker: Default::default(),
 			read_write_tracker: Default::default(),
 			whitelist: Default::default(),
+			proof_recorder: record_proof.then(Default::default),
 		};
 
 		state.add_whitelist_to_tracker();
@@ -153,7 +167,14 @@ impl<B: BlockT> BenchmarkingState<B> {
 			None => Arc::new(::kvdb_memorydb::create(1)),
 		};
 		self.db.set(Some(db.clone()));
-		let storage_db = Arc::new(StorageDb::<B> { db, _block: Default::default() });
+		if let Some(recorder) = &self.proof_recorder {
+			recorder.write().clear();
+		}
+		let storage_db = Arc::new(StorageDb::<B> {
+			db,
+			proof_recorder: self.proof_recorder.clone(),
+			_block: Default::default()
+		});
 		*self.state.borrow_mut() = Some(State::new(
 			DbState::<B>::new(storage_db, self.root.get()),
 			self.shared_cache.clone(),
@@ -495,6 +516,17 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 	fn usage_info(&self) -> sp_state_machine::UsageInfo {
 		self.state.borrow().as_ref().map_or(sp_state_machine::UsageInfo::empty(), |s| s.usage_info())
 	}
+
+	fn proof_size(&self) -> Option<u32> {
+		self.proof_recorder.as_ref().map(|recorder| {
+			let proof = StorageProof::new(recorder
+				.read()
+				.iter()
+				.filter_map(|(_k, v)| v.as_ref().map(|v| v.to_vec()))
+				.collect());
+			proof.encoded_size() as u32
+		})
+	}
 }
 
 impl<Block: BlockT> std::fmt::Debug for BenchmarkingState<Block> {
@@ -510,7 +542,7 @@ mod test {
 
 	#[test]
 	fn read_to_main_and_child_tries() {
-		let bench_state = BenchmarkingState::<crate::tests::Block>::new(Default::default(), None)
+		let bench_state = BenchmarkingState::<crate::tests::Block>::new(Default::default(), None, false)
 			.unwrap();
 
 		for _ in 0..2 {

frame/benchmarking/src/analysis.rs

@@ -38,6 +38,7 @@ pub enum BenchmarkSelector {
 	StorageRootTime,
 	Reads,
 	Writes,
+	ProofSize,
 }
 
 #[derive(Debug)]
@@ -86,6 +87,7 @@ impl Analysis {
 				BenchmarkSelector::StorageRootTime => result.storage_root_time,
 				BenchmarkSelector::Reads => result.reads.into(),
 				BenchmarkSelector::Writes => result.writes.into(),
+				BenchmarkSelector::ProofSize => result.proof_size.into(),
 			}
 		).collect();
 
@@ -126,6 +128,7 @@ impl Analysis {
 						BenchmarkSelector::StorageRootTime => result.storage_root_time,
 						BenchmarkSelector::Reads => result.reads.into(),
 						BenchmarkSelector::Writes => result.writes.into(),
+						BenchmarkSelector::ProofSize => result.proof_size.into(),
 					};
 					(result.components[i].1, data)
 				})
@@ -190,6 +193,7 @@ impl Analysis {
 					BenchmarkSelector::StorageRootTime => result.storage_root_time,
 					BenchmarkSelector::Reads => result.reads.into(),
 					BenchmarkSelector::Writes => result.writes.into(),
+					BenchmarkSelector::ProofSize => result.proof_size.into(),
 				})
 		}
 
@@ -370,6 +374,7 @@ mod tests {
 			repeat_reads: 0,
 			writes,
 			repeat_writes: 0,
+			proof_size: 0,
 		}
 	}
 

frame/benchmarking/src/lib.rs

@@ -764,12 +764,21 @@ macro_rules! impl_benchmark {
 								"Start Benchmark: {:?}", c
 							);
 
+							let start_pov = $crate::benchmarking::proof_size();
 							let start_extrinsic = $crate::benchmarking::current_time();
 
 							closure_to_benchmark()?;
 
 							let finish_extrinsic = $crate::benchmarking::current_time();
-							let elapsed_extrinsic = finish_extrinsic - start_extrinsic;
+							let end_pov = $crate::benchmarking::proof_size();
+
+							// Calculate the diff caused by the benchmark.
+							let elapsed_extrinsic = finish_extrinsic.saturating_sub(start_extrinsic);
+							let diff_pov = match (start_pov, end_pov) {
+								(Some(start), Some(end)) => end.saturating_sub(start),
+								_ => Default::default(),
+							};
+
 							// Commit the changes to get proper write count
 							$crate::benchmarking::commit_db();
 							$crate::log::trace!(
@@ -796,6 +805,7 @@ macro_rules! impl_benchmark {
 								repeat_reads: read_write_count.1,
 								writes: read_write_count.2,
 								repeat_writes: read_write_count.3,
+								proof_size: diff_pov,
 							});
 						}
 

frame/benchmarking/src/utils.rs

@@ -62,6 +62,7 @@ pub struct BenchmarkResults {
 	pub repeat_reads: u32,
 	pub writes: u32,
 	pub repeat_writes: u32,
+	pub proof_size: u32,
 }
 
 /// Configuration used to setup and run runtime benchmarks.
@@ -162,6 +163,11 @@ pub trait Benchmarking {
 		whitelist.retain(|x| x.key != remove);
 		self.set_whitelist(whitelist);
 	}
+
+	/// Get current estimated proof size.
+	fn proof_size(&self) -> Option<u32> {
+		self.proof_size()
+	}
 }
 
 /// The pallet benchmarking trait.

primitives/externalities/src/lib.rs

@@ -281,6 +281,16 @@ pub trait Externalities: ExtensionStore {
 	///
 	/// Adds new storage keys to the DB tracking whitelist.
 	fn set_whitelist(&mut self, new: Vec<TrackedStorageKey>);
+
+	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+	/// Benchmarking related functionality and shouldn't be used anywhere else!
+	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+	///
+	/// Returns estimated proof size for the state queries so far.
+	/// Proof is reset on commit and wipe.
+	fn proof_size(&self) -> Option<u32> {
+		None
+	}
 }
 
 /// Extension for the [`Externalities`] trait.

primitives/state-machine/src/backend.rs

@@ -245,6 +245,11 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 
 	/// Update the whitelist for tracking db reads/writes
 	fn set_whitelist(&self, _: Vec<TrackedStorageKey>) {}
+
+	/// Estimate proof size
+	fn proof_size(&self) -> Option<u32> {
+		unimplemented!()
+	}
 }
 
 impl<'a, T: Backend<H>, H: Hasher> Backend<H> for &'a T {

primitives/state-machine/src/ext.rs

@@ -712,6 +712,10 @@ where
 	fn set_whitelist(&mut self, new: Vec<TrackedStorageKey>) {
 		self.backend.set_whitelist(new)
 	}
+
+	fn proof_size(&self) -> Option<u32> {
+		self.backend.proof_size()
+	}
 }
 
 /// Implement `Encode` by forwarding the stored raw vec.

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

@@ -63,7 +63,7 @@ impl BenchmarkCmd {
 		let genesis_storage = spec.build_storage()?;
 		let mut changes = Default::default();
 		let cache_size = Some(self.database_cache_size as usize);
-		let state = BenchmarkingState::<BB>::new(genesis_storage, cache_size)?;
+		let state = BenchmarkingState::<BB>::new(genesis_storage, cache_size, self.record_proof)?;
 		let executor = NativeExecutor::<ExecDispatch>::new(
 			wasm_method,
 			self.heap_pages,
@@ -126,19 +126,20 @@ impl BenchmarkCmd {
 						// Print the table header
 						batch.results[0].components.iter().for_each(|param| print!("{:?},", param.0));
 
-						print!("extrinsic_time,storage_root_time,reads,repeat_reads,writes,repeat_writes\n");
+						print!("extrinsic_time_ns,storage_root_time_ns,reads,repeat_reads,writes,repeat_writes,proof_size_bytes\n");
 						// Print the values
 						batch.results.iter().for_each(|result| {
 							let parameters = &result.components;
 							parameters.iter().for_each(|param| print!("{:?},", param.1));
 							// Print extrinsic time and storage root time
-							print!("{:?},{:?},{:?},{:?},{:?},{:?}\n",
+							print!("{:?},{:?},{:?},{:?},{:?},{:?},{:?}\n",
 								result.extrinsic_time,
 								result.storage_root_time,
 								result.reads,
 								result.repeat_reads,
 								result.writes,
 								result.repeat_writes,
+								result.proof_size,
 							);
 						});
 

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

@@ -97,6 +97,10 @@ pub struct BenchmarkCmd {
 	#[structopt(long)]
 	pub extra: bool,
 
+	/// Estimate PoV size.
+	#[structopt(long)]
+	pub record_proof: bool,
+
 	#[allow(missing_docs)]
 	#[structopt(flatten)]
 	pub shared_params: sc_cli::SharedParams,

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

@@ -421,6 +421,7 @@ mod test {
 					repeat_reads: 0,
 					writes: (base + slope * i).into(),
 					repeat_writes: 0,
+					proof_size: 0,
 				}
 			)
 		}