From 0f367d903515016fe8c1017a0fde20a55a3b22b3 Mon Sep 17 00:00:00 2001
From: Bryan Stitt <bryan@llamanodes.com>
Date: Thu, 18 May 2023 00:04:30 -0700
Subject: [PATCH] more quick cache ttl

---
 Cargo.lock                       |  2 +-
 deferred-rate-limiter/Cargo.toml |  2 +-
 deferred-rate-limiter/src/lib.rs | 34 ++++++++++++++++++++------------
 web3_proxy/src/app/mod.rs        | 14 +++++--------
 4 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4292411b..43ccb657 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1310,7 +1310,7 @@ dependencies = [
  "anyhow",
  "hashbrown 0.13.2",
  "log",
- "moka",
+ "quick_cache_ttl",
  "redis-rate-limiter",
  "tokio",
 ]
diff --git a/deferred-rate-limiter/Cargo.toml b/deferred-rate-limiter/Cargo.toml
index 699c9927..fad3722f 100644
--- a/deferred-rate-limiter/Cargo.toml
+++ b/deferred-rate-limiter/Cargo.toml
@@ -5,10 +5,10 @@ authors = ["Bryan Stitt <bryan@stitthappens.com>"]
 edition = "2021"
 
 [dependencies]
+quick_cache_ttl = { path = "../quick_cache_ttl" }
 redis-rate-limiter = { path = "../redis-rate-limiter" }
 
 anyhow = "1.0.71"
 hashbrown = "0.13.2"
 log = "0.4.17"
-moka = { version = "0.11.0", default-features = false, features = ["future"] }
 tokio = "1.28.1"
diff --git a/deferred-rate-limiter/src/lib.rs b/deferred-rate-limiter/src/lib.rs
index f02c69f8..65f320a3 100644
--- a/deferred-rate-limiter/src/lib.rs
+++ b/deferred-rate-limiter/src/lib.rs
@@ -1,6 +1,6 @@
 //#![warn(missing_docs)]
 use log::error;
-use moka::future::Cache;
+use quick_cache_ttl::{CacheWithTTL, UnitWeighter};
 use redis_rate_limiter::{RedisRateLimitResult, RedisRateLimiter};
 use std::cmp::Eq;
 use std::fmt::{Debug, Display};
@@ -16,7 +16,8 @@ pub struct DeferredRateLimiter<K>
 where
     K: Send + Sync,
 {
-    local_cache: Cache<K, Arc<AtomicU64>, hashbrown::hash_map::DefaultHashBuilder>,
+    local_cache:
+        CacheWithTTL<K, Arc<AtomicU64>, UnitWeighter, hashbrown::hash_map::DefaultHashBuilder>,
     prefix: String,
     rrl: RedisRateLimiter,
     /// if None, defers to the max on rrl
@@ -33,9 +34,9 @@ impl<K> DeferredRateLimiter<K>
 where
     K: Copy + Debug + Display + Hash + Eq + Send + Sync + 'static,
 {
-    pub fn new(
+    pub async fn new(
         // TODO: change this to cache_size in bytes
-        cache_size: u64,
+        cache_size: usize,
         prefix: &str,
         rrl: RedisRateLimiter,
         default_max_requests_per_second: Option<u64>,
@@ -45,11 +46,18 @@ where
         // TODO: time to live is not exactly right. we want this ttl counter to start only after redis is down. this works for now
         // TODO: what do these weigh?
         // TODO: allow skipping max_capacity
-        let local_cache = Cache::builder()
-            .time_to_live(Duration::from_secs(ttl))
-            .max_capacity(cache_size)
-            .name(prefix)
-            .build_with_hasher(hashbrown::hash_map::DefaultHashBuilder::default());
+        let local_cache = CacheWithTTL::new(
+            cache_size,
+            cache_size as u64,
+            UnitWeighter,
+            hashbrown::hash_map::DefaultHashBuilder::default(),
+            Duration::from_secs(ttl),
+        )
+        .await;
+        // .time_to_live(Duration::from_secs(ttl))
+        // .max_capacity(cache_size)
+        // .name(prefix)
+        // .build_with_hasher(hashbrown::hash_map::DefaultHashBuilder::default());
 
         Self {
             local_cache,
@@ -87,9 +95,9 @@ where
             let redis_key = redis_key.clone();
             let rrl = Arc::new(self.rrl.clone());
 
-            // set arc_deferred_rate_limit_result and return the coun
+            // set arc_deferred_rate_limit_result and return the count
             self.local_cache
-                .get_with_by_ref(&key, async move {
+                .get_or_insert_async::<anyhow::Error, _>(&key, async move {
                     // we do not use the try operator here because we want to be okay with redis errors
                     let redis_count = match rrl
                         .throttle_label(&redis_key, Some(max_requests_per_period), count)
@@ -126,9 +134,9 @@ where
                         }
                     };
 
-                    Arc::new(AtomicU64::new(redis_count))
+                    Ok(Arc::new(AtomicU64::new(redis_count)))
                 })
-                .await
+                .await?
         };
 
         let mut locked = deferred_rate_limit_result.lock().await;
diff --git a/web3_proxy/src/app/mod.rs b/web3_proxy/src/app/mod.rs
index 70b0b6e4..ed77cfe8 100644
--- a/web3_proxy/src/app/mod.rs
+++ b/web3_proxy/src/app/mod.rs
@@ -573,15 +573,11 @@ impl Web3ProxyApp {
                 // these two rate limiters can share the base limiter
                 // these are deferred rate limiters because we don't want redis network requests on the hot path
                 // TODO: take cache_size from config
-                frontend_ip_rate_limiter = Some(DeferredRateLimiter::<IpAddr>::new(
-                    10_000,
-                    "ip",
-                    rpc_rrl.clone(),
-                    None,
-                ));
-                frontend_registered_user_rate_limiter = Some(DeferredRateLimiter::<u64>::new(
-                    10_000, "key", rpc_rrl, None,
-                ));
+                frontend_ip_rate_limiter = Some(
+                    DeferredRateLimiter::<IpAddr>::new(20_000, "ip", rpc_rrl.clone(), None).await,
+                );
+                frontend_registered_user_rate_limiter =
+                    Some(DeferredRateLimiter::<u64>::new(10_000, "key", rpc_rrl, None).await);
             }
 
             // login rate limiter