PoV Reclaim (Clawback) Node Side

cumulus/pallets/parachain-system/Cargo.toml

@@ -35,11 +35,14 @@ xcm = { package = "staging-xcm", path = "../../../polkadot/xcm", default-feature
 cumulus-pallet-parachain-system-proc-macro = { path = "proc-macro", default-features = false }
 cumulus-primitives-core = { path = "../../primitives/core", default-features = false }
 cumulus-primitives-parachain-inherent = { path = "../../primitives/parachain-inherent", default-features = false }
+cumulus-primitives-pov-reclaim = { path = "../../primitives/pov-reclaim", default-features = false }
 
 [dev-dependencies]
 assert_matches = "1.5"
 hex-literal = "0.4.1"
 lazy_static = "1.4"
+trie-standardmap = "0.16.0"
+rand = "0.8.5"
 
 # Substrate
 sc-client-api = { path = "../../../substrate/client/api" }
@@ -58,6 +61,7 @@ std = [
 	"cumulus-pallet-parachain-system-proc-macro/std",
 	"cumulus-primitives-core/std",
 	"cumulus-primitives-parachain-inherent/std",
+	"cumulus-primitives-pov-reclaim/std",
 	"environmental/std",
 	"frame-support/std",
 	"frame-system/std",

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,
 };
@@ -33,13 +33,15 @@ use sp_core::storage::{ChildInfo, StateVersion};
 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_std::{prelude::*, sync::Arc};
+use sp_trie::{MemoryDB, ProofSizeProvider, TrieRecorderProvider};
+use trie_recorder::SizeOnlyRecorderProvider;
 
 type TrieBackend<B> = sp_state_machine::TrieBackend<
 	MemoryDB<HashingFor<B>>,
 	HashingFor<B>,
 	trie_cache::CacheProvider<HashingFor<B>>,
+	SizeOnlyRecorderProvider<HashingFor<B>>,
 >;
 
 type Ext<'a, B> = sp_state_machine::Ext<'a, HashingFor<B>, TrieBackend<B>>;
@@ -48,6 +50,10 @@ 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.
+/// Static mut is fine here because we are single-threaded in WASM.
+static mut RECORDER: Option<Box<dyn ProofSizeProvider + Send + Sync>> = None;
+
 /// Validate the given parachain block.
 ///
 /// This function is doing roughly the following:
@@ -120,6 +126,7 @@ where
 
 	sp_std::mem::drop(storage_proof);
 
+	let 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,8 +135,11 @@ where
 		*parent_header.state_root(),
 		cache_provider,
 	)
+	.with_recorder(recorder.clone())
 	.build();
 
+	set_recorder(recorder);
+
 	let _guard = (
 		// Replace storage calls with our own implementations
 		sp_io::storage::host_read.replace_implementation(host_storage_read),
@@ -167,6 +177,8 @@ where
 			.replace_implementation(host_default_child_storage_next_key),
 		sp_io::offchain_index::host_set.replace_implementation(host_offchain_index_set),
 		sp_io::offchain_index::host_clear.replace_implementation(host_offchain_index_clear),
+		cumulus_primitives_pov_reclaim::pov_reclaim_host_functions::host_current_storage_proof_size
+			.replace_implementation(host_current_storage_proof_size),
 	);
 
 	run_with_externalities::<B, _, _>(&backend, || {
@@ -305,6 +317,14 @@ fn host_storage_clear(key: &[u8]) {
 	with_externalities(|ext| ext.place_storage(key.to_vec(), None))
 }
 
+fn host_current_storage_proof_size() -> u32 {
+	get_recorder_ref()
+		.map(|r| r.estimate_encoded_size())
+		.unwrap_or_default()
+		.try_into()
+		.unwrap_or_default()
+}
+
 fn host_storage_root(version: StateVersion) -> Vec<u8> {
 	with_externalities(|ext| ext.storage_root(version))
 }
@@ -406,3 +426,14 @@ fn host_default_child_storage_next_key(storage_key: &[u8], key: &[u8]) -> Option
 fn host_offchain_index_set(_key: &[u8], _value: &[u8]) {}
 
 fn host_offchain_index_clear(_key: &[u8]) {}
+
+fn set_recorder<H: sp_core::Hasher + 'static>(recorder: SizeOnlyRecorderProvider<H>) {
+	// This is safe here, there is strictly sequential access.
+	unsafe {
+		RECORDER = Some(Box::new(recorder));
+	}
+}
+
+fn get_recorder_ref() -> Option<&'static Box<dyn ProofSizeProvider + Send + Sync>> {
+	unsafe { RECORDER.as_ref() }
+}

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

@@ -26,6 +26,10 @@ mod tests;
 #[doc(hidden)]
 mod trie_cache;
 
+#[cfg(any(test, not(feature = "std")))]
+#[doc(hidden)]
+mod trie_recorder;
+
 #[cfg(not(feature = "std"))]
 #[doc(hidden)]
 pub use bytes;