PoV Reclaim Runtime Side (#3002)

# Runtime side for PoV Reclaim

## Implementation Overview
- Hostfunction to fetch the storage proof size has been added to the
PVF. It uses the size tracking recorder that was introduced in my
previous PR.
- Mechanisms to use the reclaim HostFunction have been introduced.
- 1. A SignedExtension that checks the node-reported proof size before
and after application of an extrinsic. Then it reclaims the difference.
- 2. A manual helper to make reclaiming easier when manual interaction
is required, for example in `on_idle` or other hooks.
- In order to utilize the manual reclaiming, I modified `WeightMeter` to
support the reduction of consumed weight, at least for storage proof
size.

## How to use
To enable the general functionality for a parachain:
1. Add the SignedExtension to your parachain runtime. 
2. Provide the HostFunction to the node
3. Enable proof recording during block import

## TODO
- [x] PRDoc

---------

Co-authored-by: Dmitry Markin <dmitry@markin.tech>
Co-authored-by: Davide Galassi <davxy@datawok.net>
Co-authored-by: Bastian Köcher <git@kchr.de>

.gitlab/pipeline/check.yml

@@ -108,8 +108,10 @@ check-toml-format:
       export RUST_LOG=remote-ext=debug,runtime=debug
 
       echo "---------- Downloading try-runtime CLI ----------"
-      curl -sL https://github.com/paritytech/try-runtime-cli/releases/download/v0.5.0/try-runtime-x86_64-unknown-linux-musl -o try-runtime
+      curl -sL https://github.com/paritytech/try-runtime-cli/releases/download/v0.5.4/try-runtime-x86_64-unknown-linux-musl -o try-runtime
       chmod +x ./try-runtime
+      echo "Using try-runtime-cli version:"
+      ./try-runtime --version
 
       echo "---------- Building ${PACKAGE} runtime ----------"
       time cargo build --release --locked -p "$PACKAGE" --features try-runtime

Cargo.lock

@@ -4047,6 +4047,25 @@ dependencies = [
  "sp-trie",
 ]
 
+[[package]]
+name = "cumulus-primitives-storage-weight-reclaim"
+version = "1.0.0"
+dependencies = [
+ "cumulus-primitives-core",
+ "cumulus-primitives-proof-size-hostfunction",
+ "cumulus-test-runtime",
+ "docify",
+ "frame-support",
+ "frame-system",
+ "log",
+ "parity-scale-codec",
+ "scale-info",
+ "sp-io",
+ "sp-runtime",
+ "sp-std 14.0.0",
+ "sp-trie",
+]
+
 [[package]]
 name = "cumulus-primitives-timestamp"
 version = "0.7.0"
@@ -4209,6 +4228,7 @@ dependencies = [
  "cumulus-primitives-core",
  "cumulus-primitives-parachain-inherent",
  "cumulus-primitives-proof-size-hostfunction",
+ "cumulus-primitives-storage-weight-reclaim",
  "cumulus-test-relay-sproof-builder",
  "cumulus-test-runtime",
  "cumulus-test-service",
@@ -4253,6 +4273,7 @@ version = "0.1.0"
 dependencies = [
  "cumulus-pallet-parachain-system",
  "cumulus-primitives-core",
+ "cumulus-primitives-storage-weight-reclaim",
  "frame-executive",
  "frame-support",
  "frame-system",
@@ -4295,6 +4316,7 @@ dependencies = [
  "cumulus-client-service",
  "cumulus-pallet-parachain-system",
  "cumulus-primitives-core",
+ "cumulus-primitives-storage-weight-reclaim",
  "cumulus-relay-chain-inprocess-interface",
  "cumulus-relay-chain-interface",
  "cumulus-relay-chain-minimal-node",
@@ -11340,6 +11362,7 @@ dependencies = [
  "cumulus-pallet-xcm",
  "cumulus-pallet-xcmp-queue",
  "cumulus-primitives-core",
+ "cumulus-primitives-storage-weight-reclaim",
  "cumulus-primitives-utility",
  "frame-benchmarking",
  "frame-executive",

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};
+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,9 +174,11 @@ 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, || {
+	run_with_externalities_and_recorder::<B, _, _>(&backend, &mut recorder, || {
 		let relay_chain_proof = crate::RelayChainStateProof::new(
 			PSC::SelfParaId::get(),
 			inherent_data.validation_data.relay_parent_storage_root,
@@ -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,15 +274,17 @@ fn validate_validation_data(
 	);
 }
 
-/// Run the given closure with the externalities set.
-fn run_with_externalities<B: BlockT, R, F: FnOnce() -> R>(
+/// 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();
 
-	set_and_run_with_externalities(&mut ext, || execute())
+	recorder::using(recorder, || set_and_run_with_externalities(&mut ext, || execute()))
 }
 
 fn host_storage_read(key: &[u8], value_out: &mut [u8], value_offset: u32) -> Option<u32> {
@@ -305,6 +316,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/tests.rs

@@ -59,7 +59,7 @@ fn call_validate_block(
 }
 
 fn create_test_client() -> (Client, Header) {
-	let client = TestClientBuilder::new().build();
+	let client = TestClientBuilder::new().enable_import_proof_recording().build();
 
 	let genesis_header = client
 		.header(client.chain_info().genesis_hash)

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

@@ -97,6 +97,7 @@ 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(),
@@ -104,6 +105,13 @@ impl<H: Hasher> SizeOnlyRecorderProvider<H> {
 			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

@@ -18,6 +18,7 @@
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
+#[cfg(feature = "std")]
 use sp_externalities::ExternalitiesExt;
 
 use sp_runtime_interface::runtime_interface;
@@ -35,7 +36,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,46 @@
+[package]
+name = "cumulus-primitives-storage-weight-reclaim"
+version = "1.0.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 }
+cumulus-test-runtime = { path = "../../test/runtime" }
+
+[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",
+]

cumulus/primitives/storage-weight-reclaim/src/lib.rs

@@ -0,0 +1,663 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Mechanism to reclaim PoV proof size weight after an extrinsic has been applied.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+use codec::{Decode, Encode};
+use cumulus_primitives_core::Weight;
+use cumulus_primitives_proof_size_hostfunction::{
+	storage_proof_size::storage_proof_size, PROOF_RECORDING_DISABLED,
+};
+use frame_support::{
+	dispatch::{DispatchInfo, PostDispatchInfo},
+	weights::WeightMeter,
+};
+use frame_system::Config;
+use scale_info::TypeInfo;
+use sp_runtime::{
+	traits::{DispatchInfoOf, Dispatchable, PostDispatchInfoOf, SignedExtension},
+	transaction_validity::TransactionValidityError,
+	DispatchResult,
+};
+use sp_std::marker::PhantomData;
+
+const LOG_TARGET: &'static str = "runtime::storage_reclaim";
+
+/// `StorageWeightReclaimer` is a mechanism for manually reclaiming storage weight.
+///
+/// It internally keeps track of the proof size and storage weight at initialization time. At
+/// reclaim  it computes the real consumed storage weight and refunds excess weight.
+///
+/// # Example
+#[doc = docify::embed!("src/lib.rs", simple_reclaimer_example)]
+pub struct StorageWeightReclaimer {
+	previous_remaining_proof_size: u64,
+	previous_reported_proof_size: Option<u64>,
+}
+
+impl StorageWeightReclaimer {
+	/// Creates a new `StorageWeightReclaimer` instance and initializes it with the storage
+	/// size provided by `weight_meter` and reported proof size from the node.
+	#[must_use = "Must call `reclaim_with_meter` to reclaim the weight"]
+	pub fn new(weight_meter: &WeightMeter) -> StorageWeightReclaimer {
+		let previous_remaining_proof_size = weight_meter.remaining().proof_size();
+		let previous_reported_proof_size = get_proof_size();
+		Self { previous_remaining_proof_size, previous_reported_proof_size }
+	}
+
+	/// Check the consumed storage weight and calculate the consumed excess weight.
+	fn reclaim(&mut self, remaining_weight: Weight) -> Option<Weight> {
+		let current_remaining_weight = remaining_weight.proof_size();
+		let current_storage_proof_size = get_proof_size()?;
+		let previous_storage_proof_size = self.previous_reported_proof_size?;
+		let used_weight =
+			self.previous_remaining_proof_size.saturating_sub(current_remaining_weight);
+		let reported_used_size =
+			current_storage_proof_size.saturating_sub(previous_storage_proof_size);
+		let reclaimable = used_weight.saturating_sub(reported_used_size);
+		log::trace!(
+			target: LOG_TARGET,
+			"Found reclaimable storage weight. benchmarked: {used_weight}, consumed: {reported_used_size}"
+		);
+
+		self.previous_remaining_proof_size = current_remaining_weight.saturating_add(reclaimable);
+		self.previous_reported_proof_size = Some(current_storage_proof_size);
+		Some(Weight::from_parts(0, reclaimable))
+	}
+
+	/// Check the consumed storage weight and add the reclaimed
+	/// weight budget back to `weight_meter`.
+	pub fn reclaim_with_meter(&mut self, weight_meter: &mut WeightMeter) -> Option<Weight> {
+		let reclaimed = self.reclaim(weight_meter.remaining())?;
+		weight_meter.reclaim_proof_size(reclaimed.proof_size());
+		Some(reclaimed)
+	}
+}
+
+/// Returns the current storage proof size from the host side.
+///
+/// Returns `None` if proof recording is disabled on the host.
+pub fn get_proof_size() -> Option<u64> {
+	let proof_size = storage_proof_size();
+	(proof_size != PROOF_RECORDING_DISABLED).then_some(proof_size)
+}
+
+/// Storage weight reclaim mechanism.
+///
+/// This extension checks the size of the node-side storage proof
+/// before and after executing a given extrinsic. The difference between
+/// benchmarked and spent weight can be reclaimed.
+#[derive(Encode, Decode, Clone, Eq, PartialEq, Default, TypeInfo)]
+#[scale_info(skip_type_params(T))]
+pub struct StorageWeightReclaim<T: Config + Send + Sync>(PhantomData<T>);
+
+impl<T: Config + Send + Sync> StorageWeightReclaim<T> {
+	/// Create a new `StorageWeightReclaim` instance.
+	pub fn new() -> Self {
+		Self(Default::default())
+	}
+}
+
+impl<T: Config + Send + Sync> core::fmt::Debug for StorageWeightReclaim<T> {
+	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
+		let _ = write!(f, "StorageWeightReclaim");
+		Ok(())
+	}
+}
+
+impl<T: Config + Send + Sync> SignedExtension for StorageWeightReclaim<T>
+where
+	T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
+{
+	const IDENTIFIER: &'static str = "StorageWeightReclaim";
+
+	type AccountId = T::AccountId;
+	type Call = T::RuntimeCall;
+	type AdditionalSigned = ();
+	type Pre = Option<u64>;
+
+	fn additional_signed(
+		&self,
+	) -> Result<Self::AdditionalSigned, sp_runtime::transaction_validity::TransactionValidityError>
+	{
+		Ok(())
+	}
+
+	fn pre_dispatch(
+		self,
+		_who: &Self::AccountId,
+		_call: &Self::Call,
+		_info: &sp_runtime::traits::DispatchInfoOf<Self::Call>,
+		_len: usize,
+	) -> Result<Self::Pre, sp_runtime::transaction_validity::TransactionValidityError> {
+		Ok(get_proof_size())
+	}
+
+	fn post_dispatch(
+		pre: Option<Self::Pre>,
+		info: &DispatchInfoOf<Self::Call>,
+		post_info: &PostDispatchInfoOf<Self::Call>,
+		_len: usize,
+		_result: &DispatchResult,
+	) -> Result<(), TransactionValidityError> {
+		let Some(Some(pre_dispatch_proof_size)) = pre else {
+			return Ok(());
+		};
+
+		let Some(post_dispatch_proof_size) = get_proof_size() else {
+			log::debug!(
+				target: LOG_TARGET,
+				"Proof recording enabled during pre-dispatch, now disabled. This should not happen."
+			);
+			return Ok(())
+		};
+		let benchmarked_weight = info.weight.proof_size();
+		let consumed_weight = post_dispatch_proof_size.saturating_sub(pre_dispatch_proof_size);
+
+		// Unspent weight according to the `actual_weight` from `PostDispatchInfo`
+		// This unspent weight will be refunded by the `CheckWeight` extension, so we need to
+		// account for that.
+		let unspent = post_info.calc_unspent(info).proof_size();
+		let storage_size_diff =
+			benchmarked_weight.saturating_sub(unspent).abs_diff(consumed_weight as u64);
+
+		// This value will be reclaimed by [`frame_system::CheckWeight`], so we need to calculate
+		// that in.
+		frame_system::BlockWeight::<T>::mutate(|current| {
+			if consumed_weight > benchmarked_weight {
+				log::error!(
+					target: LOG_TARGET,
+					"Benchmarked storage weight smaller than consumed storage weight. benchmarked: {benchmarked_weight} consumed: {consumed_weight} unspent: {unspent}"
+				);
+				current.accrue(Weight::from_parts(0, storage_size_diff), info.class)
+			} else {
+				log::trace!(
+					target: LOG_TARGET,
+					"Reclaiming storage weight. benchmarked: {benchmarked_weight}, consumed: {consumed_weight} unspent: {unspent}"
+				);
+				current.reduce(Weight::from_parts(0, storage_size_diff), info.class)
+			}
+		});
+		Ok(())
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use frame_support::{
+		assert_ok,
+		dispatch::DispatchClass,
+		weights::{Weight, WeightMeter},
+	};
+	use frame_system::{BlockWeight, CheckWeight};
+	use sp_runtime::{AccountId32, BuildStorage};
+	use sp_std::marker::PhantomData;
+	use sp_trie::proof_size_extension::ProofSizeExt;
+
+	type Test = cumulus_test_runtime::Runtime;
+	const CALL: &<Test as Config>::RuntimeCall =
+		&cumulus_test_runtime::RuntimeCall::System(frame_system::Call::set_heap_pages {
+			pages: 0u64,
+		});
+	const ALICE: AccountId32 = AccountId32::new([1u8; 32]);
+	const LEN: usize = 0;
+
+	pub fn new_test_ext() -> sp_io::TestExternalities {
+		let ext: sp_io::TestExternalities = cumulus_test_runtime::RuntimeGenesisConfig::default()
+			.build_storage()
+			.unwrap()
+			.into();
+		ext
+	}
+
+	struct TestRecorder {
+		return_values: Box<[usize]>,
+		counter: std::sync::atomic::AtomicUsize,
+	}
+
+	impl TestRecorder {
+		fn new(values: &[usize]) -> Self {
+			TestRecorder { return_values: values.into(), counter: Default::default() }
+		}
+	}
+
+	impl sp_trie::ProofSizeProvider for TestRecorder {
+		fn estimate_encoded_size(&self) -> usize {
+			let counter = self.counter.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+			self.return_values[counter]
+		}
+	}
+
+	fn setup_test_externalities(proof_values: &[usize]) -> sp_io::TestExternalities {
+		let mut test_ext = new_test_ext();
+		let test_recorder = TestRecorder::new(proof_values);
+		test_ext.register_extension(ProofSizeExt::new(test_recorder));
+		test_ext
+	}
+
+	fn set_current_storage_weight(new_weight: u64) {
+		BlockWeight::<Test>::mutate(|current_weight| {
+			current_weight.set(Weight::from_parts(0, new_weight), DispatchClass::Normal);
+		});
+	}
+
+	#[test]
+	fn basic_refund() {
+		// The real cost will be 100 bytes of storage size
+		let mut test_ext = setup_test_externalities(&[0, 100]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 500
+			let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(0));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			// We expect a refund of 400
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 600);
+		})
+	}
+
+	#[test]
+	fn does_nothing_without_extension() {
+		let mut test_ext = new_test_ext();
+
+		// Proof size extension not registered
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 500
+			let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, None);
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1000);
+		})
+	}
+
+	#[test]
+	fn negative_refund_is_added_to_weight() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+			// Benchmarked storage weight: 100
+			let info = DispatchInfo { weight: Weight::from_parts(0, 100), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			// We expect no refund
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1100);
+		})
+	}
+
+	#[test]
+	fn test_zero_proof_size() {
+		let mut test_ext = setup_test_externalities(&[0, 0]);
+
+		test_ext.execute_with(|| {
+			let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(0));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 0);
+		});
+	}
+
+	#[test]
+	fn test_larger_pre_dispatch_proof_size() {
+		let mut test_ext = setup_test_externalities(&[300, 100]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1300);
+
+			let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(300));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 800);
+		});
+	}
+
+	#[test]
+	fn test_incorporates_check_weight_unspent_weight() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 300
+			let info = DispatchInfo { weight: Weight::from_parts(100, 300), ..Default::default() };
+
+			// Actual weight is 50
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(50, 250)),
+				pays_fee: Default::default(),
+			};
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+			// we always need to call `post_dispatch` to verify that they interoperate correctly.
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 900);
+		})
+	}
+
+	#[test]
+	fn test_incorporates_check_weight_unspent_weight_on_negative() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+			// Benchmarked storage weight: 50
+			let info = DispatchInfo { weight: Weight::from_parts(100, 50), ..Default::default() };
+
+			// Actual weight is 25
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(50, 25)),
+				pays_fee: Default::default(),
+			};
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+			// we always need to call `post_dispatch` to verify that they interoperate correctly.
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1150);
+		})
+	}
+
+	#[test]
+	fn test_incorporates_check_weight_unspent_weight_reverse_order() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 300
+			let info = DispatchInfo { weight: Weight::from_parts(100, 300), ..Default::default() };
+
+			// Actual weight is 50
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(50, 250)),
+				pays_fee: Default::default(),
+			};
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+			// `CheckWeight` gets called after `StorageWeightReclaim` this time.
+			// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+			// we always need to call `post_dispatch` to verify that they interoperate correctly.
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 900);
+		})
+	}
+
+	#[test]
+	fn test_incorporates_check_weight_unspent_weight_on_negative_reverse_order() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+			// Benchmarked storage weight: 50
+			let info = DispatchInfo { weight: Weight::from_parts(100, 50), ..Default::default() };
+
+			// Actual weight is 25
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(50, 25)),
+				pays_fee: Default::default(),
+			};
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+			// `CheckWeight` gets called after `StorageWeightReclaim` this time.
+			// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+			// we always need to call `post_dispatch` to verify that they interoperate correctly.
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1150);
+		})
+	}
+
+	#[test]
+	fn storage_size_reported_correctly() {
+		let mut test_ext = setup_test_externalities(&[1000]);
+		test_ext.execute_with(|| {
+			assert_eq!(get_proof_size(), Some(1000));
+		});
+
+		let mut test_ext = new_test_ext();
+
+		let test_recorder = TestRecorder::new(&[0]);
+
+		test_ext.register_extension(ProofSizeExt::new(test_recorder));
+
+		test_ext.execute_with(|| {
+			assert_eq!(get_proof_size(), Some(0));
+		});
+	}
+
+	#[test]
+	fn storage_size_disabled_reported_correctly() {
+		let mut test_ext = setup_test_externalities(&[PROOF_RECORDING_DISABLED as usize]);
+
+		test_ext.execute_with(|| {
+			assert_eq!(get_proof_size(), None);
+		});
+	}
+
+	#[test]
+	fn test_reclaim_helper() {
+		let mut test_ext = setup_test_externalities(&[1000, 1300, 1800]);
+
+		test_ext.execute_with(|| {
+			let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 2000));
+			let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+			remaining_weight_meter.consume(Weight::from_parts(0, 500));
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 200)));
+
+			remaining_weight_meter.consume(Weight::from_parts(0, 800));
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 300)));
+			assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(0, 1200));
+		});
+	}
+
+	#[test]
+	fn test_reclaim_helper_does_not_reclaim_negative() {
+		// Benchmarked weight does not change at all
+		let mut test_ext = setup_test_externalities(&[1000, 1300]);
+
+		test_ext.execute_with(|| {
+			let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 1000));
+			let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 0)));
+			assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(0, 1000));
+		});
+
+		// Benchmarked weight increases less than storage proof consumes
+		let mut test_ext = setup_test_externalities(&[1000, 1300]);
+
+		test_ext.execute_with(|| {
+			let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 1000));
+			let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+			remaining_weight_meter.consume(Weight::from_parts(0, 0));
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 0)));
+		});
+	}
+
+	/// Just here for doc purposes
+	fn get_benched_weight() -> Weight {
+		Weight::from_parts(0, 5)
+	}
+
+	/// Just here for doc purposes
+	fn do_work() {}
+
+	#[docify::export_content(simple_reclaimer_example)]
+	fn reclaim_with_weight_meter() {
+		let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(10, 10));
+
+		let benched_weight = get_benched_weight();
+
+		// It is important to instantiate the `StorageWeightReclaimer` before we consume the weight
+		// for a piece of work from the weight meter.
+		let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+
+		if remaining_weight_meter.try_consume(benched_weight).is_ok() {
+			// Perform some work that takes has `benched_weight` storage weight.
+			do_work();
+
+			// Reclaimer will detect that we only consumed 2 bytes, so 3 bytes are reclaimed.
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+			// We reclaimed 3 bytes of storage size!
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 3)));
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 10);
+			assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(10, 8));
+		}
+	}
+
+	#[test]
+	fn test_reclaim_helper_works_with_meter() {
+		// The node will report 12 - 10 = 2 consumed storage size between the calls.
+		let mut test_ext = setup_test_externalities(&[10, 12]);
+
+		test_ext.execute_with(|| {
+			// Initial storage size is 10.
+			set_current_storage_weight(10);
+			reclaim_with_weight_meter();
+		});
+	}
+}

