From 350423e8060bbccecb1f77323a19dff47d50a519 Mon Sep 17 00:00:00 2001
From: Daniel Freire <daniel.freire@cloudwalk.io>
Date: Thu, 19 Dec 2024 10:30:54 -0300
Subject: [PATCH] enha: add cache config

---
 src/eth/storage/cache.rs           | 40 +++++++++++++++++++++++++-----
 src/eth/storage/mod.rs             |  8 ++++--
 src/eth/storage/stratus_storage.rs | 17 +++++++------
 3 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/src/eth/storage/cache.rs b/src/eth/storage/cache.rs
index 1de6a3b81..192eff64d 100644
--- a/src/eth/storage/cache.rs
+++ b/src/eth/storage/cache.rs
@@ -1,3 +1,5 @@
+use clap::Parser;
+use display_json::DebugAsJson;
 use quick_cache::sync::Cache;
 use quick_cache::sync::DefaultLifecycle;
 use quick_cache::UnitWeighter;
@@ -16,16 +18,42 @@ pub struct StorageCache {
     account_cache: Cache<Address, Account, UnitWeighter, FxBuildHasher>,
 }
 
-impl Default for StorageCache {
-    fn default() -> Self {
-        Self {
-            slot_cache: Cache::with(100_000, 100_000, UnitWeighter, FxBuildHasher, DefaultLifecycle::default()),
-            account_cache: Cache::with(20_000, 20_000, UnitWeighter, FxBuildHasher, DefaultLifecycle::default()),
-        }
+#[derive(DebugAsJson, Clone, Parser, serde::Serialize)]
+pub struct CacheConfig {
+    /// Capacity of slot cache
+    #[arg(long = "slot-cache-capacity", env = "SLOT_CACHE_CAPACITY", default_value = "100000")]
+    pub slot_cache_capacity: usize,
+
+    /// Capacity of account cache
+    #[arg(long = "account-cache-capacity", env = "ACCOUNT_CACHE_CAPACITY", default_value = "20000")]
+    pub account_cache_capacity: usize,
+}
+
+impl CacheConfig {
+    pub fn init(&self) -> StorageCache {
+        StorageCache::new(self)
     }
 }
 
 impl StorageCache {
+    pub fn new(config: &CacheConfig) -> Self {
+        Self {
+            slot_cache: Cache::with(
+                config.slot_cache_capacity,
+                config.slot_cache_capacity as u64,
+                UnitWeighter,
+                FxBuildHasher,
+                DefaultLifecycle::default(),
+            ),
+            account_cache: Cache::with(
+                config.account_cache_capacity,
+                config.account_cache_capacity as u64,
+                UnitWeighter,
+                FxBuildHasher,
+                DefaultLifecycle::default(),
+            ),
+        }
+    }
     pub fn clear(&self) {
         self.slot_cache.clear();
         self.account_cache.clear();
diff --git a/src/eth/storage/mod.rs b/src/eth/storage/mod.rs
index f922fde17..d849055b7 100644
--- a/src/eth/storage/mod.rs
+++ b/src/eth/storage/mod.rs
@@ -1,5 +1,6 @@
 //! Ethereum / EVM storage.
 
+use cache::CacheConfig;
 pub use cache::StorageCache;
 pub use permanent::InMemoryPermanentStorage;
 pub use permanent::PermanentStorage;
@@ -129,6 +130,9 @@ pub struct StorageConfig {
 
     #[clap(flatten)]
     pub perm_storage: PermanentStorageConfig,
+
+    #[clap(flatten)]
+    pub cache: CacheConfig,
 }
 
 impl StorageConfig {
@@ -136,9 +140,9 @@ impl StorageConfig {
     pub fn init(&self) -> Result<Arc<StratusStorage>, StratusError> {
         let perm_storage = self.perm_storage.init()?;
         let temp_storage = self.temp_storage.init(&*perm_storage)?;
-
+        let cache = self.cache.init();
         let StorageKind::StratusStorage = self.storage_kind;
-        let storage = StratusStorage::new(temp_storage, perm_storage)?;
+        let storage = StratusStorage::new(temp_storage, perm_storage, cache)?;
 
         Ok(Arc::new(storage))
     }
diff --git a/src/eth/storage/stratus_storage.rs b/src/eth/storage/stratus_storage.rs
index 503b2f1a0..44e2d299d 100644
--- a/src/eth/storage/stratus_storage.rs
+++ b/src/eth/storage/stratus_storage.rs
@@ -43,12 +43,8 @@ pub struct StratusStorage {
 
 impl StratusStorage {
     /// Creates a new storage with the specified temporary and permanent implementations.
-    pub fn new(temp: Box<dyn TemporaryStorage>, perm: Box<dyn PermanentStorage>) -> Result<Self, StratusError> {
-        let this = Self {
-            temp,
-            cache: StorageCache::default(),
-            perm,
-        };
+    pub fn new(temp: Box<dyn TemporaryStorage>, perm: Box<dyn PermanentStorage>, cache: StorageCache) -> Result<Self, StratusError> {
+        let this = Self { temp, cache, perm };
 
         // create genesis block and accounts if necessary
         #[cfg(feature = "dev")]
@@ -64,10 +60,17 @@ impl StratusStorage {
 
     #[cfg(test)]
     pub fn new_test() -> Result<Self, StratusError> {
+        use super::cache::CacheConfig;
+
         let perm = Box::new(super::InMemoryPermanentStorage::default());
         let temp = Box::new(super::InMemoryTemporaryStorage::new(0.into()));
+        let cache = CacheConfig {
+            slot_cache_capacity: 100000,
+            account_cache_capacity: 20000,
+        }
+        .init();
 
-        Self::new(temp, perm)
+        Self::new(temp, perm, cache)
     }
 }