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/client/db/src/bench.rs

@@ -18,13 +18,7 @@
 
 //! State backend that's useful for benchmarking
 
-use std::{
-	cell::{Cell, RefCell},
-	collections::HashMap,
-	sync::Arc,
-};
-
-use crate::storage_cache::{new_shared_cache, CachingState, SharedCache};
+use crate::{DbState, DbStateBuilder};
 use hash_db::{Hasher, Prefix};
 use kvdb::{DBTransaction, KeyValueDB};
 use linked_hash_map::LinkedHashMap;
@@ -37,40 +31,31 @@ use sp_runtime::{
 	StateVersion, Storage,
 };
 use sp_state_machine::{
-	backend::Backend as StateBackend, ChildStorageCollection, DBValue, ProofRecorder,
-	StorageCollection,
+	backend::Backend as StateBackend, ChildStorageCollection, DBValue, StorageCollection,
+};
+use sp_trie::{
+	cache::{CacheSize, SharedTrieCache},
+	prefixed_key, MemoryDB,
+};
+use std::{
+	cell::{Cell, RefCell},
+	collections::HashMap,
+	sync::Arc,
 };
-use sp_trie::{prefixed_key, MemoryDB};
 
-type DbState<B> =
-	sp_state_machine::TrieBackend<Arc<dyn sp_state_machine::Storage<HashFor<B>>>, HashFor<B>>;
-
-type State<B> = CachingState<DbState<B>, B>;
+type State<B> = DbState<B>;
 
 struct StorageDb<Block: BlockT> {
 	db: Arc<dyn KeyValueDB>,
-	proof_recorder: Option<ProofRecorder<Block::Hash>>,
 	_block: std::marker::PhantomData<Block>,
 }
 
 impl<Block: BlockT> sp_state_machine::Storage<HashFor<Block>> for StorageDb<Block> {
 	fn get(&self, key: &Block::Hash, prefix: Prefix) -> Result<Option<DBValue>, String> {
 		let prefixed_key = prefixed_key::<HashFor<Block>>(key, prefix);
-		if let Some(recorder) = &self.proof_recorder {
-			if let Some(v) = recorder.get(key) {
-				return Ok(v)
-			}
-			let backend_value = self
-				.db
-				.get(0, &prefixed_key)
-				.map_err(|e| format!("Database backend error: {:?}", e))?;
-			recorder.record(*key, backend_value.clone());
-			Ok(backend_value)
-		} else {
-			self.db
-				.get(0, &prefixed_key)
-				.map_err(|e| format!("Database backend error: {:?}", e))
-		}
+		self.db
+			.get(0, &prefixed_key)
+			.map_err(|e| format!("Database backend error: {:?}", e))
 	}
 }
 
