replace cached with moka for caching

Author: ashleychandy

Cargo.lock

@@ -18,11 +18,11 @@ alloy-sol-types = { workspace = true }
 alloy-transport = { workspace = true }
 anyhow = { workspace = true }
 async-trait = { workspace = true }
-cached = { workspace = true, features = ["default"] }
 clap = { workspace = true }
 configs = { workspace = true }
 contracts = { workspace = true }
 ethrpc = { workspace = true }
+moka = { workspace = true, features = ["sync"] }
 serde = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }

crates/balance-overrides/Cargo.toml

@@ -1,5 +1,8 @@
 use {
-    crate::detector::{DetectionError, SimulationError, extract_sload_slots, mapping_slot_hash},
+    crate::{
+        cache::Cache,
+        detector::{DetectionError, SimulationError, extract_sload_slots, mapping_slot_hash},
+    },
     alloy_eips::BlockId,
     alloy_primitives::{Address, B256, TxKind, U256, keccak256, map::AddressMap},
     alloy_provider::ext::DebugApi,
@@ -11,10 +14,9 @@ use {
     },
     alloy_sol_types::SolCall,
     alloy_transport::TransportErrorKind,
-    cached::{Cached, SizedCache},
     contracts::ERC20,
     ethrpc::Web3,
-    std::{iter, sync::Mutex, time::Duration},
+    std::{iter, time::Duration},
 };
 
 /// These are the solady magic bytes for user allowances
@@ -184,8 +186,6 @@ impl ApprovalStrategy {
     }
 }
 
-type Cache = SizedCache<(Address, Option<(Address, Address)>), Option<ApprovalStrategy>>;
-
 /// Heuristic approval override detector with integrated caching.
 ///
 /// Owns the Web3 handle, detection parameters, and the per-token strategy
@@ -197,7 +197,7 @@ pub(crate) struct Detector {
     web3: Web3,
     probing_depth: u8,
     verification_timeout: Duration,
-    pub(crate) cache: Mutex<Cache>,
+    pub(crate) cache: Cache<(Address, Option<(Address, Address)>), Option<ApprovalStrategy>>,
 }
 
 impl Detector {
@@ -211,7 +211,7 @@ impl Detector {
             web3,
             probing_depth,
             verification_timeout,
-            cache: Mutex::new(SizedCache::with_size(cache_size)),
+            cache: Cache::new(u64::try_from(cache_size).expect("cache_size must be non-negative")),
         }
     }
 
