From 3386377b0f88cb2c7cfafe8a5f5b3b41b85dde59 Mon Sep 17 00:00:00 2001
From: Sebastian Kunert <skunert49@gmail.com>
Date: Fri, 23 Feb 2024 15:09:49 +0100
Subject: [PATCH] PoV Reclaim Runtime Side (#3002)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
# 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 | 4 +-
Cargo.lock | 23 +
Cargo.toml | 1 +
.../src/validate_block/implementation.rs | 29 +-
.../src/validate_block/tests.rs | 2 +-
.../src/validate_block/trie_recorder.rs | 11 +
.../parachain-template/node/src/command.rs | 3 +-
.../parachain-template/node/src/service.rs | 8 +-
cumulus/parachain-template/runtime/Cargo.toml | 2 +
cumulus/parachain-template/runtime/src/lib.rs | 1 +
cumulus/polkadot-parachain/src/command.rs | 3 +-
cumulus/polkadot-parachain/src/service.rs | 13 +-
.../proof-size-hostfunction/src/lib.rs | 4 +-
.../storage-weight-reclaim/Cargo.toml | 46 ++
.../storage-weight-reclaim/src/lib.rs | 663 ++++++++++++++++++
cumulus/test/client/Cargo.toml | 1 +
cumulus/test/client/src/lib.rs | 20 +-
cumulus/test/runtime/Cargo.toml | 2 +
cumulus/test/runtime/src/lib.rs | 1 +
cumulus/test/service/Cargo.toml | 1 +
cumulus/test/service/src/lib.rs | 5 +-
prdoc/pr_3002.prdoc | 29 +
substrate/frame/system/src/lib.rs | 3 +-
.../primitives/weights/src/weight_meter.rs | 20 +
substrate/test-utils/client/src/lib.rs | 9 +
25 files changed, 875 insertions(+), 29 deletions(-)
create mode 100644 cumulus/primitives/storage-weight-reclaim/Cargo.toml
create mode 100644 cumulus/primitives/storage-weight-reclaim/src/lib.rs
create mode 100644 prdoc/pr_3002.prdoc
diff --git a/.gitlab/pipeline/check.yml b/.gitlab/pipeline/check.yml
index cdb5d1b05d0..4d71a473372 100644
--- a/.gitlab/pipeline/check.yml
+++ b/.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
diff --git a/Cargo.lock b/Cargo.lock
index d9791ec0147..fbe2729acba 100644
--- a/Cargo.lock
+++ b/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",
diff --git a/Cargo.toml b/Cargo.toml
index 774ce1b52a3..1d27cfe9539 100644
--- a/Cargo.toml
+++ b/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",
diff --git a/cumulus/pallets/parachain-system/src/validate_block/implementation.rs b/cumulus/pallets/parachain-system/src/validate_block/implementation.rs
index ce3b724420f..ecab7a9a093 100644
--- a/cumulus/pallets/parachain-system/src/validate_block/implementation.rs
+++ b/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))
}
diff --git a/cumulus/pallets/parachain-system/src/validate_block/tests.rs b/cumulus/pallets/parachain-system/src/validate_block/tests.rs
index f17ac6007a0..a9fb65e1108 100644
--- a/cumulus/pallets/parachain-system/src/validate_block/tests.rs
+++ b/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)
diff --git a/cumulus/pallets/parachain-system/src/validate_block/trie_recorder.rs b/cumulus/pallets/parachain-system/src/validate_block/trie_recorder.rs
index e73aef70aa4..48310670c07 100644
--- a/cumulus/pallets/parachain-system/src/validate_block/trie_recorder.rs
+++ b/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)
}
}
}
diff --git a/cumulus/parachain-template/node/src/command.rs b/cumulus/parachain-template/node/src/command.rs
index 72b3ab7bb4b..82624ae0be5 100644
--- a/cumulus/parachain-template/node/src/command.rs
+++ b/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`."
diff --git a/cumulus/parachain-template/node/src/service.rs b/cumulus/parachain-template/node/src/service.rs
index 830b6e82f96..4dd24803e9b 100644
--- a/cumulus/parachain-template/node/src/service.rs
+++ b/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);
diff --git a/cumulus/parachain-template/runtime/Cargo.toml b/cumulus/parachain-template/runtime/Cargo.toml
index 44d96ffc4e6..1873bd0a23e 100644
--- a/cumulus/parachain-template/runtime/Cargo.toml
+++ b/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",
diff --git a/cumulus/parachain-template/runtime/src/lib.rs b/cumulus/parachain-template/runtime/src/lib.rs
index d9bc111fcef..cee9b33bf37 100644
--- a/cumulus/parachain-template/runtime/src/lib.rs
+++ b/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.
diff --git a/cumulus/polkadot-parachain/src/command.rs b/cumulus/polkadot-parachain/src/command.rs
index a40c356dcd1..4d44879af51 100644
--- a/cumulus/polkadot-parachain/src/command.rs
+++ b/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`."
diff --git a/cumulus/polkadot-parachain/src/service.rs b/cumulus/polkadot-parachain/src/service.rs
index 553975b01a8..4e06cd38f1d 100644
--- a/cumulus/polkadot-parachain/src/service.rs
+++ b/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);
diff --git a/cumulus/primitives/proof-size-hostfunction/src/lib.rs b/cumulus/primitives/proof-size-hostfunction/src/lib.rs
index 6da6235e585..8ebc58ea450 100644
--- a/cumulus/primitives/proof-size-hostfunction/src/lib.rs
+++ b/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())
}
}
diff --git a/cumulus/primitives/storage-weight-reclaim/Cargo.toml b/cumulus/primitives/storage-weight-reclaim/Cargo.toml
new file mode 100644
index 00000000000..4835fb5192b
--- /dev/null
+++ b/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",
+]
diff --git a/cumulus/primitives/storage-weight-reclaim/src/lib.rs b/cumulus/primitives/storage-weight-reclaim/src/lib.rs
new file mode 100644
index 00000000000..5dddc92e395
--- /dev/null
+++ b/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();
+ });
+ }
+}
diff --git a/cumulus/test/client/Cargo.toml b/cumulus/test/client/Cargo.toml
index 7190172101c..028733ce235 100644
--- a/cumulus/test/client/Cargo.toml
+++ b/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 = [
diff --git a/cumulus/test/client/src/lib.rs b/cumulus/test/client/src/lib.rs
index df63f683de6..c46f4da7f67 100644
--- a/cumulus/test/client/src/lib.rs
+++ b/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(
diff --git a/cumulus/test/runtime/Cargo.toml b/cumulus/test/runtime/Cargo.toml
index 5902a62512b..449a8b819bc 100644
--- a/cumulus/test/runtime/Cargo.toml
+++ b/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",
diff --git a/cumulus/test/runtime/src/lib.rs b/cumulus/test/runtime/src/lib.rs
index 6068f895c83..5fb31410984 100644
--- a/cumulus/test/runtime/src/lib.rs
+++ b/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 =
diff --git a/cumulus/test/service/Cargo.toml b/cumulus/test/service/Cargo.toml
index b26f0b9967c..27273f4e0a8 100644
--- a/cumulus/test/service/Cargo.toml
+++ b/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]
diff --git a/cumulus/test/service/src/lib.rs b/cumulus/test/service/src/lib.rs
index 1c2e1db9741..3554a383f21 100644
--- a/cumulus/test/service/src/lib.rs
+++ b/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(
diff --git a/prdoc/pr_3002.prdoc b/prdoc/pr_3002.prdoc
new file mode 100644
index 00000000000..511a07e39c4
--- /dev/null
+++ b/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.
diff --git a/substrate/frame/system/src/lib.rs b/substrate/frame/system/src/lib.rs
index 069217bcee4..1405c7303f8 100644
--- a/substrate/frame/system/src/lib.rs
+++ b/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]
diff --git a/substrate/primitives/weights/src/weight_meter.rs b/substrate/primitives/weights/src/weight_meter.rs
index 584d22304c3..1738948e4c3 100644
--- a/substrate/primitives/weights/src/weight_meter.rs
+++ b/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")]
diff --git a/substrate/test-utils/client/src/lib.rs b/substrate/test-utils/client/src/lib.rs
index e3f06e27563..d283b24f286 100644
--- a/substrate/test-utils/client/src/lib.rs
+++ b/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()
--
GitLab