PoV Reclaim Runtime Side

Cargo.toml

@@ -127,6 +127,7 @@ members = [
 	"cumulus/primitives/core",
 	"cumulus/primitives/parachain-inherent",
 	"cumulus/primitives/proof-size-hostfunction",
+	"cumulus/primitives/storage-weight-reclaim",
 	"cumulus/primitives/timestamp",
 	"cumulus/primitives/utility",
 	"cumulus/test/client",

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

@@ -16,7 +16,7 @@
 
 //! The actual implementation of the validate block functionality.
 
-use super::{trie_cache, MemoryOptimizedValidationParams};
+use super::{trie_cache, trie_recorder, MemoryOptimizedValidationParams};
 use cumulus_primitives_core::{
 	relay_chain::Hash as RHash, ParachainBlockData, PersistedValidationData,
 };
@@ -34,12 +34,14 @@ use sp_externalities::{set_and_run_with_externalities, Externalities};
 use sp_io::KillStorageResult;
 use sp_runtime::traits::{Block as BlockT, Extrinsic, HashingFor, Header as HeaderT};
 use sp_std::prelude::*;
-use sp_trie::MemoryDB;
+use sp_trie::{MemoryDB, ProofSizeProvider, TrieRecorderProvider};
+use trie_recorder::SizeOnlyRecorderProvider;
 
 type TrieBackend<B> = sp_state_machine::TrieBackend<
 	MemoryDB<HashingFor<B>>,
 	HashingFor<B>,
 	trie_cache::CacheProvider<HashingFor<B>>,
+	SizeOnlyRecorderProvider<HashingFor<B>>,
 >;
 
 type Ext<'a, B> = sp_state_machine::Ext<'a, HashingFor<B>, TrieBackend<B>>;
@@ -48,6 +50,9 @@ fn with_externalities<F: FnOnce(&mut dyn Externalities) -> R, R>(f: F) -> R {
 	sp_externalities::with_externalities(f).expect("Environmental externalities not set.")
 }
 
+/// Recorder instance to be used during this validate_block call.
+environmental::environmental!(recorder: trait ProofSizeProvider);
+
 /// Validate the given parachain block.
 ///
 /// This function is doing roughly the following:
@@ -120,6 +125,7 @@ where
 
 	sp_std::mem::drop(storage_proof);
 
+	let mut recorder = SizeOnlyRecorderProvider::new();
 	let cache_provider = trie_cache::CacheProvider::new();
 	// We use the storage root of the `parent_head` to ensure that it is the correct root.
 	// This is already being done above while creating the in-memory db, but let's be paranoid!!
@@ -128,6 +134,7 @@ where
 		*parent_header.state_root(),
 		cache_provider,
 	)
+	.with_recorder(recorder.clone())
 	.build();
 
 	let _guard = (
@@ -167,6 +174,8 @@ where
 			.replace_implementation(host_default_child_storage_next_key),
 		sp_io::offchain_index::host_set.replace_implementation(host_offchain_index_set),
 		sp_io::offchain_index::host_clear.replace_implementation(host_offchain_index_clear),
+		cumulus_primitives_proof_size_hostfunction::storage_proof_size::host_storage_proof_size
+			.replace_implementation(host_storage_proof_size),
 	);
 
 	run_with_externalities::<B, _, _>(&backend, || {
@@ -190,7 +199,7 @@ where
 		}
 	});
 
