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

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

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

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

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

cumulus/parachain-template/runtime/Cargo.toml

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

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

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

cumulus/polkadot-parachain/src/command.rs

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

cumulus/polkadot-parachain/src/service.rs

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

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

@@ -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())
 	}
 }