Introduce trie level cache and remove state cache (#11407)

* trie state cache

* Also cache missing access on read.

* fix comp

* bis

* fix

* use has_lru

* remove local storage cache on size 0.

* No cache.

* local cache only

* trie cache and local cache

* storage cache (with local)

* trie cache no local cache

* Add state access benchmark

* Remove warnings etc

* Add trie cache benchmark

* No extra "clone" required

* Change benchmark to use multiple blocks

* Use patches

* Integrate shitty implementation

* More stuff

* Revert "Merge branch 'master' into trie_state_cache"

This reverts commit 947cd8e6d43fced10e21b76d5b92ffa57b57c318, reversing
changes made to 29ff036463.

* Improve benchmark

* Adapt to latest changes

* Adapt to changes in trie

* Add a test that uses iterator

* Start fixing it

* Remove obsolete file

* Make it compile

* Start rewriting the trie node cache

* More work on the cache

* More docs and code etc

* Make data cache an optional

* Tests

* Remove debug stuff

* Recorder

* Some docs and a simple test for the recorder

* Compile fixes

* Make it compile

* More fixes

* More fixes

* Fix fix fix

* Make sure cache and recorder work together for basic stuff

* Test that data caching and recording works

* Test `TrieDBMut` with caching

* Try something

* Fixes, fixes, fixes

* Forward the recorder

* Make it compile

* Use recorder in more places

* Switch to new `with_optional_recorder` fn

* Refactor and cleanups

* Move `ProvingBackend` tests

* Simplify

* Move over all functionality to the essence

* Fix compilation

* Implement estimate encoded size for StorageProof

* Start using the `cache` everywhere

* Use the cache everywhere

* Fix compilation

* Fix tests

* Adds `TrieBackendBuilder` and enhances the tests

* Ensure that recorder drain checks that values are found as expected

* Switch over to `TrieBackendBuilder`

* Start fixing the problem with child tries and recording

* Fix recording of child tries

* Make it compile

* Overwrite `storage_hash` in `TrieBackend`

* Add `storage_cache` to  the benchmarks

* Fix `no_std` build

* Speed up cache lookup

* Extend the state access benchmark to also hash a runtime

* Fix build

* Fix compilation

* Rewrite value cache

* Add lru cache

* Ensure that the cache lru works

* Value cache should not be optional

* Add support for keeping the shared node cache in its bounds

* Make the cache configurable

* Check that the cache respects the bounds

* Adds a new test

* Fixes

* Docs and some renamings

* More docs

* Start using the new recorder

* Fix more code

* Take `self` argument

* Remove warnings

* Fix benchmark

* Fix accounting

* Rip off the state cache

* Start fixing fallout after removing the state cache

* Make it compile after trie changes

* Fix test

* Add some logging

* Some docs

* Some fixups and clean ups

* Fix benchmark

* Remove unneeded file

* Use git for patching

* Make CI happy

* Update primitives/trie/Cargo.toml

Co-authored-by: Koute <koute@users.noreply.github.com>

* Update primitives/state-machine/src/trie_backend.rs

Co-authored-by: cheme <emericchevalier.pro@gmail.com>

* Introduce new `AsTrieBackend` trait

* Make the LocalTrieCache not clonable

* Make it work in no_std and add docs

* Remove duplicate dependency

* Switch to ahash for better performance

* Speedup value cache merge

* Output errors on underflow

* Ensure the internal LRU map doesn't grow too much

* Use const fn to calculate the value cache element size

* Remove cache configuration

* Fix

* Clear the cache in between for more testing

* Try to come up with a failing test case

* Make the test fail

* Fix the child trie recording

* Make everything compile after the changes to trie

* Adapt to latest trie-db changes

* Fix on stable

* Update primitives/trie/src/cache.rs

Co-authored-by: cheme <emericchevalier.pro@gmail.com>

* Fix wrong merge

* Docs

* Fix warnings

* Cargo.lock

* Bump pin-project

* Fix warnings

* Switch to released crate version

* More fixes

* Make clippy and rustdocs happy

* More clippy

* Print error when using deprecated `--state-cache-size`

* 🤦

* Fixes

* Fix storage_hash linkings

* Update client/rpc/src/dev/mod.rs

Co-authored-by: Arkadiy Paronyan <arkady.paronyan@gmail.com>

* Review feedback

* encode bound

* Rework the shared value cache

Instead of using a `u64` to represent the key we now use an `Arc<[u8]>`. This arc is also stored in
some extra `HashSet`. We store the key are in an extra `HashSet` to de-duplicate the keys accross
different storage roots. When the latest key usage is dropped in the lru, we also remove the key
from the `HashSet`.

* Improve of the cache by merging the old and new solution

* FMT

* Please stop coming back all the time :crying:

* Update primitives/trie/src/cache/shared_cache.rs

Co-authored-by: Arkadiy Paronyan <arkady.paronyan@gmail.com>

* Fixes

* Make clippy happy

* Ensure we don't deadlock

* Only use one lock to simplify the code

* Do not depend on `Hasher`

* Fix tests

* FMT

* Clippy 🤦

Co-authored-by: cheme <emericchevalier.pro@gmail.com>
Co-authored-by: Koute <koute@users.noreply.github.com>
Co-authored-by: Arkadiy Paronyan <arkady.paronyan@gmail.com>

substrate/utils/frame/benchmarking-cli/src/lib.rs

@@ -93,9 +93,9 @@ impl CliConfiguration for BenchmarkCmd {
 		}
 	}
 