@@ -231,19 +231,18 @@ impl Detector {
         tracing::trace!(?token, "attempting to auto-detect approval slot");
 
         {
-            let mut cache = self.cache.lock().unwrap();
-            if let Some(strategy) = cache.cache_get(&(token, None)) {
+            if let Some(strategy_opt) = self.cache.get(&(token, None)) {
                 tracing::trace!(?token, "cache hit (strategy valid for all pairs)");
-                return strategy.clone();
+                return strategy_opt;
             }
-            if let Some(strategy) = cache.cache_get(&(token, Some((owner, spender)))) {
+            if let Some(strategy_opt) = self.cache.get(&(token, Some((owner, spender)))) {
                 tracing::trace!(
                     ?token,
                     ?owner,
                     ?spender,
                     "cache hit (pair-specific strategy)"
                 );
-                return strategy.clone();
+                return strategy_opt;
             }
         }
 
@@ -256,15 +255,9 @@ impl Detector {
                     token,
                     (!strategy.is_valid_for_all_pairs()).then_some((owner, spender)),
                 );
-                self.cache
-                    .lock()
-                    .unwrap()
-                    .cache_set(cache_key, Some(strategy.clone()));
+                self.cache.insert(cache_key, Some(strategy.clone()));
             } else {
-                self.cache
-                    .lock()
-                    .unwrap()
-                    .cache_set((token, Some((owner, spender))), None);
+                self.cache.insert((token, Some((owner, spender))), None);
             }
         } else {
             tracing::warn!(?token, ?result, "error auto-detecting approval strategy");

crates/balance-overrides/src/approval/mod.rs

@@ -1,7 +1,10 @@
 pub(crate) mod aave;
 
 use {
-    crate::detector::{DetectionError, SimulationError, extract_sload_slots, mapping_slot_hash},
+    crate::{
+        cache::Cache,
+        detector::{DetectionError, SimulationError, extract_sload_slots, mapping_slot_hash},
+    },
     alloy_eips::BlockId,
     alloy_primitives::{Address, B256, TxKind, U256, keccak256, map::AddressMap},
     alloy_provider::ext::DebugApi,
@@ -13,10 +16,9 @@ use {
     },
     alloy_sol_types::SolCall,
     alloy_transport::TransportErrorKind,
-    cached::{Cached, SizedCache},
     contracts::ERC20,
     ethrpc::Web3,
-    std::{collections::HashMap, iter, sync::Mutex, time::Duration},
+    std::{collections::HashMap, iter, time::Duration},
 };
 
 /// These are the solady magic bytes for user balances
@@ -228,8 +230,6 @@ impl PartialEq for Strategy {
 
 impl Eq for Strategy {}
 
-type Cache = SizedCache<(Address, Option<Address>), Option<Strategy>>;
-
 /// Heuristic balance override detector with integrated caching.
 ///
 /// Owns the Web3 handle, detection parameters, and the per-token strategy
@@ -239,7 +239,7 @@ pub(crate) struct Detector {
     web3: Web3,
     probing_depth: u8,
     verification_timeout: Duration,
-    pub(crate) cache: Mutex<Cache>,
+    pub(crate) cache: Cache<(Address, Option<Address>), Option<Strategy>>,
 }
 
 impl Detector {
@@ -253,7 +253,7 @@ impl Detector {
             web3,
             probing_depth,
             verification_timeout,
-            cache: Mutex::new(SizedCache::with_size(cache_size)),
+            cache: Cache::new(u64::try_from(cache_size).expect("cache_size must be non-negative")),
         }
     }
 
@@ -263,14 +263,13 @@ impl Detector {
         tracing::trace!(?token, "attempting to auto-detect balance slot");
 
         {
-            let mut cache = self.cache.lock().unwrap();
-            if let Some(strategy) = cache.cache_get(&(token, None)) {
+            if let Some(strategy_opt) = self.cache.get(&(token, None)) {
                 tracing::trace!(?token, "cache hit (strategy valid for all holders)");
-                return strategy.clone();
+                return strategy_opt;
             }
-            if let Some(strategy) = cache.cache_get(&(token, Some(holder))) {
+            if let Some(strategy_opt) = self.cache.get(&(token, Some(holder))) {
                 tracing::trace!(?token, ?holder, "cache hit (holder-specific strategy)");
-                return strategy.clone();
+                return strategy_opt;
             }
         }
 
@@ -283,17 +282,11 @@ impl Detector {
                     (!strategy.can_be_applied_to_any_holder()).then_some(holder),
                 );
                 tracing::debug!(?token, ?strategy, "caching auto-detected balance strategy");
-                self.cache
-                    .lock()
-                    .unwrap()
-                    .cache_set(cache_key, Some(strategy.clone()));
+                self.cache.insert(cache_key, Some(strategy.clone()));
             }
             Err(DetectionError::NotFound) => {
                 tracing::debug!(?token, "caching token as unsupported");
-                self.cache
-                    .lock()
-                    .unwrap()
-                    .cache_set((token, Some(holder)), None);
+                self.cache.insert((token, Some(holder)), None);
             }
             Err(err) => {
                 tracing::warn!(

crates/balance-overrides/src/balance/mod.rs

@@ -0,0 +1,26 @@
+/// Note: cloning shares the underlying cache instance.
+#[derive(Clone)]
+pub(crate) struct Cache<K, V> {
+    data: moka::sync::Cache<K, V>,
+}
+
+impl<K, V> Cache<K, V>
+where
+    K: std::hash::Hash + Eq + Send + Sync + 'static,
+    V: Clone + Send + Sync + 'static,
+{
+    pub(crate) fn new(max_capacity: u64) -> Self {
+        let data = moka::sync::Cache::builder()
+            .max_capacity(max_capacity)
+            .build();
+        Self { data }
+    }
+
+    pub(crate) fn get(&self, key: &K) -> Option<V> {
+        self.data.get(key)
+    }
+
+    pub(crate) fn insert(&self, key: K, value: V) {
+        self.data.insert(key, value)
+    }
+}