@@ -82,7 +67,6 @@ pub struct BenchmarkingState<B: BlockT> {
 	db: Cell<Option<Arc<dyn KeyValueDB>>>,
 	genesis: HashMap<Vec<u8>, (Vec<u8>, i32)>,
 	record: Cell<Vec<Vec<u8>>>,
-	shared_cache: SharedCache<B>, // shared cache is always empty
 	/// Key tracker for keys in the main trie.
 	/// We track the total number of reads and writes to these keys,
 	/// not de-duplicated for repeats.
@@ -93,9 +77,10 @@ pub struct BenchmarkingState<B: BlockT> {
 	/// not de-duplicated for repeats.
 	child_key_tracker: RefCell<LinkedHashMap<Vec<u8>, LinkedHashMap<Vec<u8>, TrackedStorageKey>>>,
 	whitelist: RefCell<Vec<TrackedStorageKey>>,
-	proof_recorder: Option<ProofRecorder<B::Hash>>,
+	proof_recorder: Option<sp_trie::recorder::Recorder<HashFor<B>>>,
 	proof_recorder_root: Cell<B::Hash>,
 	enable_tracking: bool,
+	shared_trie_cache: SharedTrieCache<HashFor<B>>,
 }
 
 impl<B: BlockT> BenchmarkingState<B> {
@@ -109,7 +94,7 @@ impl<B: BlockT> BenchmarkingState<B> {
 		let state_version = sp_runtime::StateVersion::default();
 		let mut root = B::Hash::default();
 		let mut mdb = MemoryDB::<HashFor<B>>::default();
-		sp_state_machine::TrieDBMutV1::<HashFor<B>>::new(&mut mdb, &mut root);
+		sp_trie::trie_types::TrieDBMutBuilderV1::<HashFor<B>>::new(&mut mdb, &mut root).build();
 
 		let mut state = BenchmarkingState {
 			state: RefCell::new(None),
@@ -118,13 +103,14 @@ impl<B: BlockT> BenchmarkingState<B> {
 			genesis: Default::default(),
 			genesis_root: Default::default(),
 			record: Default::default(),
-			shared_cache: new_shared_cache(0, (1, 10)),
 			main_key_tracker: Default::default(),
 			child_key_tracker: Default::default(),
 			whitelist: Default::default(),
 			proof_recorder: record_proof.then(Default::default),
 			proof_recorder_root: Cell::new(root),
 			enable_tracking,
+			// Enable the cache, but do not sync anything to the shared state.
+			shared_trie_cache: SharedTrieCache::new(CacheSize::Maximum(0)),
 		};
 
 		state.add_whitelist_to_tracker();
@@ -160,16 +146,13 @@ impl<B: BlockT> BenchmarkingState<B> {
 			recorder.reset();
 			self.proof_recorder_root.set(self.root.get());
 		}
-		let storage_db = Arc::new(StorageDb::<B> {
-			db,
-			proof_recorder: self.proof_recorder.clone(),
-			_block: Default::default(),
-		});
-		*self.state.borrow_mut() = Some(State::new(
-			DbState::<B>::new(storage_db, self.root.get()),
-			self.shared_cache.clone(),
-			None,
-		));
+		let storage_db = Arc::new(StorageDb::<B> { db, _block: Default::default() });
+		*self.state.borrow_mut() = Some(
+			DbStateBuilder::<B>::new(storage_db, self.root.get())
+				.with_optional_recorder(self.proof_recorder.clone())
+				.with_cache(self.shared_trie_cache.local_cache())
+				.build(),
+		);
 		Ok(())
 	}
 
@@ -324,6 +307,19 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 			.child_storage(child_info, key)
 	}
 
+	fn child_storage_hash(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<Option<B::Hash>, Self::Error> {
+		self.add_read_key(Some(child_info.storage_key()), key);
+		self.state
+			.borrow()
+			.as_ref()
+			.ok_or_else(state_err)?
+			.child_storage_hash(child_info, key)
+	}
+
 	fn exists_storage(&self, key: &[u8]) -> Result<bool, Self::Error> {
 		self.add_read_key(None, key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.exists_storage(key)
@@ -604,22 +600,25 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 	fn proof_size(&self) -> Option<u32> {
 		self.proof_recorder.as_ref().map(|recorder| {
 			let proof_size = recorder.estimate_encoded_size() as u32;
+
 			let proof = recorder.to_storage_proof();
+
 			let proof_recorder_root = self.proof_recorder_root.get();
 			if proof_recorder_root == Default::default() || proof_size == 1 {
 				// empty trie
 				proof_size
-			} else if let Some(size) = proof.encoded_compact_size::<HashFor<B>>(proof_recorder_root)
-			{
-				size as u32
 			} else {
-				panic!(
-					"proof rec root {:?}, root {:?}, genesis {:?}, rec_len {:?}",
-					self.proof_recorder_root.get(),
-					self.root.get(),
-					self.genesis_root,
-					proof_size,
-				);
+				if let Some(size) = proof.encoded_compact_size::<HashFor<B>>(proof_recorder_root) {
+					size as u32
+				} else {
+					panic!(
+						"proof rec root {:?}, root {:?}, genesis {:?}, rec_len {:?}",
+						self.proof_recorder_root.get(),
+						self.root.get(),
+						self.genesis_root,
+						proof_size,
+					);
+				}
 			}
 		})
 	}