Make SharedTrieCache/LocalTrieCache work with entire state in memory

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

@@ -40,7 +40,7 @@ use polkadot_service::{
 use sc_cli::{RuntimeVersion, SubstrateCli};
 use sc_client_api::{
 	blockchain::BlockStatus, Backend, BlockchainEvents, HeaderBackend, ImportNotifications,
-	StorageProof,
+	StorageProof, TrieCacheContext,
 };
 use sc_telemetry::TelemetryWorkerHandle;
 use sp_api::{CallApiAt, CallApiAtParams, CallContext, ProvideRuntimeApi};
@@ -217,7 +217,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)
 	}
 
@@ -226,7 +226,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/pallets/parachain-system/src/validate_block/trie_recorder.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 =

prdoc/pr_7682.prdoc

@@ -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

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

@@ -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/client/api/src/backend.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,32 @@ 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.
+	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 +611,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/in_mem.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/basic-authorship/src/basic_authorship.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/cli/src/commands/chain_info_cmd.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/consensus/beefy/src/import.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.

substrate/client/db/Cargo.toml

@@ -32,6 +32,7 @@ linked-hash-map = { workspace = true }
 log = { workspace = true, default-features = true }
 parity-db = { workspace = true }
 parking_lot = { workspace = true, default-features = true }
+prometheus-endpoint = { workspace = true, default-features = true }
 sc-client-api = { workspace = true, default-features = true }
 sc-state-db = { workspace = true, default-features = true }
 schnellru = { workspace = true }

substrate/client/db/benches/state_access.rs

@@ -18,7 +18,9 @@
 
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use rand::{distributions::Uniform, rngs::StdRng, Rng, SeedableRng};
-use sc_client_api::{Backend as _, BlockImportOperation, NewBlockState, StateBackend};
+use sc_client_api::{
+	Backend as _, BlockImportOperation, NewBlockState, StateBackend, TrieCacheContext,
+};
 use sc_client_db::{Backend, BlocksPruning, DatabaseSettings, DatabaseSource, PruningMode};
 use sp_core::H256;
 use sp_runtime::{
@@ -83,10 +85,11 @@ fn insert_blocks(db: &Backend<Block>, storage: Vec<(Vec<u8>, Vec<u8>)>) -> H256
 			.map(|(k, v)| (k.clone(), Some(v.clone())))
 			.collect::<Vec<_>>();
 
-		let (state_root, tx) = db.state_at(parent_hash).unwrap().storage_root(
-			changes.iter().map(|(k, v)| (k.as_slice(), v.as_deref())),
-			StateVersion::V1,
-		);
+		let (state_root, tx) =
+			db.state_at(parent_hash, TrieCacheContext::Trusted).unwrap().storage_root(
+				changes.iter().map(|(k, v)| (k.as_slice(), v.as_deref())),
+				StateVersion::V1,
+			);
 		header.state_root = state_root;
 
 		op.update_db_storage(tx).unwrap();
@@ -122,6 +125,7 @@ fn create_backend(config: BenchmarkConfig, temp_dir: &TempDir) -> Backend<Block>
 		state_pruning: Some(PruningMode::ArchiveAll),
 		source: DatabaseSource::ParityDb { path },
 		blocks_pruning: BlocksPruning::KeepAll,
+		metrics_registry: None,
 	};
 
 	Backend::new(settings, 100).expect("Creates backend")
@@ -175,7 +179,7 @@ fn state_access_benchmarks(c: &mut Criterion) {
 
 		group.bench_function(desc, |b| {
 			b.iter_batched(
-				|| backend.state_at(block_hash).expect("Creates state"),
+				|| backend.state_at(block_hash, TrieCacheContext::Trusted).expect("Creates state"),
 				|state| {
 					for key in keys.iter().cycle().take(keys.len() * multiplier) {
 						let _ = state.storage(&key).expect("Doesn't fail").unwrap();
@@ -213,7 +217,7 @@ fn state_access_benchmarks(c: &mut Criterion) {
 
 		group.bench_function(desc, |b| {
 			b.iter_batched(
-				|| backend.state_at(block_hash).expect("Creates state"),
+				|| backend.state_at(block_hash, TrieCacheContext::Trusted).expect("Creates state"),
 				|state| {
 					for key in keys.iter().take(1).cycle().take(multiplier) {
 						let _ = state.storage(&key).expect("Doesn't fail").unwrap();
@@ -251,7 +255,7 @@ fn state_access_benchmarks(c: &mut Criterion) {
 
 		group.bench_function(desc, |b| {
 			b.iter_batched(
-				|| backend.state_at(block_hash).expect("Creates state"),
+				|| backend.state_at(block_hash, TrieCacheContext::Trusted).expect("Creates state"),
 				|state| {
 					for key in keys.iter().take(1).cycle().take(multiplier) {
 						let _ = state.storage_hash(&key).expect("Doesn't fail").unwrap();
@@ -289,7 +293,7 @@ fn state_access_benchmarks(c: &mut Criterion) {
 
 		group.bench_function(desc, |b| {
 			b.iter_batched(
-				|| backend.state_at(block_hash).expect("Creates state"),
+				|| backend.state_at(block_hash, TrieCacheContext::Trusted).expect("Creates state"),
 				|state| {
 					let _ = state
 						.storage_hash(sp_core::storage::well_known_keys::CODE)

substrate/client/db/src/bench.rs

@@ -154,7 +154,7 @@ impl<Hasher: Hash> BenchmarkingState<Hasher> {
 			proof_recorder: record_proof.then(Default::default),
 			proof_recorder_root: Cell::new(root),
 			// Enable the cache, but do not sync anything to the shared state.
-			shared_trie_cache: SharedTrieCache::new(CacheSize::new(0)),
+			shared_trie_cache: SharedTrieCache::new(CacheSize::new(0), None),
 		};
 
 		state.add_whitelist_to_tracker();
@@ -199,7 +199,7 @@ impl<Hasher: Hash> BenchmarkingState<Hasher> {
 		*self.state.borrow_mut() = Some(
 			DbStateBuilder::<Hasher>::new(storage_db, self.root.get())
 				.with_optional_recorder(self.proof_recorder.clone())
-				.with_cache(self.shared_trie_cache.local_cache())
+				.with_cache(self.shared_trie_cache.local_cache_untrusted())
 				.build(),
 		);
 		Ok(())