cumulus/test/client/Cargo.toml

@@ -41,6 +41,7 @@ cumulus-test-relay-sproof-builder = { path = "../relay-sproof-builder" }
 cumulus-primitives-core = { path = "../../primitives/core" }
 cumulus-primitives-proof-size-hostfunction = { path = "../../primitives/proof-size-hostfunction" }
 cumulus-primitives-parachain-inherent = { path = "../../primitives/parachain-inherent" }
+cumulus-primitives-storage-weight-reclaim = { path = "../../primitives/storage-weight-reclaim" }
 
 [features]
 runtime-benchmarks = [

cumulus/test/client/src/lib.rs

@@ -151,6 +151,7 @@ pub fn generate_extrinsic_with_pair(
 		frame_system::CheckNonce::<Runtime>::from(nonce),
 		frame_system::CheckWeight::<Runtime>::new(),
 		pallet_transaction_payment::ChargeTransactionPayment::<Runtime>::from(tip),
+		cumulus_primitives_storage_weight_reclaim::StorageWeightReclaim::<Runtime>::new(),
 	);
 
 	let function = function.into();
@@ -158,7 +159,7 @@ pub fn generate_extrinsic_with_pair(
 	let raw_payload = SignedPayload::from_raw(
 		function.clone(),
 		extra.clone(),
-		((), VERSION.spec_version, genesis_block, current_block_hash, (), (), ()),
+		((), VERSION.spec_version, genesis_block, current_block_hash, (), (), (), ()),
 	);
 	let signature = raw_payload.using_encoded(|e| origin.sign(e));
 
@@ -203,13 +204,16 @@ pub fn validate_block(
 	let mut ext_ext = ext.ext();
 
 	let heap_pages = HeapAllocStrategy::Static { extra_pages: 1024 };
-	let executor = WasmExecutor::<sp_io::SubstrateHostFunctions>::builder()
-		.with_execution_method(WasmExecutionMethod::default())
-		.with_max_runtime_instances(1)
-		.with_runtime_cache_size(2)
-		.with_onchain_heap_alloc_strategy(heap_pages)
-		.with_offchain_heap_alloc_strategy(heap_pages)
-		.build();
+	let executor = WasmExecutor::<(
+		sp_io::SubstrateHostFunctions,
+		cumulus_primitives_proof_size_hostfunction::storage_proof_size::HostFunctions,
+	)>::builder()
+	.with_execution_method(WasmExecutionMethod::default())
+	.with_max_runtime_instances(1)
+	.with_runtime_cache_size(2)
+	.with_onchain_heap_alloc_strategy(heap_pages)
+	.with_offchain_heap_alloc_strategy(heap_pages)
+	.build();
 
 	executor
 		.uncached_call(

cumulus/test/runtime/Cargo.toml

@@ -39,6 +39,7 @@ sp-version = { path = "../../../substrate/primitives/version", default-features
 # Cumulus
 cumulus-pallet-parachain-system = { path = "../../pallets/parachain-system", default-features = false, features = ["parameterized-consensus-hook"] }
 cumulus-primitives-core = { path = "../../primitives/core", default-features = false }
+cumulus-primitives-storage-weight-reclaim = { path = "../../primitives/storage-weight-reclaim", default-features = false }
 
 [build-dependencies]
 substrate-wasm-builder = { path = "../../../substrate/utils/wasm-builder", optional = true }
@@ -49,6 +50,7 @@ std = [
 	"codec/std",
 	"cumulus-pallet-parachain-system/std",
 	"cumulus-primitives-core/std",
+	"cumulus-primitives-storage-weight-reclaim/std",
 	"frame-executive/std",
 	"frame-support/std",
 	"frame-system-rpc-runtime-api/std",

cumulus/test/runtime/src/lib.rs

@@ -331,6 +331,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.
 pub type UncheckedExtrinsic =

cumulus/test/service/Cargo.toml

@@ -81,6 +81,7 @@ cumulus-relay-chain-minimal-node = { path = "../../client/relay-chain-minimal-no
 cumulus-client-pov-recovery = { path = "../../client/pov-recovery" }
 cumulus-test-relay-sproof-builder = { path = "../relay-sproof-builder" }
 cumulus-pallet-parachain-system = { path = "../../pallets/parachain-system", default-features = false, features = ["parameterized-consensus-hook"] }
+cumulus-primitives-storage-weight-reclaim = { path = "../../primitives/storage-weight-reclaim" }
 pallet-timestamp = { path = "../../../substrate/frame/timestamp" }
 
 [dev-dependencies]

cumulus/test/service/src/lib.rs

@@ -112,7 +112,7 @@ pub type AnnounceBlockFn = Arc<dyn Fn(Hash, Option<Vec<u8>>) + Send + Sync>;
 pub struct RuntimeExecutor;
 
 impl sc_executor::NativeExecutionDispatch for RuntimeExecutor {
-	type ExtendHostFunctions = ();
+	type ExtendHostFunctions = cumulus_client_service::storage_proof_size::HostFunctions;
 
 	fn dispatch(method: &str, data: &[u8]) -> Option<Vec<u8>> {
 		cumulus_test_runtime::api::dispatch(method, data)
@@ -894,11 +894,12 @@ pub fn construct_extrinsic(
 		frame_system::CheckNonce::<runtime::Runtime>::from(nonce),
 		frame_system::CheckWeight::<runtime::Runtime>::new(),
 		pallet_transaction_payment::ChargeTransactionPayment::<runtime::Runtime>::from(tip),
+		cumulus_primitives_storage_weight_reclaim::StorageWeightReclaim::<runtime::Runtime>::new(),
 	);
 	let raw_payload = runtime::SignedPayload::from_raw(
 		function.clone(),
 		extra.clone(),
-		((), runtime::VERSION.spec_version, genesis_block, current_block_hash, (), (), ()),
+		((), runtime::VERSION.spec_version, genesis_block, current_block_hash, (), (), (), ()),
 	);
 	let signature = raw_payload.using_encoded(|e| caller.sign(e));
 	runtime::UncheckedExtrinsic::new_signed(

prdoc/pr_3002.prdoc

@@ -0,0 +1,29 @@
+# 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: PoV Reclaim Runtime Side
+author: skunert
+topic: runtime
+doc:
+  - audience: Runtime Dev
+    description: |
+      Adds a mechanism to reclaim proof size weight.
+      1. Introduces a new `SignedExtension` that reclaims the difference
+      between benchmarked proof size weight and actual consumed proof size weight.
+      2. Introduces a manual mechanism, `StorageWeightReclaimer`, to reclaim excess storage weight for situations
+      that require manual weight management. The most prominent case is the `on_idle` hook.
+      3. Adds the `storage_proof_size` host function to the PVF. Parachain nodes should add it to ensure compatibility.
+
+      To enable proof size reclaiming, add the host `storage_proof_size` host function to the parachain node. Add the
+      `StorageWeightReclaim` `SignedExtension` to your runtime and enable proof recording during block import.
+
+
+crates:
+  - name: "cumulus-primitives-storage-weight-reclaim"
+host_functions:
+  - name: "storage_proof_size"
+    description: |
+      This host function is used to pass the current size of the storage proof to the runtime.
+      It was introduced before but becomes relevant now.
+      Note: This host function is intended to be used through `cumulus_primitives_storage_weight_reclaim::get_proof_size`.
+      Direct usage is not recommended.

substrate/frame/system/src/lib.rs

@@ -148,6 +148,7 @@ use sp_io::TestExternalities;
 pub mod limits;
 #[cfg(test)]
 pub(crate) mod mock;
+
 pub mod offchain;
 
 mod extensions;
@@ -847,7 +848,7 @@ pub mod pallet {
 	#[pallet::storage]
 	#[pallet::whitelist_storage]
 	#[pallet::getter(fn block_weight)]
-	pub(super) type BlockWeight<T: Config> = StorageValue<_, ConsumedWeight, ValueQuery>;
+	pub type BlockWeight<T: Config> = StorageValue<_, ConsumedWeight, ValueQuery>;
 
 	/// Total length (in bytes) for all extrinsics put together, for the current block.
 	#[pallet::storage]

substrate/primitives/weights/src/weight_meter.rs

@@ -149,6 +149,11 @@ impl WeightMeter {
 	pub fn can_consume(&self, w: Weight) -> bool {
 		self.consumed.checked_add(&w).map_or(false, |t| t.all_lte(self.limit))
 	}
+
+	/// Reclaim the given weight.
+	pub fn reclaim_proof_size(&mut self, s: u64) {
+		self.consumed.saturating_reduce(Weight::from_parts(0, s));
+	}
 }
 
 #[cfg(test)]
@@ -277,6 +282,21 @@ mod tests {
 		assert_eq!(meter.consumed(), Weight::from_parts(5, 10));
 	}
 
+	#[test]
+	#[cfg(debug_assertions)]
+	fn reclaim_works() {
+		let mut meter = WeightMeter::with_limit(Weight::from_parts(5, 10));
+
+		meter.consume(Weight::from_parts(5, 10));
+		assert_eq!(meter.consumed(), Weight::from_parts(5, 10));
+
+		meter.reclaim_proof_size(3);
+		assert_eq!(meter.consumed(), Weight::from_parts(5, 7));
+
+		meter.reclaim_proof_size(10);
+		assert_eq!(meter.consumed(), Weight::from_parts(5, 0));
+	}
+
 	#[test]
 	#[cfg(debug_assertions)]
 	#[should_panic(expected = "Weight counter overflow")]

substrate/test-utils/client/src/lib.rs

@@ -72,6 +72,7 @@ pub struct TestClientBuilder<Block: BlockT, ExecutorDispatch, Backend: 'static,
 	fork_blocks: ForkBlocks<Block>,
 	bad_blocks: BadBlocks<Block>,
 	enable_offchain_indexing_api: bool,
+	enable_import_proof_recording: bool,
 	no_genesis: bool,
 }
 
@@ -120,6 +121,7 @@ impl<Block: BlockT, ExecutorDispatch, Backend, G: GenesisInit>
 			bad_blocks: None,
 			enable_offchain_indexing_api: false,
 			no_genesis: false,
+			enable_import_proof_recording: false,
 		}
 	}
 
@@ -165,6 +167,12 @@ impl<Block: BlockT, ExecutorDispatch, Backend, G: GenesisInit>
 		self
 	}
 
+	/// Enable proof recording on import.
+	pub fn enable_import_proof_recording(mut self) -> Self {
+		self.enable_import_proof_recording = true;
+		self
+	}
+
 	/// Disable writing genesis.
 	pub fn set_no_genesis(mut self) -> Self {
 		self.no_genesis = true;
@@ -202,6 +210,7 @@ impl<Block: BlockT, ExecutorDispatch, Backend, G: GenesisInit>
 		};
 
 		let client_config = ClientConfig {
+			enable_import_proof_recording: self.enable_import_proof_recording,
 			offchain_indexing_api: self.enable_offchain_indexing_api,
 			no_genesis: self.no_genesis,
 			..Default::default()