-	run_with_externalities::<B, _, _>(&backend, || {
+	run_with_externalities_and_recorder::<B, _, _>(&backend, &mut recorder, || {
 		let head_data = HeadData(block.header().encode());
 
 		E::execute_block(block);
@@ -265,6 +274,19 @@ fn validate_validation_data(
 	);
 }
 
+/// Run the given closure with the externalities and recorder set.
+fn run_with_externalities_and_recorder<B: BlockT, R, F: FnOnce() -> R>(
+	backend: &TrieBackend<B>,
+	recorder: &mut SizeOnlyRecorderProvider<HashingFor<B>>,
+	execute: F,
+) -> R {
+	let mut overlay = sp_state_machine::OverlayedChanges::default();
+	let mut ext = Ext::<B>::new(&mut overlay, backend);
+	recorder.reset();
+
+	recorder::using(recorder, || set_and_run_with_externalities(&mut ext, || execute()))
+}
+
 /// Run the given closure with the externalities set.
 fn run_with_externalities<B: BlockT, R, F: FnOnce() -> R>(
 	backend: &TrieBackend<B>,
@@ -305,6 +327,10 @@ fn host_storage_clear(key: &[u8]) {
 	with_externalities(|ext| ext.place_storage(key.to_vec(), None))
 }
 
+fn host_storage_proof_size() -> u64 {
+	recorder::with(|rec| rec.estimate_encoded_size()).expect("Recorder is always set; qed") as _
+}
+
 fn host_storage_root(version: StateVersion) -> Vec<u8> {
 	with_externalities(|ext| ext.storage_root(version))
 }

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

@@ -97,13 +97,21 @@ pub(crate) struct SizeOnlyRecorderProvider<H: Hasher> {
 }
 
 impl<H: Hasher> SizeOnlyRecorderProvider<H> {
+	/// Create a new instance of [`SizeOnlyRecorderProvider`]
 	pub fn new() -> Self {
 		Self {
 			seen_nodes: Default::default(),
 			encoded_size: Default::default(),
 			recorded_keys: Default::default(),
 		}
 	}
+
+	/// Reset the internal state.
+	pub fn reset(&self) {
+		self.seen_nodes.borrow_mut().clear();
+		*self.encoded_size.borrow_mut() = 0;
+		self.recorded_keys.borrow_mut().clear();
+	}
 }
 
 impl<H: trie_db::Hasher> sp_trie::TrieRecorderProvider<H> for SizeOnlyRecorderProvider<H> {
@@ -281,6 +289,9 @@ mod tests {
 				reference_recorder.estimate_encoded_size(),
 				recorder_for_test.estimate_encoded_size()
 			);
+
+			recorder_for_test.reset();
+			assert_eq!(recorder_for_test.estimate_encoded_size(), 0)
 		}
 	}
 }

cumulus/parachain-template/node/src/command.rs

@@ -1,5 +1,6 @@
 use std::net::SocketAddr;
 
+use cumulus_client_service::storage_proof_size::HostFunctions as ReclaimHostFunctions;
 use cumulus_primitives_core::ParaId;
 use frame_benchmarking_cli::{BenchmarkCmd, SUBSTRATE_REFERENCE_HARDWARE};
 use log::info;
@@ -183,7 +184,7 @@ pub fn run() -> Result<()> {
 			match cmd {
 				BenchmarkCmd::Pallet(cmd) =>
 					if cfg!(feature = "runtime-benchmarks") {
-						runner.sync_run(|config| cmd.run::<sp_runtime::traits::HashingFor<Block>, ()>(config))
+						runner.sync_run(|config| cmd.run::<sp_runtime::traits::HashingFor<Block>, ReclaimHostFunctions>(config))
 					} else {
 						Err("Benchmarking wasn't enabled when building the node. \
 					You can enable it with `--features runtime-benchmarks`."

cumulus/parachain-template/node/src/service.rs

@@ -40,7 +40,10 @@ use substrate_prometheus_endpoint::Registry;
 pub struct ParachainNativeExecutor;
 
 impl sc_executor::NativeExecutionDispatch for ParachainNativeExecutor {
-	type ExtendHostFunctions = frame_benchmarking::benchmarking::HostFunctions;
+	type ExtendHostFunctions = (
+		cumulus_client_service::storage_proof_size::HostFunctions,
+		frame_benchmarking::benchmarking::HostFunctions,
+	);
 
 	fn dispatch(method: &str, data: &[u8]) -> Option<Vec<u8>> {
 		parachain_template_runtime::api::dispatch(method, data)
@@ -100,10 +103,11 @@ pub fn new_partial(config: &Configuration) -> Result<Service, sc_service::Error>
 	let executor = ParachainExecutor::new_with_wasm_executor(wasm);
 
 	let (client, backend, keystore_container, task_manager) =
-		sc_service::new_full_parts::<Block, RuntimeApi, _>(
+		sc_service::new_full_parts_record_import::<Block, RuntimeApi, _>(
 			config,
 			telemetry.as_ref().map(|(_, telemetry)| telemetry.handle()),
 			executor,
+			true,
 		)?;
 	let client = Arc::new(client);
 

cumulus/parachain-template/runtime/Cargo.toml

@@ -73,6 +73,7 @@ cumulus-pallet-xcm = { path = "../../pallets/xcm", default-features = false }
 cumulus-pallet-xcmp-queue = { path = "../../pallets/xcmp-queue", default-features = false }
 cumulus-primitives-core = { path = "../../primitives/core", default-features = false }
 cumulus-primitives-utility = { path = "../../primitives/utility", default-features = false }
+cumulus-primitives-storage-weight-reclaim = { path = "../../primitives/storage-weight-reclaim", default-features = false }
 pallet-collator-selection = { path = "../../pallets/collator-selection", default-features = false }
 parachains-common = { path = "../../parachains/common", default-features = false }
 parachain-info = { package = "staging-parachain-info", path = "../../parachains/pallets/parachain-info", default-features = false }
@@ -87,6 +88,7 @@ std = [
 	"cumulus-pallet-xcm/std",
 	"cumulus-pallet-xcmp-queue/std",
 	"cumulus-primitives-core/std",
+	"cumulus-primitives-storage-weight-reclaim/std",
 	"cumulus-primitives-utility/std",
 	"frame-benchmarking?/std",
 	"frame-executive/std",

cumulus/parachain-template/runtime/src/lib.rs

@@ -107,6 +107,7 @@ pub type SignedExtra = (
 	frame_system::CheckNonce<Runtime>,
 	frame_system::CheckWeight<Runtime>,
 	pallet_transaction_payment::ChargeTransactionPayment<Runtime>,
+	cumulus_primitives_storage_weight_reclaim::StorageWeightReclaim<Runtime>,
 );
 
 /// Unchecked extrinsic type as expected by this runtime.

cumulus/polkadot-parachain/src/command.rs

@@ -23,6 +23,7 @@ use crate::{
 	},
 	service::{new_partial, Block},
 };
+use cumulus_client_service::storage_proof_size::HostFunctions as ReclaimHostFunctions;
 use cumulus_primitives_core::ParaId;
 use frame_benchmarking_cli::{BenchmarkCmd, SUBSTRATE_REFERENCE_HARDWARE};
 use log::info;
@@ -584,7 +585,7 @@ pub fn run() -> Result<()> {
 			match cmd {
 				BenchmarkCmd::Pallet(cmd) =>
 					if cfg!(feature = "runtime-benchmarks") {
-						runner.sync_run(|config| cmd.run::<sp_runtime::traits::HashingFor<Block>, ()>(config))
+						runner.sync_run(|config| cmd.run::<sp_runtime::traits::HashingFor<Block>, ReclaimHostFunctions>(config))
 					} else {
 						Err("Benchmarking wasn't enabled when building the node. \
 				You can enable it with `--features runtime-benchmarks`."

cumulus/polkadot-parachain/src/service.rs

@@ -68,11 +68,15 @@ use substrate_prometheus_endpoint::Registry;
 use polkadot_primitives::CollatorPair;
 
 #[cfg(not(feature = "runtime-benchmarks"))]
-type HostFunctions = sp_io::SubstrateHostFunctions;
+type HostFunctions =
+	(sp_io::SubstrateHostFunctions, cumulus_client_service::storage_proof_size::HostFunctions);
 
 #[cfg(feature = "runtime-benchmarks")]
-type HostFunctions =
-	(sp_io::SubstrateHostFunctions, frame_benchmarking::benchmarking::HostFunctions);
+type HostFunctions = (
+	sp_io::SubstrateHostFunctions,
+	cumulus_client_service::storage_proof_size::HostFunctions,
+	frame_benchmarking::benchmarking::HostFunctions,
+);
 
 type ParachainClient<RuntimeApi> = TFullClient<Block, RuntimeApi, WasmExecutor<HostFunctions>>;
 
@@ -274,10 +278,11 @@ where
 		.build();
 
 	let (client, backend, keystore_container, task_manager) =
-		sc_service::new_full_parts::<Block, RuntimeApi, _>(
+		sc_service::new_full_parts_record_import::<Block, RuntimeApi, _>(
 			config,
 			telemetry.as_ref().map(|(_, telemetry)| telemetry.handle()),
 			executor,
+			true,
 		)?;
 	let client = Arc::new(client);
 

cumulus/primitives/proof-size-hostfunction/src/lib.rs

@@ -35,7 +35,8 @@ pub const PROOF_RECORDING_DISABLED: u64 = u64::MAX;
 pub trait StorageProofSize {
 	/// Returns the current storage proof size.
 	fn storage_proof_size(&mut self) -> u64 {
-		self.extension::<ProofSizeExt>().map_or(u64::MAX, |e| e.storage_proof_size())
+		self.extension::<ProofSizeExt>()
+			.map_or(PROOF_RECORDING_DISABLED, |e| e.storage_proof_size())
 	}
 }
 

cumulus/primitives/storage-weight-reclaim/Cargo.toml

@@ -0,0 +1,33 @@
+[package]
+name = "cumulus-primitives-storage-weight-reclaim"
+version = "0.2.0"
+authors.workspace = true
+edition.workspace = true
+description = "Utilities to reclaim storage weight."
+license = "Apache-2.0"
+
+[lints]
+workspace = true
+
+[dependencies]
+codec = { package = "parity-scale-codec", version = "3.0.0", default-features = false, features = ["derive"] }
+log = { workspace = true }
+scale-info = { version = "2.10.0", default-features = false, features = ["derive"] }
+
+frame-support = { path = "../../../substrate/frame/support", default-features = false }
+frame-system = { path = "../../../substrate/frame/system", default-features = false }
+
+sp-runtime = { path = "../../../substrate/primitives/runtime", default-features = false }
+sp-std = { path = "../../../substrate/primitives/std", default-features = false }
+
+cumulus-primitives-core = { path = "../../primitives/core", default-features = false }
+cumulus-primitives-proof-size-hostfunction = { path = "../../primitives/proof-size-hostfunction", default-features = false }
+docify = "0.2.7"
+
+[dev-dependencies]
+sp-trie = { path = "../../../substrate/primitives/trie", default-features = false }
+sp-io = { path = "../../../substrate/primitives/io", default-features = false }
+
+[features]
+default = ["std"]
+std = ["codec/std", "cumulus-primitives-core/std", "cumulus-primitives-proof-size-hostfunction/std", "frame-support/std", "frame-system/std", "log/std", "scale-info/std", "sp-io/std", "sp-runtime/std", "sp-std/std", "sp-trie/std"]