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>

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