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>

substrate/primitives/api/proc-macro/src/decl_runtime_apis.rs

@@ -729,7 +729,7 @@ fn decl_runtime_apis_impl_inner(api_decls: &[ItemTrait]) -> Result<TokenStream>
 	};
 
 	let decl = expander::Expander::new("decl_runtime_apis")
-		.dry(std::env::var("SP_API_EXPAND").is_err())
+		.dry(std::env::var("EXPAND_MACROS").is_err())
 		.verbose(true)
 		.write_to_out_dir(decl)
 		.expect("Does not fail because of IO in OUT_DIR; qed");

substrate/primitives/api/proc-macro/src/impl_runtime_apis.rs

@@ -846,7 +846,7 @@ fn impl_runtime_apis_impl_inner(api_impls: &[ItemImpl]) -> Result<TokenStream> {
 	);
 
 	let impl_ = expander::Expander::new("impl_runtime_apis")
-		.dry(std::env::var("SP_API_EXPAND").is_err())
+		.dry(std::env::var("EXPAND_MACROS").is_err())
 		.verbose(true)
 		.write_to_out_dir(impl_)
 		.expect("Does not fail because of IO in OUT_DIR; qed");

substrate/primitives/runtime-interface/proc-macro/Cargo.toml

@@ -20,4 +20,5 @@ Inflector = "0.11.4"
 proc-macro-crate = "1.1.3"
 proc-macro2 = "1.0.56"
 quote = "1.0.28"
+expander = "2.0.0"
 syn = { version = "2.0.38", features = ["full", "visit", "fold", "extra-traits"] }

substrate/primitives/runtime-interface/proc-macro/src/runtime_interface/mod.rs

@@ -68,5 +68,11 @@ pub fn runtime_interface_impl(
 		}
 	};
 
+	let res = expander::Expander::new("runtime_interface")
+		.dry(std::env::var("EXPAND_MACROS").is_err())
+		.verbose(true)
+		.write_to_out_dir(res)
+		.expect("Does not fail because of IO in OUT_DIR; qed");
+
 	Ok(res)
 }

substrate/primitives/state-machine/src/trie_backend.rs

@@ -33,12 +33,12 @@ use sp_core::storage::{ChildInfo, StateVersion};
 #[cfg(feature = "std")]
 use sp_trie::{
 	cache::{LocalTrieCache, TrieCache},
-	recorder::Recorder,
-	MemoryDB, StorageProof,
+	MemoryDB,
 };
 #[cfg(not(feature = "std"))]
 use sp_trie::{Error, NodeCodec};
-use sp_trie::{MerkleValue, PrefixedMemoryDB};
+use sp_trie::{MerkleValue, PrefixedMemoryDB, StorageProof, TrieRecorderProvider};
+
 use trie_db::TrieCache as TrieCacheT;
 #[cfg(not(feature = "std"))]
 use trie_db::{node::NodeOwned, CachedValue};
@@ -112,8 +112,6 @@ pub struct UnimplementedCacheProvider<H> {
 	// Not strictly necessary, but the H bound allows to use this as a drop-in
 	// replacement for the `LocalTrieCache` in no-std contexts.
 	_phantom: core::marker::PhantomData<H>,
-	// Statically prevents construction.
-	_infallible: core::convert::Infallible,
 }
 
 #[cfg(not(feature = "std"))]
@@ -156,52 +154,83 @@ impl<H: Hasher> TrieCacheProvider<H> for UnimplementedCacheProvider<H> {
 	}
 }
 
+/// Recorder provider that allows construction of a [`TrieBackend`] and satisfies the requirements,
+/// but can never be instantiated.
+#[cfg(not(feature = "std"))]
+pub struct UnimplementedRecorderProvider<H> {
+	// Not strictly necessary, but the H bound allows to use this as a drop-in
+	// replacement for the [`sp_trie::recorder::Recorder`] in no-std contexts.
+	_phantom: core::marker::PhantomData<H>,
+}
+
+#[cfg(not(feature = "std"))]
+impl<H: Hasher> trie_db::TrieRecorder<H::Out> for UnimplementedRecorderProvider<H> {
+	fn record<'a>(&mut self, _access: trie_db::TrieAccess<'a, H::Out>) {
+		unimplemented!()
+	}
+
+	fn trie_nodes_recorded_for_key(&self, _key: &[u8]) -> trie_db::RecordedForKey {
+		unimplemented!()
+	}
+}
+
+#[cfg(not(feature = "std"))]
+impl<H: Hasher> TrieRecorderProvider<H> for UnimplementedRecorderProvider<H> {
+	type Recorder<'a> = UnimplementedRecorderProvider<H> where H: 'a;
+
+	fn drain_storage_proof(self) -> Option<StorageProof> {
+		unimplemented!()
+	}
+
+	fn as_trie_recorder(&self, _storage_root: H::Out) -> Self::Recorder<'_> {
+		unimplemented!()
+	}
+}
+
 #[cfg(feature = "std")]
 type DefaultCache<H> = LocalTrieCache<H>;
 
 #[cfg(not(feature = "std"))]
 type DefaultCache<H> = UnimplementedCacheProvider<H>;
 