-	fn state_cache_size(&self) -> Result<usize> {
+	fn trie_cache_maximum_size(&self) -> Result<Option<usize>> {
 		unwrap_cmd! {
-			self, cmd, cmd.state_cache_size()
+			self, cmd, cmd.trie_cache_maximum_size()
 		}
 	}
 

substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs

@@ -96,9 +96,11 @@ pub struct StorageParams {
 	#[clap(long, possible_values = ["0", "1"])]
 	pub state_version: u8,
 
-	/// State cache size.
-	#[clap(long, default_value = "0")]
-	pub state_cache_size: usize,
+	/// Trie cache size in bytes.
+	///
+	/// Providing `0` will disable the cache.
+	#[clap(long, default_value = "1024")]
+	pub trie_cache_size: usize,
 
 	/// Include child trees in benchmark.
 	#[clap(long)]
@@ -211,7 +213,11 @@ impl CliConfiguration for StorageCmd {
 		Some(&self.pruning_params)
 	}
 
-	fn state_cache_size(&self) -> Result<usize> {
-		Ok(self.params.state_cache_size)
+	fn trie_cache_maximum_size(&self) -> Result<Option<usize>> {
+		if self.params.trie_cache_size == 0 {
+			Ok(None)
+		} else {
+			Ok(Some(self.params.trie_cache_size))
+		}
 	}
 }

substrate/utils/frame/benchmarking-cli/src/storage/write.rs

@@ -17,7 +17,7 @@
 
 use sc_cli::Result;
 use sc_client_api::{Backend as ClientBackend, StorageProvider, UsageProvider};
-use sc_client_db::{DbHash, DbState};
+use sc_client_db::{DbHash, DbState, DbStateBuilder};
 use sp_api::StateBackend;
 use sp_blockchain::HeaderBackend;
 use sp_database::{ColumnId, Transaction};
@@ -60,7 +60,7 @@ impl StorageCmd {
 		let block = BlockId::Number(client.usage_info().chain.best_number);
 		let header = client.header(block)?.ok_or("Header not found")?;
 		let original_root = *header.state_root();
-		let trie = DbState::<Block>::new(storage.clone(), original_root);
+		let trie = DbStateBuilder::<Block>::new(storage.clone(), original_root).build();
 
 		info!("Preparing keys from block {}", block);
 		// Load all KV pairs and randomly shuffle them.

substrate/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml

@@ -23,7 +23,7 @@ sp-io = { path = "../../../../primitives/io" }
 sp-core = { path = "../../../../primitives/core" }
 sp-state-machine = { path = "../../../../primitives/state-machine" }
 sp-trie = { path = "../../../../primitives/trie" }
-trie-db = { version = "0.23.1" }
+trie-db = "0.24.0"
 
 jsonrpsee = { version = "0.15.1", features = ["server", "macros"] }
 

substrate/utils/frame/rpc/state-trie-migration-rpc/src/lib.rs

@@ -31,8 +31,11 @@ use sp_core::{
 	storage::{ChildInfo, ChildType, PrefixedStorageKey},
 	Hasher,
 };
-use sp_state_machine::Backend;
-use sp_trie::{trie_types::TrieDB, KeySpacedDB, Trie};
+use sp_state_machine::backend::AsTrieBackend;
+use sp_trie::{
+	trie_types::{TrieDB, TrieDBBuilder},
+	KeySpacedDB, Trie,
+};
 use trie_db::{
 	node::{NodePlan, ValuePlan},
 	TrieDBNodeIterator,
@@ -41,9 +44,9 @@ use trie_db::{
 fn count_migrate<'a, H: Hasher>(
 	storage: &'a dyn trie_db::HashDBRef<H, Vec<u8>>,
 	root: &'a H::Out,
-) -> std::result::Result<(u64, TrieDB<'a, H>), String> {
+) -> std::result::Result<(u64, TrieDB<'a, 'a, H>), String> {
 	let mut nb = 0u64;
-	let trie = TrieDB::new(storage, root).map_err(|e| format!("TrieDB creation error: {}", e))?;
+	let trie = TrieDBBuilder::new(storage, root).build();
 	let iter_node =
 		TrieDBNodeIterator::new(&trie).map_err(|e| format!("TrieDB node iterator error: {}", e))?;
 	for node in iter_node {
@@ -68,13 +71,9 @@ pub fn migration_status<H, B>(backend: &B) -> std::result::Result<(u64, u64), St
 where
 	H: Hasher,
 	H::Out: codec::Codec,
-	B: Backend<H>,
+	B: AsTrieBackend<H>,
 {
-	let trie_backend = if let Some(backend) = backend.as_trie_backend() {
-		backend
-	} else {
-		return Err("No access to trie from backend.".to_string())
-	};
+	let trie_backend = backend.as_trie_backend();
 	let essence = trie_backend.essence();
 	let (nb_to_migrate, trie) = count_migrate(essence, essence.root())?;
 

substrate/utils/frame/try-runtime/cli/src/lib.rs

@@ -293,7 +293,7 @@ use sp_runtime::{
 	traits::{Block as BlockT, NumberFor},
 	DeserializeOwned,
 };
-use sp_state_machine::{InMemoryProvingBackend, OverlayedChanges, StateMachine};
+use sp_state_machine::{OverlayedChanges, StateMachine, TrieBackendBuilder};
 use std::{fmt::Debug, path::PathBuf, str::FromStr};
 
 mod commands;
@@ -746,9 +746,11 @@ pub(crate) fn state_machine_call_with_proof<Block: BlockT, D: NativeExecutionDis
 
 	let mut changes = Default::default();
 	let backend = ext.backend.clone();
-	let proving_backend = InMemoryProvingBackend::new(&backend);
+	let runtime_code_backend = sp_state_machine::backend::BackendRuntimeCode::new(&backend);
+
+	let proving_backend =
+		TrieBackendBuilder::wrap(&backend).with_recorder(Default::default()).build();
 
-	let runtime_code_backend = sp_state_machine::backend::BackendRuntimeCode::new(&proving_backend);
 	let runtime_code = runtime_code_backend.runtime_code()?;
 
 	let pre_root = *backend.root();
@@ -767,7 +769,9 @@ pub(crate) fn state_machine_call_with_proof<Block: BlockT, D: NativeExecutionDis
 	.map_err(|e| format!("failed to execute {}: {}", method, e))
 	.map_err::<sc_cli::Error, _>(Into::into)?;
 
-	let proof = proving_backend.extract_proof();
+	let proof = proving_backend
+		.extract_proof()
+		.expect("A recorder was set and thus, a storage proof can be extracted; qed");
 	let proof_size = proof.encoded_size();
 	let compact_proof = proof
 		.clone()