PoV Reclaim (Clawback) Node Side (#1462)

This PR provides the infrastructure for the pov-reclaim mechanism
discussed in #209. The goal is to provide the current proof size to the
runtime so it can be used to reclaim storage weight.

## New Host Function
- A new host function is provided
[here](https://github.com/skunert/polkadot-sdk/blob/5b317fda3be205f4136f10d4490387ccd4f9765d/cumulus/primitives/pov-reclaim/src/lib.rs#L23).
It returns the size of the current proof size to the runtime. If
recording is not enabled, it returns 0.

## Implementation Overview
- Implement option to enable proof recording during import in the
client. This is currently enabled for `polkadot-parachain`,
`parachain-template` and the cumulus test node.
- Make the proof recorder ready for no-std. It was previously only
enabled for std environments, but we need to record the proof size in
`validate_block` too.
- Provide a recorder implementation that only the records the size of
incoming nodes and does not store the nodes itself.
- Fix benchmarks that were broken by async backing changes
- Provide new externalities extension that is registered by default if
proof recording is enabled.
- I think we should discuss the naming, pov-reclaim was more intuitive
to me, but we could also go with clawback like in the issue.

## Impact of proof recording during import
With proof recording: 6.3058 Kelem/s
Without proof recording: 6.3427 Kelem/s

The measured impact on the importing performance is quite low on my
machine using the block import benchmark. With proof recording I am
seeing a performance hit of 0.585%.

---------

Co-authored-by: command-bot <>
Co-authored-by: Davide Galassi <davxy@datawok.net>
Co-authored-by: Bastian Köcher <git@kchr.de>

cumulus/test/client/Cargo.toml

@@ -36,6 +36,7 @@ cumulus-test-runtime = { path = "../runtime" }
 cumulus-test-service = { path = "../service" }
 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" }
 
 [features]

cumulus/test/client/src/lib.rs

@@ -44,7 +44,8 @@ mod local_executor {
 	pub struct LocalExecutor;
 
 	impl sc_executor::NativeExecutionDispatch for LocalExecutor {
-		type ExtendHostFunctions = ();
+		type ExtendHostFunctions =
+			cumulus_primitives_proof_size_hostfunction::storage_proof_size::HostFunctions;
 
 		fn dispatch(method: &str, data: &[u8]) -> Option<Vec<u8>> {
 			cumulus_test_runtime::api::dispatch(method, data)

cumulus/test/service/benches/block_import.rs

@@ -24,7 +24,7 @@ use core::time::Duration;
 use cumulus_primitives_core::ParaId;
 
 use sp_api::{Core, ProvideRuntimeApi};
-use sp_keyring::Sr25519Keyring::Alice;
+use sp_keyring::Sr25519Keyring::{Alice, Bob};
 
 use cumulus_test_service::bench_utils as utils;
 
@@ -32,51 +32,69 @@ fn benchmark_block_import(c: &mut Criterion) {
 	sp_tracing::try_init_simple();
 
 	let runtime = tokio::runtime::Runtime::new().expect("creating tokio runtime doesn't fail; qed");
-	let para_id = ParaId::from(100);
+
+	let para_id = ParaId::from(cumulus_test_runtime::PARACHAIN_ID);
 	let tokio_handle = runtime.handle();
 
 	// Create enough accounts to fill the block with transactions.
 	// Each account should only be included in one transfer.
 	let (src_accounts, dst_accounts, account_ids) = utils::create_benchmark_accounts();
 
-	let alice = runtime.block_on(
-		cumulus_test_service::TestNodeBuilder::new(para_id, tokio_handle.clone(), Alice)
+	for bench_parameters in &[(true, Alice), (false, Bob)] {
+		let node = runtime.block_on(
+			cumulus_test_service::TestNodeBuilder::new(
+				para_id,
+				tokio_handle.clone(),
+				bench_parameters.1,
+			)
 			// Preload all accounts with funds for the transfers
-			.endowed_accounts(account_ids)
+			.endowed_accounts(account_ids.clone())
+			.import_proof_recording(bench_parameters.0)
 			.build(),
-	);
+		);
 
-	let client = alice.client;
+		let client = node.client;
+		let backend = node.backend;
 
-	let (max_transfer_count, extrinsics) =
-		utils::create_benchmarking_transfer_extrinsics(&client, &src_accounts, &dst_accounts);
+		let (max_transfer_count, extrinsics) =
+			utils::create_benchmarking_transfer_extrinsics(&client, &src_accounts, &dst_accounts);
 
-	let parent_hash = client.usage_info().chain.best_hash;
-	let mut block_builder = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(parent_hash)
-		.fetch_parent_block_number(&*client)
-		.unwrap()
-		.build()
-		.unwrap();
-	for extrinsic in extrinsics {
-		block_builder.push(extrinsic).unwrap();
-	}
-	let benchmark_block = block_builder.build().unwrap();
+		let parent_hash = client.usage_info().chain.best_hash;
+		let mut block_builder = BlockBuilderBuilder::new(&*client)
+			.on_parent_block(parent_hash)
+			.fetch_parent_block_number(&*client)
+			.unwrap()
+			.build()
+			.unwrap();
+		for extrinsic in extrinsics {
+			block_builder.push(extrinsic).unwrap();
+		}
+		let benchmark_block = block_builder.build().unwrap();
 
-	let mut group = c.benchmark_group("Block import");
-	group.sample_size(20);
-	group.measurement_time(Duration::from_secs(120));
-	group.throughput(Throughput::Elements(max_transfer_count as u64));
+		let mut group = c.benchmark_group("Block import");
+		group.sample_size(20);
+		group.measurement_time(Duration::from_secs(120));
+		group.throughput(Throughput::Elements(max_transfer_count as u64));
 
-	group.bench_function(format!("(transfers = {}) block import", max_transfer_count), |b| {
-		b.iter_batched(
-			|| benchmark_block.block.clone(),
-			|block| {
-				client.runtime_api().execute_block(parent_hash, block).unwrap();
+		group.bench_function(
+			format!(
+				"(transfers = {max_transfer_count}, proof_recording = {}) block import",
+				bench_parameters.0
+			),
+			|b| {
+				b.iter_batched(
+					|| {
+						backend.reset_trie_cache();
+						benchmark_block.block.clone()
+					},
+					|block| {
+						client.runtime_api().execute_block(parent_hash, block).unwrap();
+					},
+					BatchSize::SmallInput,
+				)
 			},
-			BatchSize::SmallInput,
-		)
-	});
+		);
+	}
 }
 
 criterion_group!(benches, benchmark_block_import);

cumulus/test/service/benches/block_import_glutton.rs

@@ -27,7 +27,7 @@ use core::time::Duration;
 use cumulus_primitives_core::ParaId;
 
 use sc_block_builder::BlockBuilderBuilder;
-use sp_keyring::Sr25519Keyring::Alice;
+use sp_keyring::Sr25519Keyring::{Alice, Bob, Charlie, Ferdie};
 
 use cumulus_test_service::bench_utils as utils;
 
@@ -38,17 +38,29 @@ fn benchmark_block_import(c: &mut Criterion) {
 	let para_id = ParaId::from(100);
 	let tokio_handle = runtime.handle();
 
-	let alice = runtime.block_on(
-		cumulus_test_service::TestNodeBuilder::new(para_id, tokio_handle.clone(), Alice).build(),
-	);
-	let client = alice.client;
-
-	let mut group = c.benchmark_group("Block import");
-	group.sample_size(20);
-	group.measurement_time(Duration::from_secs(120));
-
 	let mut initialize_glutton_pallet = true;
-	for (compute_ratio, storage_ratio) in &[(One::one(), Zero::zero()), (One::one(), One::one())] {
+	for (compute_ratio, storage_ratio, proof_on_import, keyring_identity) in &[
+		(One::one(), Zero::zero(), true, Alice),
+		(One::one(), One::one(), true, Bob),
+		(One::one(), Zero::zero(), false, Charlie),
+		(One::one(), One::one(), false, Ferdie),
+	] {
+		let node = runtime.block_on(
+			cumulus_test_service::TestNodeBuilder::new(
+				para_id,
+				tokio_handle.clone(),
+				*keyring_identity,
+			)
+			.import_proof_recording(*proof_on_import)
+			.build(),
+		);
+		let client = node.client;
+		let backend = node.backend;
+
+		let mut group = c.benchmark_group("Block import");
+		group.sample_size(20);
+		group.measurement_time(Duration::from_secs(120));
+
 		let block = utils::set_glutton_parameters(
 			&client,
 			initialize_glutton_pallet,
@@ -82,7 +94,10 @@ fn benchmark_block_import(c: &mut Criterion) {
 			),
 			|b| {
 				b.iter_batched(
-					|| benchmark_block.block.clone(),
+					|| {
+						backend.reset_trie_cache();
+						benchmark_block.block.clone()
+					},
 					|block| {
 						client.runtime_api().execute_block(parent_hash, block).unwrap();
 					},

cumulus/test/service/benches/validate_block.rs

@@ -18,7 +18,9 @@
 use codec::{Decode, Encode};
 use core::time::Duration;
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
-use cumulus_primitives_core::{relay_chain::AccountId, PersistedValidationData, ValidationParams};
+use cumulus_primitives_core::{
+	relay_chain::AccountId, ParaId, PersistedValidationData, ValidationParams,
+};
 use cumulus_test_client::{
 	generate_extrinsic_with_pair, BuildParachainBlockData, InitBlockBuilder, TestClientBuilder,
 	ValidationResult,
@@ -83,6 +85,7 @@ fn benchmark_block_validation(c: &mut Criterion) {
 	// Each account should only be included in one transfer.
 	let (src_accounts, dst_accounts, account_ids) = utils::create_benchmark_accounts();
 
+	let para_id = ParaId::from(cumulus_test_runtime::PARACHAIN_ID);
 	let mut test_client_builder = TestClientBuilder::with_default_backend();
 	let genesis_init = test_client_builder.genesis_init_mut();
 	*genesis_init = cumulus_test_client::GenesisParameters { endowed_accounts: account_ids };
@@ -98,7 +101,14 @@ fn benchmark_block_validation(c: &mut Criterion) {
 		..Default::default()
 	};
 
-	let mut block_builder = client.init_block_builder(Some(validation_data), Default::default());
+	let sproof_builder = RelayStateSproofBuilder {
+		included_para_head: Some(parent_header.clone().encode().into()),
+		para_id,
+		..Default::default()
+	};
+
+	let mut block_builder =
+		client.init_block_builder(Some(validation_data), sproof_builder.clone());
 	for extrinsic in extrinsics {
 		block_builder.push(extrinsic).unwrap();
 	}
@@ -108,7 +118,6 @@ fn benchmark_block_validation(c: &mut Criterion) {
 	let proof_size_in_kb = parachain_block.storage_proof().encode().len() as f64 / 1024f64;
 	let runtime = utils::get_wasm_module();
 
-	let sproof_builder: RelayStateSproofBuilder = Default::default();
 	let (relay_parent_storage_root, _) = sproof_builder.into_state_root_and_proof();
 	let encoded_params = ValidationParams {
 		block_data: cumulus_test_client::BlockData(parachain_block.encode()),

cumulus/test/service/src/bench_utils.rs

@@ -81,8 +81,13 @@ pub fn extrinsic_set_time(client: &TestClient) -> OpaqueExtrinsic {
 pub fn extrinsic_set_validation_data(
 	parent_header: cumulus_test_runtime::Header,
 ) -> OpaqueExtrinsic {
-	let sproof_builder = RelayStateSproofBuilder { para_id: 100.into(), ..Default::default() };
 	let parent_head = HeadData(parent_header.encode());
+	let sproof_builder = RelayStateSproofBuilder {
+		para_id: cumulus_test_runtime::PARACHAIN_ID.into(),
+		included_para_head: parent_head.clone().into(),
+		..Default::default()
+	};
+
 	let (relay_parent_storage_root, relay_chain_state) = sproof_builder.into_state_root_and_proof();
 	let data = ParachainInherentData {
 		validation_data: PersistedValidationData {

cumulus/test/service/src/lib.rs

@@ -187,6 +187,7 @@ impl RecoveryHandle for FailingRecoveryHandle {
 /// be able to perform chain operations.
 pub fn new_partial(
 	config: &mut Configuration,
+	enable_import_proof_record: bool,
 ) -> Result<
 	PartialComponents<
 		Client,
@@ -214,7 +215,12 @@ pub fn new_partial(
 		sc_executor::NativeElseWasmExecutor::<RuntimeExecutor>::new_with_wasm_executor(wasm);
 
 	let (client, backend, keystore_container, task_manager) =
-		sc_service::new_full_parts::<Block, RuntimeApi, _>(config, None, executor)?;
+		sc_service::new_full_parts_record_import::<Block, RuntimeApi, _>(
+			config,
+			None,
+			executor,
+			enable_import_proof_record,
+		)?;
 	let client = Arc::new(client);
 
 	let block_import =
@@ -309,19 +315,21 @@ pub async fn start_node_impl<RB>(
 	rpc_ext_builder: RB,
 	consensus: Consensus,
 	collator_options: CollatorOptions,
+	proof_recording_during_import: bool,
 ) -> sc_service::error::Result<(
 	TaskManager,
 	Arc<Client>,
 	Arc<NetworkService<Block, H256>>,
 	RpcHandlers,
 	TransactionPool,
+	Arc<Backend>,
 )>
 where
 	RB: Fn(Arc<Client>) -> Result<jsonrpsee::RpcModule<()>, sc_service::Error> + Send + 'static,
 {
 	let mut parachain_config = prepare_node_config(parachain_config);
 
-	let params = new_partial(&mut parachain_config)?;
+	let params = new_partial(&mut parachain_config, proof_recording_during_import)?;
 
 	let transaction_pool = params.transaction_pool.clone();
 	let mut task_manager = params.task_manager;
@@ -477,7 +485,7 @@ where
 
 	start_network.start_network();
 
-	Ok((task_manager, client, network, rpc_handlers, transaction_pool))
+	Ok((task_manager, client, network, rpc_handlers, transaction_pool, backend))
 }
 
 /// A Cumulus test node instance used for testing.
@@ -495,6 +503,8 @@ pub struct TestNode {
 	pub rpc_handlers: RpcHandlers,
 	/// Node's transaction pool
 	pub transaction_pool: TransactionPool,
+	/// Node's backend
+	pub backend: Arc<Backend>,
 }
 
 #[allow(missing_docs)]
@@ -520,6 +530,7 @@ pub struct TestNodeBuilder {
 	consensus: Consensus,
 	relay_chain_mode: RelayChainMode,
 	endowed_accounts: Vec<AccountId>,
+	record_proof_during_import: bool,
 }
 
 impl TestNodeBuilder {
@@ -544,6 +555,7 @@ impl TestNodeBuilder {
 			consensus: Consensus::RelayChain,
 			endowed_accounts: Default::default(),
 			relay_chain_mode: RelayChainMode::Embedded,
+			record_proof_during_import: true,
 		}
 	}
 
@@ -656,6 +668,12 @@ impl TestNodeBuilder {
 		self
 	}
 
+	/// Record proofs during import.
+	pub fn import_proof_recording(mut self, should_record_proof: bool) -> TestNodeBuilder {
+		self.record_proof_during_import = should_record_proof;
+		self
+	}
+
 	/// Build the [`TestNode`].
 	pub async fn build(self) -> TestNode {
 		let parachain_config = node_config(
@@ -684,24 +702,26 @@ impl TestNodeBuilder {
 			format!("{} (relay chain)", relay_chain_config.network.node_name);
 
 		let multiaddr = parachain_config.network.listen_addresses[0].clone();
-		let (task_manager, client, network, rpc_handlers, transaction_pool) = start_node_impl(
-			parachain_config,
-			self.collator_key,
-			relay_chain_config,
-			self.para_id,
-			self.wrap_announce_block,
-			false,
-			|_| Ok(jsonrpsee::RpcModule::new(())),
-			self.consensus,
-			collator_options,
-		)
-		.await
-		.expect("could not create Cumulus test service");
+		let (task_manager, client, network, rpc_handlers, transaction_pool, backend) =
+			start_node_impl(
+				parachain_config,
+				self.collator_key,
+				relay_chain_config,
+				self.para_id,
+				self.wrap_announce_block,
+				false,
+				|_| Ok(jsonrpsee::RpcModule::new(())),
+				self.consensus,
+				collator_options,
+				self.record_proof_during_import,
+			)
+			.await
+			.expect("could not create Cumulus test service");
 
 		let peer_id = network.local_peer_id();
 		let addr = MultiaddrWithPeerId { multiaddr, peer_id };
 
-		TestNode { task_manager, client, network, addr, rpc_handlers, transaction_pool }
+		TestNode { task_manager, client, network, addr, rpc_handlers, transaction_pool, backend }
 	}
 }
 

cumulus/test/service/src/main.rs

@@ -128,7 +128,7 @@ fn main() -> Result<(), sc_cli::Error> {
 				})
 				.unwrap_or(cumulus_test_service::Consensus::RelayChain);
 
-			let (mut task_manager, _, _, _, _) = tokio_runtime
+			let (mut task_manager, _, _, _, _, _) = tokio_runtime
 				.block_on(cumulus_test_service::start_node_impl(
 					config,
 					collator_key,
@@ -139,6 +139,7 @@ fn main() -> Result<(), sc_cli::Error> {
 					|_| Ok(jsonrpsee::RpcModule::new(())),
 					consensus,
 					collator_options,
+					true,
 				))
 				.expect("could not create Cumulus test service");