refactor: trie cache configuration (#7566)

This is a pure refactor, including two main refactors:

- replace `TrieCacheFactory` with `TrieConfig`
- move config constants into `TrieConfig`

The motivation is to facilitate wider configuration and changes to the
config format in follow-up PRs.

See also #7564.

core/store/src/config.rs

@@ -41,7 +41,7 @@ pub struct StoreConfig {
     /// Default value: ShardUId {version: 1, shard_id: 3} -> 45_000_000
     /// We're still experimenting with this parameter and it seems decreasing its value can improve
     /// the performance of the storage
-    pub trie_cache_capacities: Vec<(ShardUId, usize)>,
+    pub trie_cache_capacities: Vec<(ShardUId, u64)>,
 
     /// Path where to create RocksDB checkpoints during database migrations or
     /// `false` to disable that feature.

core/store/src/lib.rs

@@ -32,7 +32,7 @@ pub use crate::trie::iterator::TrieIterator;
 pub use crate::trie::update::{TrieUpdate, TrieUpdateIterator, TrieUpdateValuePtr};
 pub use crate::trie::{
     estimator, split_state, ApplyStatePartResult, KeyForStateChanges, PartialStorage, ShardTries,
-    Trie, TrieAccess, TrieCache, TrieCacheFactory, TrieCachingStorage, TrieChanges, TrieStorage,
+    Trie, TrieAccess, TrieCache, TrieCachingStorage, TrieChanges, TrieConfig, TrieStorage,
     WrappedTrieChanges,
 };
 

core/store/src/test_utils.rs

@@ -4,7 +4,7 @@ use rand::seq::SliceRandom;
 use rand::Rng;
 
 use crate::db::TestDB;
-use crate::{NodeStorage, ShardTries, Store, TrieCacheFactory};
+use crate::{NodeStorage, ShardTries, Store};
 use near_primitives::account::id::AccountId;
 use near_primitives::hash::CryptoHash;
 use near_primitives::receipt::{DataReceipt, Receipt, ReceiptEnum};
@@ -31,8 +31,7 @@ pub fn create_tries() -> ShardTries {
 
 pub fn create_tries_complex(shard_version: ShardVersion, num_shards: NumShards) -> ShardTries {
     let store = create_test_store();
-    let trie_cache_factory = TrieCacheFactory::new(Default::default(), shard_version, num_shards);
-    ShardTries::new(store, trie_cache_factory)
+    ShardTries::test_shard_version(store, shard_version, num_shards)
 }
 
 pub fn test_populate_trie(

core/store/src/trie/config.rs

@@ -0,0 +1,99 @@
+use near_primitives::shard_layout::ShardUId;
+use std::collections::HashMap;
+
+/// Default number of cache entries.
+/// It was chosen to fit into RAM well. RAM spend on trie cache should not exceed 50_000 * 4 (number of shards) *
+/// TRIE_LIMIT_CACHED_VALUE_SIZE * 2 (number of caches - for regular and view client) = 0.4 GB.
+/// In our tests on a single shard, it barely occupied 40 MB, which is dominated by state cache size
+/// with 512 MB limit. The total RAM usage for a single shard was 1 GB.
+const TRIE_DEFAULT_SHARD_CACHE_SIZE: u64 = if cfg!(feature = "no_cache") { 1 } else { 50000 };
+
+/// Default total size of values which may simultaneously exist the cache.
+/// It is chosen by the estimation of the largest contract storage size we are aware as of 23/08/2022.
+const DEFAULT_SHARD_CACHE_TOTAL_SIZE_LIMIT: u64 =
+    if cfg!(feature = "no_cache") { 1 } else { 3_000_000_000 };
+
+/// Capacity for the deletions queue.
+/// It is chosen to fit all hashes of deleted nodes for 3 completely full blocks.
+const DEFAULT_SHARD_CACHE_DELETIONS_QUEUE_CAPACITY: usize =
+    if cfg!(feature = "no_cache") { 1 } else { 100_000 };
+
+/// Values above this size (in bytes) are never cached.
+/// Note that most of Trie inner nodes are smaller than this - e.g. branches use around 32 * 16 = 512 bytes.
+const TRIE_LIMIT_CACHED_VALUE_SIZE: usize = 1000;
+
+/// Stores necessary configuration for the creation of tries.
+#[derive(Default)]
+pub struct TrieConfig {
+    pub shard_cache_config: ShardCacheConfig,
+    pub view_shard_cache_config: ShardCacheConfig,
+}
+
+pub struct ShardCacheConfig {
+    /// Shard cache capacity in number of trie nodes.
+    pub default_max_entries: u64,
+    /// Limits the sum of all cached value sizes.
+    ///
+    /// This is useful to limit total memory consumption. However, crucially this
+    /// is not a hard limit. It only limits the sum of all cached values, not
+    /// factoring in the overhead for each entry.
+    pub default_max_total_bytes: u64,
+    /// Overrides `default_max_entries` per shard.
+    pub override_max_entries: HashMap<ShardUId, u64>,
+    /// Overrides `default_max_total_bytes` per shard.
+    pub override_max_total_bytes: HashMap<ShardUId, u64>,
+}
+
+impl TrieConfig {
+    /// Shard cache capacity in number of trie nodes.
+    pub fn shard_cache_capacity(&self, shard_uid: ShardUId, is_view: bool) -> u64 {
+        if is_view { &self.view_shard_cache_config } else { &self.shard_cache_config }
+            .capacity(shard_uid)
+    }
+
+    /// Shard cache capacity in total bytes.
+    pub fn shard_cache_total_size_limit(&self, shard_uid: ShardUId, is_view: bool) -> u64 {
+        if is_view { &self.view_shard_cache_config } else { &self.shard_cache_config }
+            .total_size_limit(shard_uid)
+    }
+
+    /// Size limit in bytes per single value for caching in shard caches.
+    pub fn max_cached_value_size() -> usize {
+        TRIE_LIMIT_CACHED_VALUE_SIZE
+    }
+
+    /// Capacity for deletion queue in which nodes are after unforced eviction.
+    ///
+    /// The shard cache uses LRU eviction policy for forced evictions. But when a
+    /// trie value is overwritten or deleted, the associated nodes are no longer
+    /// useful, with the exception of forks.
+    /// Thus, deleted and overwritten values are evicted to the deletion queue which
+    /// delays the actual eviction.
+    pub fn deletions_queue_capacity(&self) -> usize {
+        DEFAULT_SHARD_CACHE_DELETIONS_QUEUE_CAPACITY
+    }
+}
+
+impl ShardCacheConfig {
+    fn capacity(&self, shard_uid: ShardUId) -> u64 {
+        self.override_max_entries.get(&shard_uid).cloned().unwrap_or(self.default_max_entries)
+    }
+
+    fn total_size_limit(&self, shard_uid: ShardUId) -> u64 {
+        self.override_max_total_bytes
+            .get(&shard_uid)
+            .cloned()
+            .unwrap_or(self.default_max_total_bytes)
+    }
+}
+
+impl Default for ShardCacheConfig {
+    fn default() -> Self {
+        Self {
+            default_max_entries: TRIE_DEFAULT_SHARD_CACHE_SIZE,
+            default_max_total_bytes: DEFAULT_SHARD_CACHE_TOTAL_SIZE_LIMIT,
+            override_max_entries: HashMap::default(),
+            override_max_total_bytes: HashMap::default(),
+        }
+    }
+}

core/store/src/trie/mod.rs

@@ -15,17 +15,17 @@ use near_primitives::trie_key::TrieKey;
 use near_primitives::types::{StateRoot, StateRootNode};
 
 use crate::flat_state::FlatState;
+pub use crate::trie::config::TrieConfig;
 use crate::trie::insert_delete::NodesStorage;
 use crate::trie::iterator::TrieIterator;
 use crate::trie::nibble_slice::NibbleSlice;
-pub use crate::trie::shard_tries::{
-    KeyForStateChanges, ShardTries, TrieCacheFactory, WrappedTrieChanges,
-};
+pub use crate::trie::shard_tries::{KeyForStateChanges, ShardTries, WrappedTrieChanges};
 pub use crate::trie::trie_storage::{TrieCache, TrieCachingStorage, TrieStorage};
 use crate::trie::trie_storage::{TrieMemoryPartialStorage, TrieRecordingStorage};
 use crate::StorageError;
 pub use near_primitives::types::TrieNodesCount;
 
+mod config;
 mod insert_delete;
 pub mod iterator;
 mod nibble_slice;

core/store/src/trie/shard_tries.rs

@@ -5,60 +5,21 @@ use std::sync::{Arc, RwLock};
 use borsh::BorshSerialize;
 use near_primitives::borsh::maybestd::collections::HashMap;
 use near_primitives::hash::CryptoHash;
-use near_primitives::shard_layout;
-use near_primitives::shard_layout::{ShardUId, ShardVersion};
+use near_primitives::shard_layout::{self, ShardUId, ShardVersion};
 use near_primitives::trie_key::TrieKey;
 use near_primitives::types::{
     NumShards, RawStateChange, RawStateChangesWithTrieKey, StateChangeCause, StateRoot,
 };
 
+use crate::trie::config::TrieConfig;
 use crate::trie::trie_storage::{TrieCache, TrieCachingStorage};
 use crate::trie::{TrieRefcountChange, POISONED_LOCK_ERR};
 use crate::{metrics, DBCol, DBOp, DBTransaction};
 use crate::{Store, StoreUpdate, Trie, TrieChanges, TrieUpdate};
 
-/// Responsible for creation of trie caches, stores necessary configuration for it.
-#[derive(Default)]
-pub struct TrieCacheFactory {
-    capacities: HashMap<ShardUId, usize>,
-    shard_version: ShardVersion,
-    num_shards: NumShards,
-}
-
-impl TrieCacheFactory {
-    pub fn new(
-        capacities: HashMap<ShardUId, usize>,
-        shard_version: ShardVersion,
-        num_shards: NumShards,
-    ) -> Self {
-        Self { capacities, shard_version, num_shards }
-    }
-
-    /// Create new cache for the given shard uid.
-    pub fn create_cache(&self, shard_uid: &ShardUId, is_view: bool) -> TrieCache {
-        let capacity = if is_view { None } else { self.capacities.get(shard_uid) };
-        match capacity {
-            Some(capacity) => TrieCache::with_capacities(*capacity, shard_uid.shard_id, is_view),
-            None => TrieCache::new(shard_uid.shard_id, is_view),
-        }
-    }
-
-    /// Create caches on the initialization of storage structures.
-    pub fn create_initial_caches(&self, is_view: bool) -> HashMap<ShardUId, TrieCache> {
-        assert_ne!(self.num_shards, 0);
-        let shards: Vec<_> = (0..self.num_shards)
-            .map(|shard_id| ShardUId { version: self.shard_version, shard_id: shard_id as u32 })
-            .collect();
-        shards
-            .iter()
-            .map(|&shard_uid| (shard_uid, self.create_cache(&shard_uid, is_view)))
-            .collect()
-    }
-}
-
 struct ShardTriesInner {
     store: Store,
-    trie_cache_factory: TrieCacheFactory,
+    trie_config: TrieConfig,
     /// Cache reserved for client actor to use
     caches: RwLock<HashMap<ShardUId, TrieCache>>,
     /// Cache for readers.
@@ -69,19 +30,43 @@ struct ShardTriesInner {
 pub struct ShardTries(Arc<ShardTriesInner>);
 
 impl ShardTries {
-    pub fn new(store: Store, trie_cache_factory: TrieCacheFactory) -> Self {
-        let caches = trie_cache_factory.create_initial_caches(false);
-        let view_caches = trie_cache_factory.create_initial_caches(true);
+    pub fn new(store: Store, trie_config: TrieConfig, shard_uids: &[ShardUId]) -> Self {
+        let caches = Self::create_initial_caches(&trie_config, &shard_uids, false);
+        let view_caches = Self::create_initial_caches(&trie_config, &shard_uids, true);
         ShardTries(Arc::new(ShardTriesInner {
             store,
-            trie_cache_factory,
+            trie_config,
             caches: RwLock::new(caches),
             view_caches: RwLock::new(view_caches),
         }))
     }
 
+    /// Create `ShardTries` with a fixed number of shards with shard version 0.
+    ///
+    /// If your test cares about the shard version, use `test_shard_version` instead.
     pub fn test(store: Store, num_shards: NumShards) -> Self {
-        Self::new(store, TrieCacheFactory::new(Default::default(), 0, num_shards))
+        let shard_version = 0;
+        Self::test_shard_version(store, shard_version, num_shards)
+    }
+
+    pub fn test_shard_version(store: Store, version: ShardVersion, num_shards: NumShards) -> Self {
+        assert_ne!(0, num_shards);
+        let shard_uids: Vec<ShardUId> =
+            (0..num_shards as u32).map(|shard_id| ShardUId { shard_id, version }).collect();
+        let trie_config = TrieConfig::default();
+        ShardTries::new(store, trie_config, &shard_uids)
+    }
+
+    /// Create caches for all shards according to the trie config.
+    fn create_initial_caches(
+        config: &TrieConfig,
+        shard_uids: &[ShardUId],
+        is_view: bool,
+    ) -> HashMap<ShardUId, TrieCache> {
+        shard_uids
+            .iter()
+            .map(|&shard_uid| (shard_uid, TrieCache::new(config, shard_uid, is_view)))
+            .collect()
     }
 
     pub fn is_same(&self, other: &Self) -> bool {
@@ -109,7 +94,7 @@ impl ShardTries {
             let mut caches = caches_to_use.write().expect(POISONED_LOCK_ERR);
             caches
                 .entry(shard_uid)
-                .or_insert_with(|| self.0.trie_cache_factory.create_cache(&shard_uid, is_view))
+                .or_insert_with(|| TrieCache::new(&self.0.trie_config, shard_uid, is_view))
                 .clone()
         };
         let storage =
@@ -169,7 +154,7 @@ impl ShardTries {
         for (shard_uid, ops) in shards {
             let cache = caches
                 .entry(shard_uid)
-                .or_insert_with(|| self.0.trie_cache_factory.create_cache(&shard_uid, false))
+                .or_insert_with(|| TrieCache::new(&self.0.trie_config, shard_uid, false))
                 .clone();
             cache.update_cache(ops);
         }