Benchmark storage access on block validation

Cargo.toml

@@ -554,6 +554,7 @@ members = [
 	"substrate/utils/frame/rpc/state-trie-migration-rpc",
 	"substrate/utils/frame/rpc/support",
 	"substrate/utils/frame/rpc/system",
+	"substrate/utils/frame/storage-access-test-runtime",
 	"substrate/utils/prometheus",
 	"substrate/utils/substrate-bip39",
 	"substrate/utils/wasm-builder",
@@ -806,6 +807,7 @@ frame-election-provider-support = { path = "substrate/frame/election-provider-su
 frame-executive = { path = "substrate/frame/executive", default-features = false }
 frame-metadata = { version = "23.0.0", default-features = false }
 frame-metadata-hash-extension = { path = "substrate/frame/metadata-hash-extension", default-features = false }
+frame-storage-access-test-runtime = { path = "substrate/utils/frame/storage-access-test-runtime", default-features = false }
 frame-support = { path = "substrate/frame/support", default-features = false }
 frame-support-procedural = { path = "substrate/frame/support/procedural", default-features = false }
 frame-support-procedural-tools = { path = "substrate/frame/support/procedural/tools", default-features = false }

cumulus/pallets/parachain-system/src/validate_block/mod.rs

@@ -24,11 +24,11 @@ mod tests;
 
 #[cfg(not(feature = "std"))]
 #[doc(hidden)]
-mod trie_cache;
+pub mod trie_cache;
 
 #[cfg(any(test, not(feature = "std")))]
 #[doc(hidden)]
-mod trie_recorder;
+pub mod trie_recorder;
 
 #[cfg(not(feature = "std"))]
 #[doc(hidden)]

cumulus/pallets/parachain-system/src/validate_block/trie_cache.rs

@@ -27,7 +27,7 @@ use trie_db::{node::NodeOwned, Hasher};
 /// Special purpose trie cache implementation that is able to cache an unlimited number
 /// of values. To be used in `validate_block` to serve values and nodes that
 /// have already been loaded and decoded from the storage proof.
-pub(crate) struct TrieCache<'a, H: Hasher> {
+pub struct TrieCache<'a, H: Hasher> {
 	node_cache: RefMut<'a, BTreeMap<H::Out, NodeOwned<H::Out>>>,
 	value_cache: Option<RefMut<'a, BTreeMap<Box<[u8]>, trie_db::CachedValue<H::Out>>>>,
 }
@@ -65,7 +65,7 @@ impl<'a, H: Hasher> trie_db::TrieCache<NodeCodec<H>> for TrieCache<'a, H> {
 }
 
 /// Provider of [`TrieCache`] instances.
-pub(crate) struct CacheProvider<H: Hasher> {
+pub struct CacheProvider<H: Hasher> {
 	node_cache: RefCell<BTreeMap<H::Out, NodeOwned<H::Out>>>,
 	/// Cache: `storage_root` => `storage_key` => `value`.
 	///

cumulus/pallets/parachain-system/src/validate_block/trie_recorder.rs

@@ -34,7 +34,7 @@ use trie_db::{Hasher, RecordedForKey, TrieAccess};
 ///
 /// The internal size counting logic should align
 /// with ['sp_trie::recorder::Recorder'].
-pub(crate) struct SizeOnlyRecorder<'a, H: Hasher> {
+pub struct SizeOnlyRecorder<'a, H: Hasher> {
 	seen_nodes: RefMut<'a, BTreeSet<H::Out>>,
 	encoded_size: RefMut<'a, usize>,
 	recorded_keys: RefMut<'a, BTreeMap<Rc<[u8]>, RecordedForKey>>,
@@ -90,7 +90,7 @@ impl<'a, H: trie_db::Hasher> trie_db::TrieRecorder<H::Out> for SizeOnlyRecorder<
 }
 
 #[derive(Clone)]
-pub(crate) struct SizeOnlyRecorderProvider<H: Hasher> {
+pub struct SizeOnlyRecorderProvider<H: Hasher> {
 	seen_nodes: Rc<RefCell<BTreeSet<H::Out>>>,
 	encoded_size: Rc<RefCell<usize>>,
 	recorded_keys: Rc<RefCell<BTreeMap<Rc<[u8]>, RecordedForKey>>>,

prdoc/pr_8069.prdoc

@@ -0,0 +1,17 @@
+# 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: Benchmark storage access on block validation
+
+doc:
+  - audience: [Runtime Dev, Node Dev]
+    description: |
+      Adds checking storage weights on block validation for both read and write benchmarks.
+
+crates:
+  - name: cumulus-pallet-parachain-system
+    bump: minor
+  - name: frame-benchmarking-cli
+    bump: minor
+  - name: frame-storage-access-test-runtime
+    bump: major

substrate/utils/frame/benchmarking-cli/Cargo.toml

@@ -25,6 +25,7 @@ comfy-table = { workspace = true }
 cumulus-client-parachain-inherent = { workspace = true, default-features = true }
 cumulus-primitives-proof-size-hostfunction = { workspace = true, default-features = true }
 frame-benchmarking = { workspace = true, default-features = true }
+frame-storage-access-test-runtime = { workspace = true, default-features = true }
 frame-support = { workspace = true, default-features = true }
 frame-system = { workspace = true, default-features = true }
 gethostname = { workspace = true }
@@ -42,6 +43,8 @@ sc-cli = { workspace = true, default-features = false }
 sc-client-api = { workspace = true, default-features = true }
 sc-client-db = { workspace = true, default-features = false }
 sc-executor = { workspace = true, default-features = true }
+sc-executor-common = { workspace = true }
+sc-executor-wasmtime = { workspace = true }
 sc-runtime-utilities = { workspace = true, default-features = true }
 sc-service = { workspace = true, default-features = false }
 sc-sysinfo = { workspace = true, default-features = true }
@@ -79,6 +82,7 @@ westend-runtime = { workspace = true, default-features = true }
 default = []
 runtime-benchmarks = [
 	"frame-benchmarking/runtime-benchmarks",
+	"frame-storage-access-test-runtime/runtime-benchmarks",
 	"frame-support/runtime-benchmarks",
 	"frame-system/runtime-benchmarks",
 	"polkadot-parachain-primitives/runtime-benchmarks",

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

@@ -26,7 +26,7 @@ use sp_runtime::traits::{Block as BlockT, HashingFor};
 use sp_state_machine::Storage;
 use sp_storage::{ChildInfo, ChildType, PrefixedStorageKey, StateVersion};
 
-use clap::{Args, Parser};
+use clap::{Args, Parser, ValueEnum};
 use log::info;
 use rand::prelude::*;
 use serde::Serialize;
@@ -36,6 +36,16 @@ use std::{fmt::Debug, path::PathBuf, sync::Arc};
 use super::template::TemplateData;
 use crate::shared::{new_rng, HostInfoParams, WeightParams};
 
+/// The mode in which to run the storage benchmark.
+#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Serialize, ValueEnum)]
+pub enum StorageBenchmarkMode {
+	/// Run the benchmark for block import.
+	#[default]
+	ImportBlock,
+	/// Run the benchmark for block validation.
+	ValidateBlock,
+}
+
 /// Benchmark the storage speed of a chain snapshot.
 #[derive(Debug, Parser)]
 pub struct StorageCmd {
@@ -129,13 +139,36 @@ pub struct StorageParams {
 	#[arg(long, default_value = "false")]
 	pub disable_pov_recorder: bool,
 
-	/// The batch size for the write benchmark.
+	/// The batch size for the read/write benchmark.
 	///
 	/// Since the write size needs to also include the cost of computing the storage root, which is
 	/// done once at the end of the block, the batch size is used to simulate multiple writes in a
 	/// block.
 	#[arg(long, default_value_t = 100_000)]
 	pub batch_size: usize,
+
+	/// The mode in which to run the storage benchmark.
+	///
+	/// PoV recorder must be activated to provide a storage proof for block validation at runtime.
+	#[arg(long, value_enum, default_value_t = StorageBenchmarkMode::ImportBlock)]
+	pub mode: StorageBenchmarkMode,
+
+	/// Number of rounds to execute block validation during the benchmark.
+	///
+	/// We need to run the benchmark several times to avoid fluctuations during runtime setup.
+	/// This is only used when `mode` is `validate-block`.
+	#[arg(long, default_value_t = 20)]
+	pub validate_block_rounds: u32,
+}
+
+impl StorageParams {
+	pub fn is_import_block_mode(&self) -> bool {
+		matches!(self.mode, StorageBenchmarkMode::ImportBlock)
+	}
+
+	pub fn is_validate_block_mode(&self) -> bool {
+		matches!(self.mode, StorageBenchmarkMode::ValidateBlock)
+	}
 }
 
 impl StorageCmd {

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

@@ -21,3 +21,37 @@ pub mod template;
 pub mod write;
 
 pub use cmd::StorageCmd;
+
+/// Empirically, the maximum batch size for block validation should be no more than 10,000.
+/// Bigger sizes may cause problems with runtime memory allocation.
+pub(crate) const MAX_BATCH_SIZE_FOR_BLOCK_VALIDATION: usize = 10_000;
+
+pub(crate) fn get_wasm_module() -> Box<dyn sc_executor_common::wasm_runtime::WasmModule> {
+	let blob = sc_executor_common::runtime_blob::RuntimeBlob::uncompress_if_needed(
+		frame_storage_access_test_runtime::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the benchmark!"),
+	)
+	.expect("Failed to create runtime blob");
+	let config = sc_executor_wasmtime::Config {
+		allow_missing_func_imports: true,
+		cache_path: None,
+		semantics: sc_executor_wasmtime::Semantics {
+			heap_alloc_strategy: sc_executor_common::wasm_runtime::HeapAllocStrategy::Dynamic {
+				maximum_pages: Some(4096),
+			},
+			instantiation_strategy: sc_executor::WasmtimeInstantiationStrategy::PoolingCopyOnWrite,
+			deterministic_stack_limit: None,
+			canonicalize_nans: false,
+			parallel_compilation: false,
+			wasm_multi_value: false,
+			wasm_bulk_memory: false,
+			wasm_reference_types: false,
+			wasm_simd: false,
+		},
+	};
+
+	Box::new(
+		sc_executor_wasmtime::create_runtime::<sp_io::SubstrateHostFunctions>(blob, config)
+			.expect("Unable to create wasm module."),
+	)
+}