Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 4 additions & 82 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion crates/balance-overrides/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
33 changes: 13 additions & 20 deletions crates/balance-overrides/src/approval/mod.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -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")),
}
}

Expand All @@ -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;
}
}

Expand All @@ -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");
Expand Down
33 changes: 13 additions & 20 deletions crates/balance-overrides/src/balance/mod.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -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")),
}
}

Expand All @@ -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;
}
}

Expand All @@ -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!(
Expand Down
26 changes: 26 additions & 0 deletions crates/balance-overrides/src/cache.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
Loading
Loading