Add ProofSizeExt to benchmarks (#5257)

I propose to have `ProofSizeExt` available during benchmarking so we can
improve the accuracy for extensions using it.

Another thing we could do is to also enable recording for the timing
benchmark here:

https://github.com/paritytech/polkadot-sdk/blob/035211d707d0a74a2a768fd658160721f09d5b44/substrate/utils/frame/benchmarking-cli/src/pallet/command.rs#L232

Parachains will need to have recording enabled during import for
reclaim, so we could enable it here and provide a flag
`--disable-proof-recording` for scenarios where one does not want it.
Happy to hear opinions about this.

prdoc/pr_5257.prdoc

@@ -0,0 +1,20 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Enable proof-recording in benchmarking
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      We now enable proof recording in the timing benchmarks. This affects the standalone `frame-omni-bencher` as well
+      as the integrated benchmarking commands of the node. For parachains on recent versions of polkadot-sdk this is the
+      correct default setting, since PoV-reclaim requires proof recording to be enabled on block import.
+      This comes with a slight increase in timing weight due to the additional overhead. Relay- or solo-chains
+      which do not need proof recording can restore the old behaviour by passing `--disable-proof-recording` to the
+      benchmarking command.
+
+      In addition, the `ProofSizeExt` extension is available during benchmarking.
+
+crates:
+  - name: frame-benchmarking-cli
+    bump: minor

substrate/client/db/src/bench.rs

@@ -179,6 +179,11 @@ impl<Hasher: Hash> BenchmarkingState<Hasher> {
 		Ok(state)
 	}
 
+	/// Get the proof recorder for this state
+	pub fn recorder(&self) -> Option<sp_trie::recorder::Recorder<Hasher>> {
+		self.proof_recorder.clone()
+	}
+
 	fn reopen(&self) -> Result<(), String> {
 		*self.state.borrow_mut() = None;
 		let db = match self.db.take() {

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

@@ -37,12 +37,14 @@ use sp_core::{
 		OffchainDbExt, OffchainWorkerExt, TransactionPoolExt,
 	},
 	traits::{CallContext, CodeExecutor, ReadRuntimeVersionExt, WrappedRuntimeCode},
+	Hasher,
 };
 use sp_externalities::Extensions;
 use sp_genesis_builder::{PresetId, Result as GenesisBuildResult};
 use sp_keystore::{testing::MemoryKeystore, KeystoreExt};
 use sp_runtime::traits::Hash;
 use sp_state_machine::{OverlayedChanges, StateMachine};
+use sp_trie::{proof_size_extension::ProofSizeExt, recorder::Recorder};
 use sp_wasm_interface::HostFunctions;
 use std::{
 	borrow::Cow,
@@ -225,11 +227,12 @@ impl PalletCmd {
 			// Enable storage tracking
 			true,
 		)?;
+
 		let state_without_tracking = BenchmarkingState::<Hasher>::new(
 			genesis_storage,
 			cache_size,
-			// Do not record proof size
-			false,
+			// Proof recording depends on CLI settings
+			!self.disable_proof_recording,
 			// Do not enable storage tracking
 			false,
 		)?;
@@ -263,7 +266,7 @@ impl PalletCmd {
 					&executor,
 					"Benchmark_benchmark_metadata",
 					&(self.extra).encode(),
-					&mut Self::build_extensions(executor.clone()),
+					&mut Self::build_extensions(executor.clone(), state.recorder()),
 					&runtime_code,
 					CallContext::Offchain,
 				),
@@ -365,7 +368,7 @@ impl PalletCmd {
 								1,    // no need to do internal repeats
 							)
 								.encode(),
-							&mut Self::build_extensions(executor.clone()),
+							&mut Self::build_extensions(executor.clone(), state.recorder()),
 							&runtime_code,
 							CallContext::Offchain,
 						),
@@ -406,7 +409,7 @@ impl PalletCmd {
 								self.repeat,
 							)
 								.encode(),
-							&mut Self::build_extensions(executor.clone()),
+							&mut Self::build_extensions(executor.clone(), state.recorder()),
 							&runtime_code,
 							CallContext::Offchain,
 						),
@@ -449,7 +452,7 @@ impl PalletCmd {
 								self.repeat,
 							)
 								.encode(),
-							&mut Self::build_extensions(executor.clone()),
+							&mut Self::build_extensions(executor.clone(), state.recorder()),
 							&runtime_code,
 							CallContext::Offchain,
 						),
@@ -626,7 +629,7 @@ impl PalletCmd {
 				&executor,
 				"GenesisBuilder_get_preset",
 				&None::<PresetId>.encode(), // Use the default preset
-				&mut Self::build_extensions(executor.clone()),
+				&mut Self::build_extensions(executor.clone(), state.recorder()),
 				&runtime_code,
 				CallContext::Offchain,
 			),
@@ -647,7 +650,7 @@ impl PalletCmd {
 				&executor,
 				"GenesisBuilder_get_preset",
 				&Some::<PresetId>(GENESIS_PRESET.into()).encode(), // Use the default preset
-				&mut Self::build_extensions(executor.clone()),
+				&mut Self::build_extensions(executor.clone(), state.recorder()),
 				&runtime_code,
 				CallContext::Offchain,
 			),
@@ -707,7 +710,10 @@ impl PalletCmd {
 	}
 
 	/// Build the extension that are available for pallet benchmarks.
-	fn build_extensions<E: CodeExecutor>(exe: E) -> Extensions {
+	fn build_extensions<E: CodeExecutor, H: Hasher + 'static>(
+		exe: E,
+		maybe_recorder: Option<Recorder<H>>,
+	) -> Extensions {
 		let mut extensions = Extensions::default();
 		let (offchain, _) = TestOffchainExt::new();
 		let (pool, _) = TestTransactionPoolExt::new();
@@ -717,6 +723,9 @@ impl PalletCmd {
 		extensions.register(OffchainDbExt::new(offchain));
 		extensions.register(TransactionPoolExt::new(pool));
 		extensions.register(ReadRuntimeVersionExt::new(exe));
+		if let Some(recorder) = maybe_recorder {
+			extensions.register(ProofSizeExt::new(recorder));
+		}
 		extensions
 	}
 

substrate/utils/frame/benchmarking-cli/src/pallet/mod.rs

@@ -243,4 +243,13 @@ pub struct PalletCmd {
 	/// use-cases, this option reduces the noise.
 	#[arg(long)]
 	quiet: bool,
+
+	/// Do not enable proof recording during time benchmarking.
+	///
+	/// By default, proof recording is enabled during benchmark execution. This can slightly
+	/// inflate the resulting time weights. For parachains using PoV-reclaim, this is typically the
+	/// correct setting. Chains that ignore the proof size dimension of weight (e.g. relay chain,
+	/// solo-chains) can disable proof recording to get more accurate results.
+	#[arg(long)]
+	disable_proof_recording: bool,
 }