use crate::{Error, NodeCodec};
use hash_db::Hasher;
use nohash_hasher::BuildNoHashHasher;
use parking_lot::{Mutex, MutexGuard};
use schnellru::LruMap;
use shared_cache::{ValueCacheKey, ValueCacheRef};
use std::{
collections::HashMap,
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
time::Duration,
};
use trie_db::{node::NodeOwned, CachedValue};
mod shared_cache;
pub use shared_cache::SharedTrieCache;
use self::shared_cache::ValueCacheKeyHash;
const LOG_TARGET: &str = "trie-cache";
const SHARED_CACHE_WRITE_LOCK_TIMEOUT: Duration = Duration::from_millis(100);
const SHARED_NODE_CACHE_MAX_PROMOTED_KEYS: u32 = 1792;
const SHARED_VALUE_CACHE_MAX_PROMOTED_KEYS: u32 = 1792;
const SHARED_NODE_CACHE_MAX_REPLACE_PERCENT: usize = 33;
const SHARED_VALUE_CACHE_MAX_REPLACE_PERCENT: usize = 33;
const LOCAL_NODE_CACHE_MAX_INLINE_SIZE: usize = 512 * 1024;
const LOCAL_VALUE_CACHE_MAX_INLINE_SIZE: usize = 512 * 1024;
const LOCAL_NODE_CACHE_MAX_HEAP_SIZE: usize = 8 * 1024 * 1024;
const LOCAL_VALUE_CACHE_MAX_HEAP_SIZE: usize = 2 * 1024 * 1024;
#[derive(Debug, Clone, Copy)]
pub struct CacheSize(usize);
impl CacheSize {
pub const fn unlimited() -> Self {
CacheSize(usize::MAX)
}
pub const fn new(bytes: usize) -> Self {
CacheSize(bytes)
}
}
#[derive(Default)]
pub struct LocalNodeCacheLimiter {
current_heap_size: usize,
}
impl<H> schnellru::Limiter<H, NodeCached<H>> for LocalNodeCacheLimiter
where
H: AsRef<[u8]> + std::fmt::Debug,
{
type KeyToInsert<'a> = H;
type LinkType = u32;
#[inline]
fn is_over_the_limit(&self, length: usize) -> bool {
if length <= 1 {
return false
}
self.current_heap_size > LOCAL_NODE_CACHE_MAX_HEAP_SIZE
}
#[inline]
fn on_insert<'a>(
&mut self,
_length: usize,
key: H,
cached_node: NodeCached<H>,
) -> Option<(H, NodeCached<H>)> {
self.current_heap_size += cached_node.heap_size();
Some((key, cached_node))
}
#[inline]
fn on_replace(
&mut self,
_length: usize,
_old_key: &mut H,
_new_key: H,
old_node: &mut NodeCached<H>,
new_node: &mut NodeCached<H>,
) -> bool {
debug_assert_eq!(_old_key.as_ref().len(), _new_key.as_ref().len());
self.current_heap_size =
self.current_heap_size + new_node.heap_size() - old_node.heap_size();
true
}
#[inline]
fn on_removed(&mut self, _key: &mut H, cached_node: &mut NodeCached<H>) {
self.current_heap_size -= cached_node.heap_size();
}
#[inline]
fn on_cleared(&mut self) {
self.current_heap_size = 0;
}
#[inline]
fn on_grow(&mut self, new_memory_usage: usize) -> bool {
new_memory_usage <= LOCAL_NODE_CACHE_MAX_INLINE_SIZE
}
}
#[derive(Default)]
pub struct LocalValueCacheLimiter {
current_heap_size: usize,
}
impl<H> schnellru::Limiter<ValueCacheKey<H>, CachedValue<H>> for LocalValueCacheLimiter
where
H: AsRef<[u8]>,
{
type KeyToInsert<'a> = ValueCacheRef<'a, H>;
type LinkType = u32;
#[inline]
fn is_over_the_limit(&self, length: usize) -> bool {
if length <= 1 {
return false
}
self.current_heap_size > LOCAL_VALUE_CACHE_MAX_HEAP_SIZE
}
#[inline]
fn on_insert(
&mut self,
_length: usize,
key: Self::KeyToInsert<'_>,
value: CachedValue<H>,
) -> Option<(ValueCacheKey<H>, CachedValue<H>)> {
self.current_heap_size += key.storage_key.len();
Some((key.into(), value))
}
#[inline]
fn on_replace(
&mut self,
_length: usize,
_old_key: &mut ValueCacheKey<H>,
_new_key: ValueCacheRef<H>,
_old_value: &mut CachedValue<H>,
_new_value: &mut CachedValue<H>,
) -> bool {
debug_assert_eq!(_old_key.storage_key.len(), _new_key.storage_key.len());
true
}
#[inline]
fn on_removed(&mut self, key: &mut ValueCacheKey<H>, _: &mut CachedValue<H>) {
self.current_heap_size -= key.storage_key.len();
}
#[inline]
fn on_cleared(&mut self) {
self.current_heap_size = 0;
}
#[inline]
fn on_grow(&mut self, new_memory_usage: usize) -> bool {
new_memory_usage <= LOCAL_VALUE_CACHE_MAX_INLINE_SIZE
}
}
#[derive(Default)]
struct HitStats {
shared_hits: AtomicU64,
shared_fetch_attempts: AtomicU64,
local_hits: AtomicU64,
local_fetch_attempts: AtomicU64,
}
impl std::fmt::Display for HitStats {
fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
let shared_hits = self.shared_hits.load(Ordering::Relaxed);
let shared_fetch_attempts = self.shared_fetch_attempts.load(Ordering::Relaxed);
let local_hits = self.local_hits.load(Ordering::Relaxed);
let local_fetch_attempts = self.local_fetch_attempts.load(Ordering::Relaxed);
if shared_fetch_attempts == 0 && local_hits == 0 {
write!(fmt, "empty")
} else {
let percent_local = (local_hits as f32 / local_fetch_attempts as f32) * 100.0;
let percent_shared = (shared_hits as f32 / shared_fetch_attempts as f32) * 100.0;
write!(
fmt,
"local hit rate = {}% [{}/{}], shared hit rate = {}% [{}/{}]",
percent_local as u32,
local_hits,
local_fetch_attempts,
percent_shared as u32,
shared_hits,
shared_fetch_attempts
)
}
}
}
#[derive(Default)]
struct TrieHitStats {
node_cache: HitStats,
value_cache: HitStats,
}
pub(crate) struct NodeCached<H> {
pub node: NodeOwned<H>,
pub is_from_shared_cache: bool,
}
impl<H> NodeCached<H> {
fn heap_size(&self) -> usize {
self.node.size_in_bytes() - std::mem::size_of::<NodeOwned<H>>()
}
}
type NodeCacheMap<H> = LruMap<H, NodeCached<H>, LocalNodeCacheLimiter, schnellru::RandomState>;
type ValueCacheMap<H> = LruMap<
ValueCacheKey<H>,
CachedValue<H>,
LocalValueCacheLimiter,
BuildNoHashHasher<ValueCacheKey<H>>,
>;
type ValueAccessSet =
LruMap<ValueCacheKeyHash, (), schnellru::ByLength, BuildNoHashHasher<ValueCacheKeyHash>>;
pub struct LocalTrieCache<H: Hasher> {
shared: SharedTrieCache<H>,
node_cache: Mutex<NodeCacheMap<H::Out>>,
value_cache: Mutex<ValueCacheMap<H::Out>>,
shared_value_cache_access: Mutex<ValueAccessSet>,
stats: TrieHitStats,
}
impl<H: Hasher> LocalTrieCache<H> {
pub fn as_trie_db_cache(&self, storage_root: H::Out) -> TrieCache<'_, H> {
let value_cache = ValueCache::ForStorageRoot {
storage_root,
local_value_cache: self.value_cache.lock(),
shared_value_cache_access: self.shared_value_cache_access.lock(),
buffered_value: None,
};
TrieCache {
shared_cache: self.shared.clone(),
local_cache: self.node_cache.lock(),
value_cache,
stats: &self.stats,
}
}
pub fn as_trie_db_mut_cache(&self) -> TrieCache<'_, H> {
TrieCache {
shared_cache: self.shared.clone(),
local_cache: self.node_cache.lock(),
value_cache: ValueCache::Fresh(Default::default()),
stats: &self.stats,
}
}
}
impl<H: Hasher> Drop for LocalTrieCache<H> {
fn drop(&mut self) {
tracing::debug!(
target: LOG_TARGET,
"Local node trie cache dropped: {}",
self.stats.node_cache
);
tracing::debug!(
target: LOG_TARGET,
"Local value trie cache dropped: {}",
self.stats.value_cache
);
let mut shared_inner = match self.shared.write_lock_inner() {
Some(inner) => inner,
None => {
tracing::warn!(
target: LOG_TARGET,
"Timeout while trying to acquire a write lock for the shared trie cache"
);
return
},
};
shared_inner.node_cache_mut().update(self.node_cache.get_mut().drain());
shared_inner.value_cache_mut().update(
self.value_cache.get_mut().drain(),
self.shared_value_cache_access.get_mut().drain().map(|(key, ())| key),
);
}
}
enum ValueCache<'a, H: Hasher> {
Fresh(HashMap<Arc<[u8]>, CachedValue<H::Out>>),
ForStorageRoot {
shared_value_cache_access: MutexGuard<'a, ValueAccessSet>,
local_value_cache: MutexGuard<'a, ValueCacheMap<H::Out>>,
storage_root: H::Out,
buffered_value: Option<CachedValue<H::Out>>,
},
}
impl<H: Hasher> ValueCache<'_, H> {
fn get(
&mut self,
key: &[u8],
shared_cache: &SharedTrieCache<H>,
stats: &HitStats,
) -> Option<&CachedValue<H::Out>> {
stats.local_fetch_attempts.fetch_add(1, Ordering::Relaxed);
match self {
Self::Fresh(map) =>
if let Some(value) = map.get(key) {
stats.local_hits.fetch_add(1, Ordering::Relaxed);
Some(value)
} else {
None
},
Self::ForStorageRoot {
local_value_cache,
shared_value_cache_access,
storage_root,
buffered_value,
} => {
let hash = ValueCacheKey::hash_data(key, storage_root);
if let Some(value) = local_value_cache
.peek_by_hash(hash.raw(), |existing_key, _| {
existing_key.is_eq(storage_root, key)
}) {
stats.local_hits.fetch_add(1, Ordering::Relaxed);
return Some(value)
}
stats.shared_fetch_attempts.fetch_add(1, Ordering::Relaxed);
if let Some(value) = shared_cache.peek_value_by_hash(hash, storage_root, key) {
stats.shared_hits.fetch_add(1, Ordering::Relaxed);
shared_value_cache_access.insert(hash, ());
*buffered_value = Some(value.clone());
return buffered_value.as_ref()
}
None
},
}
}
fn insert(&mut self, key: &[u8], value: CachedValue<H::Out>) {
match self {
Self::Fresh(map) => {
map.insert(key.into(), value);
},
Self::ForStorageRoot { local_value_cache, storage_root, .. } => {
local_value_cache.insert(ValueCacheRef::new(key, *storage_root), value);
},
}
}
}
pub struct TrieCache<'a, H: Hasher> {
shared_cache: SharedTrieCache<H>,
local_cache: MutexGuard<'a, NodeCacheMap<H::Out>>,
value_cache: ValueCache<'a, H>,
stats: &'a TrieHitStats,
}
impl<'a, H: Hasher> TrieCache<'a, H> {
pub fn merge_into(self, local: &LocalTrieCache<H>, storage_root: H::Out) {
let ValueCache::Fresh(cache) = self.value_cache else { return };
if !cache.is_empty() {
let mut value_cache = local.value_cache.lock();
let partial_hash = ValueCacheKey::hash_partial_data(&storage_root);
cache.into_iter().for_each(|(k, v)| {
let hash = ValueCacheKeyHash::from_hasher_and_storage_key(partial_hash.clone(), &k);
let k = ValueCacheRef { storage_root, storage_key: &k, hash };
value_cache.insert(k, v);
});
}
}
}
impl<'a, H: Hasher> trie_db::TrieCache<NodeCodec<H>> for TrieCache<'a, H> {
fn get_or_insert_node(
&mut self,
hash: H::Out,
fetch_node: &mut dyn FnMut() -> trie_db::Result<NodeOwned<H::Out>, H::Out, Error<H::Out>>,
) -> trie_db::Result<&NodeOwned<H::Out>, H::Out, Error<H::Out>> {
let mut is_local_cache_hit = true;
self.stats.node_cache.local_fetch_attempts.fetch_add(1, Ordering::Relaxed);
let node = self.local_cache.get_or_insert_fallible(hash, || {
is_local_cache_hit = false;
self.stats.node_cache.shared_fetch_attempts.fetch_add(1, Ordering::Relaxed);
if let Some(node) = self.shared_cache.peek_node(&hash) {
self.stats.node_cache.shared_hits.fetch_add(1, Ordering::Relaxed);
tracing::trace!(target: LOG_TARGET, ?hash, "Serving node from shared cache");
return Ok(NodeCached::<H::Out> { node: node.clone(), is_from_shared_cache: true })
}
match fetch_node() {
Ok(node) => {
tracing::trace!(target: LOG_TARGET, ?hash, "Serving node from database");
Ok(NodeCached::<H::Out> { node, is_from_shared_cache: false })
},
Err(error) => {
tracing::trace!(target: LOG_TARGET, ?hash, "Serving node from database failed");
Err(error)
},
}
});
if is_local_cache_hit {
tracing::trace!(target: LOG_TARGET, ?hash, "Serving node from local cache");
self.stats.node_cache.local_hits.fetch_add(1, Ordering::Relaxed);
}
Ok(&node?
.expect("you can always insert at least one element into the local cache; qed")
.node)
}
fn get_node(&mut self, hash: &H::Out) -> Option<&NodeOwned<H::Out>> {
let mut is_local_cache_hit = true;
self.stats.node_cache.local_fetch_attempts.fetch_add(1, Ordering::Relaxed);
let cached_node = self.local_cache.get_or_insert_fallible(*hash, || {
is_local_cache_hit = false;
self.stats.node_cache.shared_fetch_attempts.fetch_add(1, Ordering::Relaxed);
if let Some(node) = self.shared_cache.peek_node(&hash) {
self.stats.node_cache.shared_hits.fetch_add(1, Ordering::Relaxed);
tracing::trace!(target: LOG_TARGET, ?hash, "Serving node from shared cache");
Ok(NodeCached::<H::Out> { node: node.clone(), is_from_shared_cache: true })
} else {
tracing::trace!(target: LOG_TARGET, ?hash, "Serving node from cache failed");
Err(())
}
});
if is_local_cache_hit {
tracing::trace!(target: LOG_TARGET, ?hash, "Serving node from local cache");
self.stats.node_cache.local_hits.fetch_add(1, Ordering::Relaxed);
}
match cached_node {
Ok(Some(cached_node)) => Some(&cached_node.node),
Ok(None) => {
unreachable!(
"you can always insert at least one element into the local cache; qed"
);
},
Err(()) => None,
}
}
fn lookup_value_for_key(&mut self, key: &[u8]) -> Option<&CachedValue<H::Out>> {
let res = self.value_cache.get(key, &self.shared_cache, &self.stats.value_cache);
tracing::trace!(
target: LOG_TARGET,
key = ?sp_core::hexdisplay::HexDisplay::from(&key),
found = res.is_some(),
"Looked up value for key",
);
res
}
fn cache_value_for_key(&mut self, key: &[u8], data: CachedValue<H::Out>) {
tracing::trace!(
target: LOG_TARGET,
key = ?sp_core::hexdisplay::HexDisplay::from(&key),
"Caching value for key",
);
self.value_cache.insert(key, data);
}
}
#[cfg(test)]
mod tests {
use super::*;
use trie_db::{Bytes, Trie, TrieDBBuilder, TrieDBMutBuilder, TrieHash, TrieMut};
type MemoryDB = crate::MemoryDB<sp_core::Blake2Hasher>;
type Layout = crate::LayoutV1<sp_core::Blake2Hasher>;
type Cache = super::SharedTrieCache<sp_core::Blake2Hasher>;
type Recorder = crate::recorder::Recorder<sp_core::Blake2Hasher>;
const TEST_DATA: &[(&[u8], &[u8])] =
&[(b"key1", b"val1"), (b"key2", &[2; 64]), (b"key3", b"val3"), (b"key4", &[4; 64])];
const CACHE_SIZE_RAW: usize = 1024 * 10;
const CACHE_SIZE: CacheSize = CacheSize::new(CACHE_SIZE_RAW);
fn create_trie() -> (MemoryDB, TrieHash<Layout>) {
let mut db = MemoryDB::default();
let mut root = Default::default();
{
let mut trie = TrieDBMutBuilder::<Layout>::new(&mut db, &mut root).build();
for (k, v) in TEST_DATA {
trie.insert(k, v).expect("Inserts data");
}
}
(db, root)
}
#[test]
fn basic_cache_works() {
let (db, root) = create_trie();
let shared_cache = Cache::new(CACHE_SIZE);
let local_cache = shared_cache.local_cache();
{
let mut cache = local_cache.as_trie_db_cache(root);
let trie = TrieDBBuilder::<Layout>::new(&db, &root).with_cache(&mut cache).build();
assert_eq!(TEST_DATA[0].1.to_vec(), trie.get(TEST_DATA[0].0).unwrap().unwrap());
}
assert!(shared_cache.read_lock_inner().value_cache().lru.is_empty());
assert!(shared_cache.read_lock_inner().node_cache().lru.is_empty());
drop(local_cache);
assert!(shared_cache.read_lock_inner().node_cache().lru.len() >= 1);
let cached_data = shared_cache
.read_lock_inner()
.value_cache()
.lru
.peek(&ValueCacheKey::new_value(TEST_DATA[0].0, root))
.unwrap()
.clone();
assert_eq!(Bytes::from(TEST_DATA[0].1.to_vec()), cached_data.data().flatten().unwrap());
let fake_data = Bytes::from(&b"fake_data"[..]);
let local_cache = shared_cache.local_cache();
shared_cache.write_lock_inner().unwrap().value_cache_mut().lru.insert(
ValueCacheKey::new_value(TEST_DATA[1].0, root),
(fake_data.clone(), Default::default()).into(),
);
{
let mut cache = local_cache.as_trie_db_cache(root);
let trie = TrieDBBuilder::<Layout>::new(&db, &root).with_cache(&mut cache).build();
assert_eq!(b"fake_data".to_vec(), trie.get(TEST_DATA[1].0).unwrap().unwrap());
}
}
#[test]
fn trie_db_mut_cache_works() {
let (mut db, root) = create_trie();
let new_key = b"new_key".to_vec();
let new_value = vec![23; 64];
let shared_cache = Cache::new(CACHE_SIZE);
let mut new_root = root;
{
let local_cache = shared_cache.local_cache();
let mut cache = local_cache.as_trie_db_mut_cache();
{
let mut trie = TrieDBMutBuilder::<Layout>::from_existing(&mut db, &mut new_root)
.with_cache(&mut cache)
.build();
trie.insert(&new_key, &new_value).unwrap();
}
cache.merge_into(&local_cache, new_root);
}
let cached_data = shared_cache
.read_lock_inner()
.value_cache()
.lru
.peek(&ValueCacheKey::new_value(new_key, new_root))
.unwrap()
.clone();
assert_eq!(Bytes::from(new_value), cached_data.data().flatten().unwrap());
}
#[test]
fn trie_db_cache_and_recorder_work_together() {
let (db, root) = create_trie();
let shared_cache = Cache::new(CACHE_SIZE);
for i in 0..5 {
if i == 2 {
shared_cache.reset_node_cache();
} else if i == 3 {
shared_cache.reset_value_cache();
}
let local_cache = shared_cache.local_cache();
let recorder = Recorder::default();
{
let mut cache = local_cache.as_trie_db_cache(root);
let mut recorder = recorder.as_trie_recorder(root);
let trie = TrieDBBuilder::<Layout>::new(&db, &root)
.with_cache(&mut cache)
.with_recorder(&mut recorder)
.build();
for (key, value) in TEST_DATA {
assert_eq!(*value, trie.get(&key).unwrap().unwrap());
}
}
let storage_proof = recorder.drain_storage_proof();
let memory_db: MemoryDB = storage_proof.into_memory_db();
{
let trie = TrieDBBuilder::<Layout>::new(&memory_db, &root).build();
for (key, value) in TEST_DATA {
assert_eq!(*value, trie.get(&key).unwrap().unwrap());
}
}
}
}
#[test]
fn trie_db_mut_cache_and_recorder_work_together() {
const DATA_TO_ADD: &[(&[u8], &[u8])] = &[(b"key11", &[45; 78]), (b"key33", &[78; 89])];
let (db, root) = create_trie();
let shared_cache = Cache::new(CACHE_SIZE);
for i in 0..5 {
if i == 2 {
shared_cache.reset_node_cache();
} else if i == 3 {
shared_cache.reset_value_cache();
}
let recorder = Recorder::default();
let local_cache = shared_cache.local_cache();
let mut new_root = root;
{
let mut db = db.clone();
let mut cache = local_cache.as_trie_db_cache(root);
let mut recorder = recorder.as_trie_recorder(root);
let mut trie = TrieDBMutBuilder::<Layout>::from_existing(&mut db, &mut new_root)
.with_cache(&mut cache)
.with_recorder(&mut recorder)
.build();
for (key, value) in DATA_TO_ADD {
trie.insert(key, value).unwrap();
}
}
let storage_proof = recorder.drain_storage_proof();
let mut memory_db: MemoryDB = storage_proof.into_memory_db();
let mut proof_root = root;
{
let mut trie =
TrieDBMutBuilder::<Layout>::from_existing(&mut memory_db, &mut proof_root)
.build();
for (key, value) in DATA_TO_ADD {
trie.insert(key, value).unwrap();
}
}
assert_eq!(new_root, proof_root)
}
}
#[test]
fn cache_lru_works() {
let (db, root) = create_trie();
let shared_cache = Cache::new(CACHE_SIZE);
{
let local_cache = shared_cache.local_cache();
let mut cache = local_cache.as_trie_db_cache(root);
let trie = TrieDBBuilder::<Layout>::new(&db, &root).with_cache(&mut cache).build();
for (k, _) in TEST_DATA {
trie.get(k).unwrap().unwrap();
}
}
assert!(shared_cache
.read_lock_inner()
.value_cache()
.lru
.iter()
.map(|d| d.0)
.all(|l| TEST_DATA.iter().any(|d| &*l.storage_key == d.0)));
for _ in 0..2 {
{
let local_cache = shared_cache.local_cache();
let mut cache = local_cache.as_trie_db_cache(root);
let trie = TrieDBBuilder::<Layout>::new(&db, &root).with_cache(&mut cache).build();
for (k, _) in TEST_DATA.iter().take(2) {
trie.get(k).unwrap().unwrap();
}
}
assert!(shared_cache
.read_lock_inner()
.value_cache()
.lru
.iter()
.take(2)
.map(|d| d.0)
.all(|l| { TEST_DATA.iter().take(2).any(|d| &*l.storage_key == d.0) }));
shared_cache.reset_value_cache();
}
let most_recently_used_nodes = shared_cache
.read_lock_inner()
.node_cache()
.lru
.iter()
.map(|d| *d.0)
.collect::<Vec<_>>();
{
let local_cache = shared_cache.local_cache();
let mut cache = local_cache.as_trie_db_cache(root);
let trie = TrieDBBuilder::<Layout>::new(&db, &root).with_cache(&mut cache).build();
for (k, _) in TEST_DATA.iter().skip(2) {
trie.get(k).unwrap().unwrap();
}
}
assert_ne!(
most_recently_used_nodes,
shared_cache
.read_lock_inner()
.node_cache()
.lru
.iter()
.map(|d| *d.0)
.collect::<Vec<_>>()
);
}
#[test]
fn cache_respects_bounds() {
let (mut db, root) = create_trie();
let shared_cache = Cache::new(CACHE_SIZE);
{
let local_cache = shared_cache.local_cache();
let mut new_root = root;
{
let mut cache = local_cache.as_trie_db_cache(root);
{
let mut trie =
TrieDBMutBuilder::<Layout>::from_existing(&mut db, &mut new_root)
.with_cache(&mut cache)
.build();
let value = vec![10u8; 100];
for i in 0..CACHE_SIZE_RAW / 100 * 2 {
trie.insert(format!("key{}", i).as_bytes(), &value).unwrap();
}
}
cache.merge_into(&local_cache, new_root);
}
}
assert!(shared_cache.used_memory_size() < CACHE_SIZE_RAW);
}
}