Make SharedTrieCache/LocalTrieCache work with entire state in memory (#7682)

Part of
#6131 (comment),
we need to make sure that the TrieCache is able to work with the
assumption that more or less the entire current state fits into memory,
this can be split into a few parts:

- [x] Remove the LocalTrieCache hard coded size and make it configurable
from outside.
- [x] Decided to have two flavours of building the LocalTrieCache, for
the import and block authoring we are using a trusted local cache that
grow to hold everything in the block and then propagated everything into
the shared trie cache.
- ~[x] Everything from LocalTrieCache needs to be propagated back to the
SharedTrieCache in a fast and safe manner, so that it doesn't become a
bottleneck. Currently, that is done on drop, with the function holding
the ShareTrieCache write lock while promoting all keys. Decided to move
this on a separate worker thread, so that the hot-path does not have to
wait for drop function to propagate everything from the LocalTrieCache
to the SharedTrieCache, the flushing of the worker thread happens after
the block production and import happens.~ Update this is not need
because authoring and imports is bounded, so the numbers is not
exaggeratedly big, to make the drop too heavy, see numbers here
#7682 (comment).

# Todo
- [x] Stats to prove most of reads/writes hit the shared or the local
trie cache.
- [x] Unit testing  the new changes.

Fixes: #7534

---------

Signed-off-by: Alexandru Gheorghe <alexandru.gheorghe@parity.io>
Co-authored-by: Bastian Köcher <git@kchr.de>

Author: alexggh

Cargo.lock

@@ -43,7 +43,7 @@ use polkadot_service::{
 use sc_cli::{RuntimeVersion, SubstrateCli};
 use sc_client_api::{
 	blockchain::BlockStatus, Backend, BlockchainEvents, HeaderBackend, ImportNotifications,
-	StorageProof,
+	StorageProof, TrieCacheContext,
 };
 use sc_network::{
 	config::NetworkBackendType,
@@ -225,7 +225,7 @@ impl RelayChainInterface for RelayChainInProcessInterface {
 		relay_parent: PHash,
 		key: &[u8],
 	) -> RelayChainResult<Option<StorageValue>> {
-		let state = self.backend.state_at(relay_parent)?;
+		let state = self.backend.state_at(relay_parent, TrieCacheContext::Untrusted)?;
 		state.storage(key).map_err(RelayChainError::GenericError)
 	}
 
@@ -234,7 +234,7 @@ impl RelayChainInterface for RelayChainInProcessInterface {
 		relay_parent: PHash,
 		relevant_keys: &Vec<Vec<u8>>,
 	) -> RelayChainResult<StorageProof> {
-		let state_backend = self.backend.state_at(relay_parent)?;
+		let state_backend = self.backend.state_at(relay_parent, TrieCacheContext::Untrusted)?;
 
 		sp_state_machine::prove_read(state_backend, relevant_keys)
 			.map_err(RelayChainError::StateMachineError)

cumulus/client/relay-chain-inprocess-interface/src/lib.rs

@@ -198,11 +198,11 @@ mod tests {
 			let recorder_for_test: SizeOnlyRecorderProvider<sp_core::Blake2Hasher> =
 				SizeOnlyRecorderProvider::default();
 			let reference_cache: SharedTrieCache<sp_core::Blake2Hasher> =
-				SharedTrieCache::new(CacheSize::new(1024 * 5));
+				SharedTrieCache::new(CacheSize::new(1024 * 5), None);
 			let cache_for_test: SharedTrieCache<sp_core::Blake2Hasher> =
-				SharedTrieCache::new(CacheSize::new(1024 * 5));
+				SharedTrieCache::new(CacheSize::new(1024 * 5), None);
 			{
-				let local_cache = cache_for_test.local_cache();
+				let local_cache = cache_for_test.local_cache_untrusted();
 				let mut trie_cache_for_reference = local_cache.as_trie_db_cache(root);
 				let mut reference_trie_recorder = reference_recorder.as_trie_recorder(root);
 				let reference_trie =
@@ -211,7 +211,7 @@ mod tests {
 						.with_cache(&mut trie_cache_for_reference)
 						.build();
 
-				let local_cache_for_test = reference_cache.local_cache();
+				let local_cache_for_test = reference_cache.local_cache_untrusted();
 				let mut trie_cache_for_test = local_cache_for_test.as_trie_db_cache(root);
 				let mut trie_recorder_under_test = recorder_for_test.as_trie_recorder(root);
 				let test_trie =

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

@@ -0,0 +1,32 @@
+title: Make SharedTrieCache/LocalTrieCache work with entire state in memory
+doc:
+- audience: Node Dev
+  description: |-
+    Extended LocalTrieCache with a trusted configuration that expands to hold everything in memory and 
+    then propagates all the data back to the SharedTrieCache. This new configuration is used on authoring 
+    and import paths.
+crates:
+- name: sp-state-machine
+  bump: major
+- name: sp-trie
+  bump: major
+- name: sc-client-api
+  bump: major
+- name: sc-client-db
+  bump: major
+- name: cumulus-pallet-parachain-system
+  bump: major
+- name: sc-cli
+  bump: major
+- name: sc-service
+  bump: major
+- name: frame-benchmarking-cli 
+  bump: major
+- name: substrate-state-trie-migration-rpc
+  bump: major
+- name: sc-consensus-beefy
+  bump: major
+- name: sc-basic-authorship
+  bump: major
+- name: cumulus-relay-chain-inprocess-interface
+  bump: major

prdoc/pr_7682.prdoc

@@ -389,6 +389,7 @@ impl BenchDb {
 			state_pruning: Some(PruningMode::ArchiveAll),
 			source: database_type.into_settings(dir.into()),
 			blocks_pruning: sc_client_db::BlocksPruning::KeepAll,
+			metrics_registry: None,
 		};
 		let task_executor = TaskExecutor::new();
 

substrate/bin/node/testing/src/bench.rs

@@ -22,6 +22,7 @@ use std::collections::HashSet;
 
 use parking_lot::RwLock;
 
+use sp_api::CallContext;
 use sp_consensus::BlockOrigin;
 use sp_core::offchain::OffchainStorage;
 use sp_runtime::{
@@ -492,6 +493,34 @@ pub trait StorageProvider<Block: BlockT, B: Backend<Block>> {
 	) -> sp_blockchain::Result<Option<MerkleValue<Block::Hash>>>;
 }
 
+/// Specify the desired trie cache context when calling [`Backend::state_at`].
+///
+/// This is used to determine the size of the local trie cache.
+#[derive(Debug, Clone, Copy)]
+pub enum TrieCacheContext {
+	/// This is used when calling [`Backend::state_at`] in a trusted context.
+	///
+	/// A trusted context is for example the building or importing of a block.
+	/// In this case the local trie cache can grow unlimited and all the cached data
+	/// will be propagated back to the shared trie cache. It is safe to let the local
+	/// cache grow to hold the entire data, because importing and building blocks is
+	/// bounded by the block size limit.
+	Trusted,
+	/// This is used when calling [`Backend::state_at`] in from untrusted context.
+	///
+	/// The local trie cache will be bounded by its preconfigured size.
+	Untrusted,
+}
+
+impl From<CallContext> for TrieCacheContext {
+	fn from(call_context: CallContext) -> Self {
+		match call_context {
+			CallContext::Onchain => TrieCacheContext::Trusted,
+			CallContext::Offchain => TrieCacheContext::Untrusted,
+		}
+	}
+}
+
 /// Client backend.
 ///
 /// Manages the data layer.
@@ -584,11 +613,15 @@ pub trait Backend<Block: BlockT>: AuxStore + Send + Sync {
 
 	/// Returns true if state for given block is available.
 	fn have_state_at(&self, hash: Block::Hash, _number: NumberFor<Block>) -> bool {
-		self.state_at(hash).is_ok()
+		self.state_at(hash, TrieCacheContext::Untrusted).is_ok()
 	}
 
 	/// Returns state backend with post-state of given block.
-	fn state_at(&self, hash: Block::Hash) -> sp_blockchain::Result<Self::State>;
+	fn state_at(
+		&self,
+		hash: Block::Hash,
+		trie_cache_context: TrieCacheContext,
+	) -> sp_blockchain::Result<Self::State>;
 
 	/// Attempts to revert the chain by `n` blocks. If `revert_finalized` is set it will attempt to
 	/// revert past any finalized block, this is unsafe and can potentially leave the node in an

substrate/client/api/src/backend.rs

@@ -42,7 +42,7 @@ use crate::{
 	backend::{self, NewBlockState},
 	blockchain::{self, BlockStatus, HeaderBackend},
 	leaves::LeafSet,
-	UsageInfo,
+	TrieCacheContext, UsageInfo,
 };
 
 struct PendingBlock<B: BlockT> {
@@ -652,7 +652,7 @@ impl<Block: BlockT> backend::Backend<Block> for Backend<Block> {
 	type OffchainStorage = OffchainStorage;
 
 	fn begin_operation(&self) -> sp_blockchain::Result<Self::BlockImportOperation> {
-		let old_state = self.state_at(Default::default())?;
+		let old_state = self.state_at(Default::default(), TrieCacheContext::Untrusted)?;
 		Ok(BlockImportOperation {
 			pending_block: None,
 			old_state,
@@ -668,7 +668,7 @@ impl<Block: BlockT> backend::Backend<Block> for Backend<Block> {
 		operation: &mut Self::BlockImportOperation,
 		block: Block::Hash,
 	) -> sp_blockchain::Result<()> {
-		operation.old_state = self.state_at(block)?;
+		operation.old_state = self.state_at(block, TrieCacheContext::Untrusted)?;
 		Ok(())
 	}
 
@@ -734,7 +734,11 @@ impl<Block: BlockT> backend::Backend<Block> for Backend<Block> {
 		None
 	}
 
-	fn state_at(&self, hash: Block::Hash) -> sp_blockchain::Result<Self::State> {
+	fn state_at(
+		&self,
+		hash: Block::Hash,
+		_trie_cache_context: TrieCacheContext,
+	) -> sp_blockchain::Result<Self::State> {
 		if hash == Default::default() {
 			return Ok(Self::State::default())
 		}

substrate/client/api/src/in_mem.rs

@@ -606,7 +606,7 @@ mod tests {
 
 	use futures::executor::block_on;
 	use parking_lot::Mutex;
-	use sc_client_api::Backend;
+	use sc_client_api::{Backend, TrieCacheContext};
 	use sc_transaction_pool::BasicPool;
 	use sc_transaction_pool_api::{ChainEvent, MaintainedTransactionPool, TransactionSource};
 	use sp_api::Core;
@@ -777,7 +777,7 @@ mod tests {
 		let api = client.runtime_api();
 		api.execute_block(genesis_hash, proposal.block).unwrap();
 
-		let state = backend.state_at(genesis_hash).unwrap();
+		let state = backend.state_at(genesis_hash, TrieCacheContext::Untrusted).unwrap();
 
 		let storage_changes = api.into_storage_changes(&state, genesis_hash).unwrap();
 

substrate/client/basic-authorship/src/basic_authorship.rs

@@ -77,6 +77,7 @@ impl ChainInfoCmd {
 			state_pruning: config.state_pruning.clone(),
 			source: config.database.clone(),
 			blocks_pruning: config.blocks_pruning,
+			metrics_registry: None,
 		};
 		let backend = sc_service::new_db_backend::<B>(db_config)?;
 		let info: ChainInfo<B> = backend.blockchain().info().into();

substrate/client/cli/src/commands/chain_info_cmd.rs

@@ -28,7 +28,7 @@ use sp_runtime::{
 	EncodedJustification,
 };
 
-use sc_client_api::backend::Backend;
+use sc_client_api::{backend::Backend, TrieCacheContext};
 use sc_consensus::{BlockCheckParams, BlockImport, BlockImportParams, ImportResult};
 
 use crate::{
@@ -149,7 +149,7 @@ where
 		// Run inner block import.
 		let inner_import_result = self.inner.import_block(block).await?;
 
-		match self.backend.state_at(hash) {
+		match self.backend.state_at(hash, TrieCacheContext::Untrusted) {
 			Ok(_) => {},
 			Err(_) => {
 				// The block is imported as part of some chain sync.