refactor: use bytesize::ByteSize for TrieCacheConfig

core/store/src/config.rs

@@ -240,23 +240,22 @@ impl Default for StoreConfig {
             block_size: bytesize::ByteSize::kib(16),
 
             trie_cache: TrieCacheConfig {
-                default_max_bytes: 500_000_000,
+                default_max_bytes: bytesize::ByteSize::mb(500),
                 // TODO(resharding) The cache size needs to adjusted for every resharding.
                 // Make that automatic e.g. by defining the minimum cache size per account rather than shard.
                 per_shard_max_bytes: HashMap::from_iter([
                     // Temporary solution to make contracts with heavy trie access
                     // patterns on shard 3 more stable. It was chosen by the estimation
                     // of the largest contract storage size we are aware as of 23/08/2022.
-                    // Note: on >= 1.34 nearcore version use 1_000_000_000 if you have
-                    // minimal hardware.
+                    // Note: on >= 1.34 nearcore version use 1gb if you have minimal hardware.
                     // In simple nightshade the heavy contract "token.sweat" is in shard 3
-                    (ShardUId { version: 1, shard_id: 3 }, 3_000_000_000),
+                    (ShardUId { version: 1, shard_id: 3 }, bytesize::ByteSize::gb(3)),
                     // In simple nightshade v2 the heavy contract "token.sweat" is in shard 4
-                    (ShardUId { version: 2, shard_id: 4 }, 3_000_000_000),
+                    (ShardUId { version: 2, shard_id: 4 }, bytesize::ByteSize::gb(3)),
                     // Shard 1 is dedicated to aurora and it had very few cache
                     // misses even with cache size of only 50MB
-                    (ShardUId { version: 1, shard_id: 1 }, 50_000_000),
-                    (ShardUId { version: 2, shard_id: 1 }, 50_000_000),
+                    (ShardUId { version: 1, shard_id: 1 }, bytesize::ByteSize::mb(50)),
+                    (ShardUId { version: 2, shard_id: 1 }, bytesize::ByteSize::mb(50)),
                 ]),
                 shard_cache_deletions_queue_capacity: DEFAULT_SHARD_CACHE_DELETIONS_QUEUE_CAPACITY,
             },
@@ -351,9 +350,9 @@ pub struct TrieCacheConfig {
     ///
     /// This is an approximate limit that attempts to factor in data structure
     /// overhead also. It is supposed to be fairly accurate in the limit.
-    pub default_max_bytes: u64,
+    pub default_max_bytes: bytesize::ByteSize,
     /// Overwrites `default_max_bytes` for specific shards.
-    pub per_shard_max_bytes: HashMap<ShardUId, u64>,
+    pub per_shard_max_bytes: HashMap<ShardUId, bytesize::ByteSize>,
     /// Limit the number of elements in caches deletions queue for specific
     /// shard
     pub shard_cache_deletions_queue_capacity: usize,

core/store/src/trie/config.rs

@@ -8,8 +8,8 @@ use tracing::error;
 /// Default memory limit, if nothing else is configured.
 /// It is chosen to correspond roughly to the old limit, which was
 /// 50k entries * TRIE_LIMIT_CACHED_VALUE_SIZE.
-pub(crate) const DEFAULT_SHARD_CACHE_TOTAL_SIZE_LIMIT: u64 =
-    if cfg!(feature = "no_cache") { 1 } else { 50_000_000 };
+pub(crate) const DEFAULT_SHARD_CACHE_TOTAL_SIZE_LIMIT: bytesize::ByteSize =
+    if cfg!(feature = "no_cache") { bytesize::ByteSize(1) } else { bytesize::ByteSize::mb(50) };
 
 /// Capacity for the deletions queue.
 /// It is chosen to fit all hashes of deleted nodes for 3 completely full blocks.

core/store/src/trie/trie_storage.rs

@@ -237,7 +237,8 @@ impl TrieCache {
             .per_shard_max_bytes
             .get(&shard_uid)
             .copied()
-            .unwrap_or(cache_config.default_max_bytes);
+            .unwrap_or(cache_config.default_max_bytes)
+            .as_u64();
         let queue_capacity = config.deletions_queue_capacity();
         Self(Arc::new(Mutex::new(TrieCacheInner::new(
             queue_capacity,
@@ -687,10 +688,10 @@ mod trie_cache_tests {
     fn test_trie_config() {
         let mut store_config = StoreConfig::default();
 
-        const DEFAULT_SIZE: u64 = 1;
-        const S0_SIZE: u64 = 2;
-        const DEFAULT_VIEW_SIZE: u64 = 3;
-        const S0_VIEW_SIZE: u64 = 4;
+        const DEFAULT_SIZE: bytesize::ByteSize = bytesize::ByteSize(1);
+        const S0_SIZE: bytesize::ByteSize = bytesize::ByteSize(2);
+        const DEFAULT_VIEW_SIZE: bytesize::ByteSize = bytesize::ByteSize(3);
+        const S0_VIEW_SIZE: bytesize::ByteSize = bytesize::ByteSize(4);
 
         let s0 = ShardUId::single_shard();
         store_config.trie_cache.default_max_bytes = DEFAULT_SIZE;
@@ -710,11 +711,11 @@ mod trie_cache_tests {
         trie_config: &TrieConfig,
         shard_id: ShardId,
         is_view: bool,
-        expected_size: u64,
+        expected_size: bytesize::ByteSize,
     ) {
         let shard_uid = ShardUId { version: 0, shard_id: shard_id as u32 };
         let trie_cache = TrieCache::new(&trie_config, shard_uid, is_view);
-        assert_eq!(expected_size, trie_cache.lock().total_size_limit,);
-        assert_eq!(is_view, trie_cache.lock().is_view,);
+        assert_eq!(expected_size.as_u64(), trie_cache.lock().total_size_limit);
+        assert_eq!(is_view, trie_cache.lock().is_view);
     }
 }

core/store/src/trie/trie_tests.rs

@@ -380,7 +380,10 @@ mod trie_storage_tests {
         let shard_uid = ShardUId::single_shard();
         let store = create_store_with_values(&values, shard_uid);
         let mut trie_config = TrieConfig::default();
-        trie_config.shard_cache_config.per_shard_max_bytes.insert(shard_uid, shard_cache_size);
+        trie_config
+            .shard_cache_config
+            .per_shard_max_bytes
+            .insert(shard_uid, bytesize::ByteSize(shard_cache_size));
         let trie_cache = TrieCache::new(&trie_config, shard_uid, false);
         let trie_caching_storage =
             TrieCachingStorage::new(store, trie_cache.clone(), shard_uid, false, None);

integration-tests/Cargo.toml

@@ -15,6 +15,7 @@ workspace = true
 actix-rt.workspace = true
 actix.workspace = true
 anyhow.workspace = true
+bytesize.workspace = true
 borsh.workspace = true
 chrono.workspace = true
 clap.workspace = true

integration-tests/src/tests/client/state_snapshot.rs

@@ -35,7 +35,7 @@ impl StateSnaptshotTestEnv {
         store: &Store,
     ) -> Self {
         let trie_cache_config = TrieCacheConfig {
-            default_max_bytes: 50_000_000,
+            default_max_bytes: bytesize::ByteSize::mb(50),
             per_shard_max_bytes: Default::default(),
             shard_cache_deletions_queue_capacity: 0,
         };