+#[cfg(feature = "std")]
+type DefaultRecorder<H> = sp_trie::recorder::Recorder<H>;
+
+#[cfg(not(feature = "std"))]
+type DefaultRecorder<H> = UnimplementedRecorderProvider<H>;
+
 /// Builder for creating a [`TrieBackend`].
-pub struct TrieBackendBuilder<S: TrieBackendStorage<H>, H: Hasher, C = DefaultCache<H>> {
+pub struct TrieBackendBuilder<
+	S: TrieBackendStorage<H>,
+	H: Hasher,
+	C = DefaultCache<H>,
+	R = DefaultRecorder<H>,
+> {
 	storage: S,
 	root: H::Out,
-	#[cfg(feature = "std")]
-	recorder: Option<Recorder<H>>,
+	recorder: Option<R>,
 	cache: Option<C>,
 }
 
-impl<S, H> TrieBackendBuilder<S, H, DefaultCache<H>>
+impl<S, H> TrieBackendBuilder<S, H>
 where
 	S: TrieBackendStorage<H>,
 	H: Hasher,
 {
 	/// Create a new builder instance.
 	pub fn new(storage: S, root: H::Out) -> Self {
-		Self {
-			storage,
-			root,
-			#[cfg(feature = "std")]
-			recorder: None,
-			cache: None,
-		}
+		Self { storage, root, recorder: None, cache: None }
 	}
 }
 
-impl<S, H, C> TrieBackendBuilder<S, H, C>
+impl<S, H, C, R> TrieBackendBuilder<S, H, C, R>
 where
 	S: TrieBackendStorage<H>,
 	H: Hasher,
 {
 	/// Create a new builder instance.
 	pub fn new_with_cache(storage: S, root: H::Out, cache: C) -> Self {
-		Self {
-			storage,
-			root,
-			#[cfg(feature = "std")]
-			recorder: None,
-			cache: Some(cache),
-		}
+		Self { storage, root, recorder: None, cache: Some(cache) }
 	}
 	/// Wrap the given [`TrieBackend`].
 	///
@@ -210,53 +239,47 @@ where
 	/// backend.
 	///
 	/// The backend storage and the cache will be taken from `other`.
-	pub fn wrap(other: &TrieBackend<S, H, C>) -> TrieBackendBuilder<&S, H, &C> {
+	pub fn wrap(other: &TrieBackend<S, H, C, R>) -> TrieBackendBuilder<&S, H, &C, R> {
 		TrieBackendBuilder {
 			storage: other.essence.backend_storage(),
 			root: *other.essence.root(),
-			#[cfg(feature = "std")]
 			recorder: None,
 			cache: other.essence.trie_node_cache.as_ref(),
 		}
 	}
 
 	/// Use the given optional `recorder` for the to be configured [`TrieBackend`].
-	#[cfg(feature = "std")]
-	pub fn with_optional_recorder(self, recorder: Option<Recorder<H>>) -> Self {
+	pub fn with_optional_recorder(self, recorder: Option<R>) -> Self {
 		Self { recorder, ..self }
 	}
 
 	/// Use the given `recorder` for the to be configured [`TrieBackend`].
-	#[cfg(feature = "std")]
-	pub fn with_recorder(self, recorder: Recorder<H>) -> Self {
+	pub fn with_recorder(self, recorder: R) -> Self {
 		Self { recorder: Some(recorder), ..self }
 	}
 
 	/// Use the given optional `cache` for the to be configured [`TrieBackend`].
-	pub fn with_optional_cache<LC>(self, cache: Option<LC>) -> TrieBackendBuilder<S, H, LC> {
+	pub fn with_optional_cache<LC>(self, cache: Option<LC>) -> TrieBackendBuilder<S, H, LC, R> {
 		TrieBackendBuilder {
 			cache,
 			root: self.root,
 			storage: self.storage,
-			#[cfg(feature = "std")]
 			recorder: self.recorder,
 		}
 	}
 
 	/// Use the given `cache` for the to be configured [`TrieBackend`].
-	pub fn with_cache<LC>(self, cache: LC) -> TrieBackendBuilder<S, H, LC> {
+	pub fn with_cache<LC>(self, cache: LC) -> TrieBackendBuilder<S, H, LC, R> {
 		TrieBackendBuilder {
 			cache: Some(cache),
 			root: self.root,
 			storage: self.storage,
-			#[cfg(feature = "std")]
 			recorder: self.recorder,
 		}
 	}
 
 	/// Build the configured [`TrieBackend`].
-	#[cfg(feature = "std")]
-	pub fn build(self) -> TrieBackend<S, H, C> {
+	pub fn build(self) -> TrieBackend<S, H, C, R> {
 		TrieBackend {
 			essence: TrieBackendEssence::new_with_cache_and_recorder(
 				self.storage,
@@ -267,27 +290,18 @@ where
 			next_storage_key_cache: Default::default(),
 		}
 	}
-
-	/// Build the configured [`TrieBackend`].
-	#[cfg(not(feature = "std"))]
-	pub fn build(self) -> TrieBackend<S, H, C> {
-		TrieBackend {
-			essence: TrieBackendEssence::new_with_cache(self.storage, self.root, self.cache),
-			next_storage_key_cache: Default::default(),
-		}
-	}
 }
 
 /// A cached iterator.
-struct CachedIter<S, H, C>
+struct CachedIter<S, H, C, R>
 where
 	H: Hasher,
 {
 	last_key: sp_std::vec::Vec<u8>,
-	iter: RawIter<S, H, C>,
+	iter: RawIter<S, H, C, R>,
 }
 
-impl<S, H, C> Default for CachedIter<S, H, C>
+impl<S, H, C, R> Default for CachedIter<S, H, C, R>
 where
 	H: Hasher,
 {
@@ -313,23 +327,32 @@ fn access_cache<T, R>(cell: &CacheCell<T>, callback: impl FnOnce(&mut T) -> R) -
 }
 
 /// Patricia trie-based backend. Transaction type is an overlay of changes to commit.
-pub struct TrieBackend<S: TrieBackendStorage<H>, H: Hasher, C = DefaultCache<H>> {
-	pub(crate) essence: TrieBackendEssence<S, H, C>,
-	next_storage_key_cache: CacheCell<Option<CachedIter<S, H, C>>>,
+pub struct TrieBackend<
+	S: TrieBackendStorage<H>,
+	H: Hasher,
+	C = DefaultCache<H>,
+	R = DefaultRecorder<H>,
+> {
+	pub(crate) essence: TrieBackendEssence<S, H, C, R>,
+	next_storage_key_cache: CacheCell<Option<CachedIter<S, H, C, R>>>,
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
-	TrieBackend<S, H, C>
+impl<
+		S: TrieBackendStorage<H>,
+		H: Hasher,
+		C: TrieCacheProvider<H> + Send + Sync,
+		R: TrieRecorderProvider<H> + Send + Sync,
+	> TrieBackend<S, H, C, R>
 where
 	H::Out: Codec,
 {
 	#[cfg(test)]
-	pub(crate) fn from_essence(essence: TrieBackendEssence<S, H, C>) -> Self {
+	pub(crate) fn from_essence(essence: TrieBackendEssence<S, H, C, R>) -> Self {
 		Self { essence, next_storage_key_cache: Default::default() }
 	}
 
 	/// Get backend essence reference.
-	pub fn essence(&self) -> &TrieBackendEssence<S, H, C> {
+	pub fn essence(&self) -> &TrieBackendEssence<S, H, C, R> {
 		&self.essence
 	}
 
@@ -361,28 +384,31 @@ where
 	/// Extract the [`StorageProof`].
 	///
 	/// This only returns `Some` when there was a recorder set.
-	#[cfg(feature = "std")]
 	pub fn extract_proof(mut self) -> Option<StorageProof> {
-		self.essence.recorder.take().map(|r| r.drain_storage_proof())
+		self.essence.recorder.take().and_then(|r| r.drain_storage_proof())
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H>> sp_std::fmt::Debug
-	for TrieBackend<S, H, C>
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H>, R: TrieRecorderProvider<H>>
+	sp_std::fmt::Debug for TrieBackend<S, H, C, R>
 {
 	fn fmt(&self, f: &mut sp_std::fmt::Formatter<'_>) -> sp_std::fmt::Result {
 		write!(f, "TrieBackend")
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync> Backend<H>
-	for TrieBackend<S, H, C>
+impl<
+		S: TrieBackendStorage<H>,
+		H: Hasher,
+		C: TrieCacheProvider<H> + Send + Sync,
+		R: TrieRecorderProvider<H> + Send + Sync,
+	> Backend<H> for TrieBackend<S, H, C, R>
 where
 	H::Out: Ord + Codec,
 {
 	type Error = crate::DefaultError;
 	type TrieBackendStorage = S;
-	type RawIter = crate::trie_backend_essence::RawIter<S, H, C>;
+	type RawIter = crate::trie_backend_essence::RawIter<S, H, C, R>;
 
 	fn storage_hash(&self, key: &[u8]) -> Result<Option<H::Out>, Self::Error> {
 		self.essence.storage_hash(key)

substrate/primitives/state-machine/src/trie_backend_essence.rs

@@ -28,19 +28,19 @@ use hash_db::{self, AsHashDB, HashDB, HashDBRef, Hasher, Prefix};
 #[cfg(feature = "std")]
 use parking_lot::RwLock;
 use sp_core::storage::{ChildInfo, ChildType, StateVersion};
-use sp_std::{boxed::Box, marker::PhantomData, vec::Vec};
 #[cfg(feature = "std")]
-use sp_trie::recorder::Recorder;
+use sp_std::sync::Arc;
+use sp_std::{boxed::Box, marker::PhantomData, vec::Vec};
 use sp_trie::{
 	child_delta_trie_root, delta_trie_root, empty_child_trie_root,
 	read_child_trie_first_descedant_value, read_child_trie_hash, read_child_trie_value,
 	read_trie_first_descedant_value, read_trie_value,
 	trie_types::{TrieDBBuilder, TrieError},
 	DBValue, KeySpacedDB, MerkleValue, NodeCodec, PrefixedMemoryDB, Trie, TrieCache,
-	TrieDBRawIterator, TrieRecorder,
+	TrieDBRawIterator, TrieRecorder, TrieRecorderProvider,
 };
 #[cfg(feature = "std")]
-use std::{collections::HashMap, sync::Arc};
+use std::collections::HashMap;
 // In this module, we only use layout for read operation and empty root,
 // where V1 and V0 are equivalent.
 use sp_trie::LayoutV1 as Layout;
@@ -83,7 +83,7 @@ enum IterState {
 }
 
 /// A raw iterator over the storage.
-pub struct RawIter<S, H, C>
+pub struct RawIter<S, H, C, R>
 where
 	H: Hasher,
 {
@@ -93,25 +93,26 @@ where
 	child_info: Option<ChildInfo>,
 	trie_iter: TrieDBRawIterator<Layout<H>>,
 	state: IterState,
-	_phantom: PhantomData<(S, C)>,
+	_phantom: PhantomData<(S, C, R)>,
 }
 
-impl<S, H, C> RawIter<S, H, C>
+impl<S, H, C, R> RawIter<S, H, C, R>
 where
 	H: Hasher,
 	S: TrieBackendStorage<H>,
 	H::Out: Codec + Ord,
 	C: TrieCacheProvider<H> + Send + Sync,
+	R: TrieRecorderProvider<H> + Send + Sync,
 {
 	#[inline]
-	fn prepare<R>(
+	fn prepare<RE>(
 		&mut self,
-		backend: &TrieBackendEssence<S, H, C>,
+		backend: &TrieBackendEssence<S, H, C, R>,
 		callback: impl FnOnce(
 			&sp_trie::TrieDB<Layout<H>>,
 			&mut TrieDBRawIterator<Layout<H>>,
-		) -> Option<core::result::Result<R, Box<TrieError<<H as Hasher>::Out>>>>,
-	) -> Option<Result<R>> {
+		) -> Option<core::result::Result<RE, Box<TrieError<<H as Hasher>::Out>>>>,
+	) -> Option<Result<RE>> {
 		if !matches!(self.state, IterState::Pending) {
 			return None
 		}
@@ -139,7 +140,7 @@ where
 	}
 }
 
-impl<S, H, C> Default for RawIter<S, H, C>
+impl<S, H, C, R> Default for RawIter<S, H, C, R>
 where
 	H: Hasher,
 {
@@ -156,14 +157,15 @@ where
 	}
 }
 
-impl<S, H, C> StorageIterator<H> for RawIter<S, H, C>
+impl<S, H, C, R> StorageIterator<H> for RawIter<S, H, C, R>
 where
 	H: Hasher,
 	S: TrieBackendStorage<H>,
 	H::Out: Codec + Ord,
 	C: TrieCacheProvider<H> + Send + Sync,
+	R: TrieRecorderProvider<H> + Send + Sync,
 {
-	type Backend = crate::TrieBackend<S, H, C>;
+	type Backend = crate::TrieBackend<S, H, C, R>;
 	type Error = crate::DefaultError;
 
 	#[inline]
@@ -204,18 +206,17 @@ where
 }
 
 /// Patricia trie-based pairs storage essence.
-pub struct TrieBackendEssence<S: TrieBackendStorage<H>, H: Hasher, C> {
+pub struct TrieBackendEssence<S: TrieBackendStorage<H>, H: Hasher, C, R> {
 	storage: S,
 	root: H::Out,
 	empty: H::Out,
 	#[cfg(feature = "std")]
 	pub(crate) cache: Arc<RwLock<Cache<H::Out>>>,
 	pub(crate) trie_node_cache: Option<C>,
-	#[cfg(feature = "std")]
-	pub(crate) recorder: Option<Recorder<H>>,
+	pub(crate) recorder: Option<R>,
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C> TrieBackendEssence<S, H, C> {
+impl<S: TrieBackendStorage<H>, H: Hasher, C, R> TrieBackendEssence<S, H, C, R> {
 	/// Create new trie-based backend.
 	pub fn new(storage: S, root: H::Out) -> Self {
 		Self::new_with_cache(storage, root, None)
@@ -230,23 +231,22 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C> TrieBackendEssence<S, H, C> {
 			#[cfg(feature = "std")]
 			cache: Arc::new(RwLock::new(Cache::new())),
 			trie_node_cache: cache,
-			#[cfg(feature = "std")]
 			recorder: None,
 		}
 	}
 
 	/// Create new trie-based backend.
-	#[cfg(feature = "std")]
 	pub fn new_with_cache_and_recorder(
 		storage: S,
 		root: H::Out,
 		cache: Option<C>,
-		recorder: Option<Recorder<H>>,
+		recorder: Option<R>,
 	) -> Self {
 		TrieBackendEssence {
 			storage,
 			root,
 			empty: H::hash(&[0u8]),
+			#[cfg(feature = "std")]
 			cache: Arc::new(RwLock::new(Cache::new())),
 			trie_node_cache: cache,
 			recorder,
@@ -289,37 +289,31 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C> TrieBackendEssence<S, H, C> {
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H>> TrieBackendEssence<S, H, C> {
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H>, R: TrieRecorderProvider<H>>
+	TrieBackendEssence<S, H, C, R>
+{
 	/// Call the given closure passing it the recorder and the cache.
 	///
 	/// If the given `storage_root` is `None`, `self.root` will be used.
 	#[inline]
-	fn with_recorder_and_cache<R>(
+	fn with_recorder_and_cache<RE>(
 		&self,
 		storage_root: Option<H::Out>,
 		callback: impl FnOnce(
 			Option<&mut dyn TrieRecorder<H::Out>>,
 			Option<&mut dyn TrieCache<NodeCodec<H>>>,
-		) -> R,
-	) -> R {
+		) -> RE,
+	) -> RE {
 		let storage_root = storage_root.unwrap_or_else(|| self.root);
 		let mut cache = self.trie_node_cache.as_ref().map(|c| c.as_trie_db_cache(storage_root));
 		let cache = cache.as_mut().map(|c| c as _);
 
-		#[cfg(feature = "std")]
-		{
-			let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root));
-			let recorder = match recorder.as_mut() {
-				Some(recorder) => Some(recorder as &mut dyn TrieRecorder<H::Out>),
-				None => None,
-			};
-			callback(recorder, cache)
-		}
-
-		#[cfg(not(feature = "std"))]
-		{
-			callback(None, cache)
-		}
+		let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root));
+		let recorder = match recorder.as_mut() {
+			Some(recorder) => Some(recorder as &mut dyn TrieRecorder<H::Out>),
+			None => None,
+		};
+		callback(recorder, cache)
 	}
 
 	/// Call the given closure passing it the recorder and the cache.
@@ -329,15 +323,14 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H>> TrieBackendEs
 	/// the new storage root. This is required to register the changes in the cache
 	/// for the correct storage root. The given `storage_root` corresponds to the root of the "old"
 	/// trie. If the value is not given, `self.root` is used.
-	#[cfg(feature = "std")]
-	fn with_recorder_and_cache_for_storage_root<R>(
+	fn with_recorder_and_cache_for_storage_root<RE>(
 		&self,
 		storage_root: Option<H::Out>,
 		callback: impl FnOnce(
 			Option<&mut dyn TrieRecorder<H::Out>>,
 			Option<&mut dyn TrieCache<NodeCodec<H>>>,
-		) -> (Option<H::Out>, R),
-	) -> R {
+		) -> (Option<H::Out>, RE),
+	) -> RE {
 		let storage_root = storage_root.unwrap_or_else(|| self.root);
 		let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root));
 		let recorder = match recorder.as_mut() {
@@ -361,46 +354,26 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H>> TrieBackendEs
 
 		result
 	}
-
-	#[cfg(not(feature = "std"))]
-	fn with_recorder_and_cache_for_storage_root<R>(
-		&self,
-		_storage_root: Option<H::Out>,
-		callback: impl FnOnce(
-			Option<&mut dyn TrieRecorder<H::Out>>,
-			Option<&mut dyn TrieCache<NodeCodec<H>>>,
-		) -> (Option<H::Out>, R),
-	) -> R {
-		if let Some(local_cache) = self.trie_node_cache.as_ref() {
-			let mut cache = local_cache.as_trie_db_mut_cache();
-
-			let (new_root, r) = callback(None, Some(&mut cache));
-
-			if let Some(new_root) = new_root {
-				local_cache.merge(cache, new_root);
-			}
-
-			r
-		} else {
-			callback(None, None).1
-		}
-	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
-	TrieBackendEssence<S, H, C>
+impl<
+		S: TrieBackendStorage<H>,
+		H: Hasher,
+		C: TrieCacheProvider<H> + Send + Sync,
+		R: TrieRecorderProvider<H> + Send + Sync,
+	> TrieBackendEssence<S, H, C, R>
 where
 	H::Out: Codec + Ord,
 {
 	/// Calls the given closure with a [`TrieDb`] constructed for the given
 	/// storage root and (optionally) child trie.
 	#[inline]
-	fn with_trie_db<R>(
+	fn with_trie_db<RE>(
 		&self,
 		root: H::Out,
 		child_info: Option<&ChildInfo>,
-		callback: impl FnOnce(&sp_trie::TrieDB<Layout<H>>) -> R,
-	) -> R {
+		callback: impl FnOnce(&sp_trie::TrieDB<Layout<H>>) -> RE,
+	) -> RE {
 		let backend = self as &dyn HashDBRef<H, Vec<u8>>;
 		let db = child_info
 			.as_ref()
@@ -609,7 +582,7 @@ where
 	}
 
 	/// Create a raw iterator over the storage.
-	pub fn raw_iter(&self, args: IterArgs) -> Result<RawIter<S, H, C>> {
+	pub fn raw_iter(&self, args: IterArgs) -> Result<RawIter<S, H, C, R>> {
 		let root = if let Some(child_info) = args.child_info.as_ref() {
 			let root = match self.child_root(&child_info)? {
 				Some(root) => root,
@@ -831,19 +804,28 @@ where
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
-	AsHashDB<H, DBValue> for TrieBackendEssence<S, H, C>
+impl<
+		S: TrieBackendStorage<H>,
+		H: Hasher,
+		C: TrieCacheProvider<H> + Send + Sync,
+		R: TrieRecorderProvider<H> + Send + Sync,
+	> AsHashDB<H, DBValue> for TrieBackendEssence<S, H, C, R>
 {
 	fn as_hash_db<'b>(&'b self) -> &'b (dyn HashDB<H, DBValue> + 'b) {
 		self
 	}
+
 	fn as_hash_db_mut<'b>(&'b mut self) -> &'b mut (dyn HashDB<H, DBValue> + 'b) {
 		self
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync> HashDB<H, DBValue>
-	for TrieBackendEssence<S, H, C>
+impl<
+		S: TrieBackendStorage<H>,
+		H: Hasher,
+		C: TrieCacheProvider<H> + Send + Sync,
+		R: TrieRecorderProvider<H> + Send + Sync,
+	> HashDB<H, DBValue> for TrieBackendEssence<S, H, C, R>
 {
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Option<DBValue> {
 		if *key == self.empty {
@@ -875,8 +857,12 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
-	HashDBRef<H, DBValue> for TrieBackendEssence<S, H, C>
+impl<
+		S: TrieBackendStorage<H>,
+		H: Hasher,
+		C: TrieCacheProvider<H> + Send + Sync,
+		R: TrieRecorderProvider<H> + Send + Sync,
+	> HashDBRef<H, DBValue> for TrieBackendEssence<S, H, C, R>
 {
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Option<DBValue> {
 		HashDB::get(self, key, prefix)
@@ -928,7 +914,10 @@ mod test {
 				.expect("insert failed");
 		};
 
-		let essence_1 = TrieBackendEssence::<_, _, LocalTrieCache<_>>::new(mdb, root_1);
+		let essence_1 =
+			TrieBackendEssence::<_, _, LocalTrieCache<_>, sp_trie::recorder::Recorder<_>>::new(
+				mdb, root_1,
+			);
 		let mdb = essence_1.backend_storage().clone();
 		let essence_1 = TrieBackend::from_essence(essence_1);
 
@@ -938,7 +927,10 @@ mod test {
 		assert_eq!(essence_1.next_storage_key(b"5"), Ok(Some(b"6".to_vec())));
 		assert_eq!(essence_1.next_storage_key(b"6"), Ok(None));
 
-		let essence_2 = TrieBackendEssence::<_, _, LocalTrieCache<_>>::new(mdb, root_2);
+		let essence_2 =
+			TrieBackendEssence::<_, _, LocalTrieCache<_>, sp_trie::recorder::Recorder<_>>::new(
+				mdb, root_2,
+			);
 
 		assert_eq!(essence_2.next_child_storage_key(child_info, b"2"), Ok(Some(b"3".to_vec())));
 		assert_eq!(essence_2.next_child_storage_key(child_info, b"3"), Ok(Some(b"4".to_vec())));

substrate/primitives/trie/Cargo.toml

@@ -34,6 +34,7 @@ trie-db = { version = "0.28.0", default-features = false }
 trie-root = { version = "0.18.0", default-features = false }
 sp-core = { path = "../core", default-features = false}
 sp-std = { path = "../std", default-features = false}
+sp-externalities = { path = "../externalities", default-features = false }
 schnellru = { version = "0.2.1", optional = true }
 
 [dev-dependencies]
@@ -58,6 +59,7 @@ std = [
 	"scale-info/std",
 	"schnellru",
 	"sp-core/std",
+	"sp-externalities/std",
 	"sp-runtime/std",
 	"sp-std/std",
 	"thiserror",

substrate/primitives/trie/src/lib.rs

@@ -30,6 +30,9 @@ mod storage_proof;
 mod trie_codec;
 mod trie_stream;
 
+#[cfg(feature = "std")]
+pub mod proof_size_extension;
+
 /// Our `NodeCodec`-specific error.
 pub use error::Error;
 /// Various re-exports from the `hash-db` crate.
@@ -146,6 +149,29 @@ where
 	}
 }
 
+/// Type that is able to provide a [`trie_db::TrieRecorder`].
+///
+/// Types implementing this trait can be used to maintain recorded state
+/// across operations on different [`trie_db::TrieDB`] instances.
+pub trait TrieRecorderProvider<H: Hasher> {
+	/// Recorder type that is going to be returned by implementors of this trait.
+	type Recorder<'a>: trie_db::TrieRecorder<H::Out> + 'a
+	where
+		Self: 'a;
+
+	/// Create a [`StorageProof`] derived from the internal state.
+	fn drain_storage_proof(self) -> Option<StorageProof>;
+
+	/// Provide a recorder implementing [`trie_db::TrieRecorder`].
+	fn as_trie_recorder(&self, storage_root: H::Out) -> Self::Recorder<'_>;
+}
+
+/// Type that is able to provide a proof size estimation.
+pub trait ProofSizeProvider {
+	/// Returns the storage proof size.
+	fn estimate_encoded_size(&self) -> usize;
+}
+
 /// TrieDB error over `TrieConfiguration` trait.
 pub type TrieError<L> = trie_db::TrieError<TrieHash<L>, CError<L>>;
 /// Reexport from `hash_db`, with genericity set for `Hasher` trait.

substrate/primitives/trie/src/proof_size_extension.rs

@@ -0,0 +1,39 @@
+// This file is part of Substrate.
+
+// 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.
+
+//! Externalities extension that provides access to the current proof size
+//! of the underlying recorder.
+
+use crate::ProofSizeProvider;
+
+sp_externalities::decl_extension! {
+	/// The proof size extension to fetch the current storage proof size
+	/// in externalities.
+	pub struct ProofSizeExt(Box<dyn ProofSizeProvider + 'static + Sync + Send>);
+}
+
+impl ProofSizeExt {
+	/// Creates a new instance of [`ProofSizeExt`].
+	pub fn new<T: ProofSizeProvider + Sync + Send + 'static>(recorder: T) -> Self {
+		ProofSizeExt(Box::new(recorder))
+	}
+
+	/// Returns the storage proof size.
+	pub fn storage_proof_size(&self) -> u64 {
+		self.0.estimate_encoded_size() as _
+	}
+}

substrate/primitives/trie/src/recorder.rs

@@ -23,7 +23,7 @@
 use crate::{NodeCodec, StorageProof};
 use codec::Encode;
 use hash_db::Hasher;
-use parking_lot::Mutex;
+use parking_lot::{Mutex, MutexGuard};
 use std::{
 	collections::{HashMap, HashSet},
 	marker::PhantomData,
@@ -80,7 +80,9 @@ impl<H> Default for RecorderInner<H> {
 
 /// The trie recorder.
 ///
-/// It can be used to record accesses to the trie and then to convert them into a [`StorageProof`].
+/// Owns the recorded data. Is used to transform data into a storage
+/// proof and to provide transaction support. The `as_trie_recorder` method provides a
+/// [`trie_db::TrieDB`] compatible recorder that implements the actual recording logic.
 pub struct Recorder<H: Hasher> {
 	inner: Arc<Mutex<RecorderInner<H::Out>>>,
 	/// The estimated encoded size of the storage proof this recorder will produce.
@@ -112,11 +114,8 @@ impl<H: Hasher> Recorder<H> {
 	///
 	/// NOTE: This locks a mutex that stays locked until the return value is dropped.
 	#[inline]
-	pub fn as_trie_recorder(
-		&self,
-		storage_root: H::Out,
-	) -> impl trie_db::TrieRecorder<H::Out> + '_ {
-		TrieRecorder::<H, _> {
+	pub fn as_trie_recorder(&self, storage_root: H::Out) -> TrieRecorder<'_, H> {
+		TrieRecorder::<H> {
 			inner: self.inner.lock(),
 			storage_root,
 			encoded_size_estimation: self.encoded_size_estimation.clone(),
@@ -231,15 +230,33 @@ impl<H: Hasher> Recorder<H> {
 	}
 }
 
+impl<H: Hasher> crate::ProofSizeProvider for Recorder<H> {
+	fn estimate_encoded_size(&self) -> usize {
+		Recorder::estimate_encoded_size(self)
+	}
+}
+
 /// The [`TrieRecorder`](trie_db::TrieRecorder) implementation.
-struct TrieRecorder<H: Hasher, I> {
-	inner: I,
+pub struct TrieRecorder<'a, H: Hasher> {
+	inner: MutexGuard<'a, RecorderInner<H::Out>>,
 	storage_root: H::Out,
 	encoded_size_estimation: Arc<AtomicUsize>,
 	_phantom: PhantomData<H>,
 }
 
-impl<H: Hasher, I: DerefMut<Target = RecorderInner<H::Out>>> TrieRecorder<H, I> {
+impl<H: Hasher> crate::TrieRecorderProvider<H> for Recorder<H> {
+	type Recorder<'a> = TrieRecorder<'a, H> where H: 'a;
+
+	fn drain_storage_proof(self) -> Option<StorageProof> {
+		Some(Recorder::drain_storage_proof(self))
+	}
+
+	fn as_trie_recorder(&self, storage_root: H::Out) -> Self::Recorder<'_> {
+		Recorder::as_trie_recorder(&self, storage_root)
+	}
+}
+
+impl<'a, H: Hasher> TrieRecorder<'a, H> {
 	/// Update the recorded keys entry for the given `full_key`.
 	fn update_recorded_keys(&mut self, full_key: &[u8], access: RecordedForKey) {
 		let inner = self.inner.deref_mut();
@@ -283,9 +300,7 @@ impl<H: Hasher, I: DerefMut<Target = RecorderInner<H::Out>>> TrieRecorder<H, I>
 	}
 }
 
-impl<H: Hasher, I: DerefMut<Target = RecorderInner<H::Out>>> trie_db::TrieRecorder<H::Out>
-	for TrieRecorder<H, I>
-{
+impl<'a, H: Hasher> trie_db::TrieRecorder<H::Out> for TrieRecorder<'a, H> {
 	fn record(&mut self, access: TrieAccess<H::Out>) {
 		let mut encoded_size_update = 0;