feat(balance): adding the balance response caching

src/handlers/balance.rs

@@ -26,6 +26,7 @@ pub const H160_EMPTY_ADDRESS: H160 = H160::repeat_byte(0xee);
 
 const PROVIDER_MAX_CALLS: usize = 2;
 const METADATA_CACHE_TTL: Duration = Duration::from_secs(60 * 60 * 24); // 1 day
+const BALANCE_CACHE_TTL: Duration = Duration::from_secs(10); // 10 seconds
 
 #[derive(Debug, Deserialize, Clone)]
 #[serde(rename_all = "camelCase")]
@@ -78,6 +79,10 @@ fn token_metadata_cache_key(caip10_token_address: &str) -> String {
     format!("token_metadata/{}", caip10_token_address)
 }
 
+fn address_balance_cache_key(address: &str) -> String {
+    format!("address_balance/{}", address)
+}
+
 pub async fn get_cached_metadata(
     cache: &Option<Arc<dyn KeyValueStorage<TokenMetadataCacheItem>>>,
     caip10_token_address: &str,
@@ -106,6 +111,34 @@ pub async fn set_cached_metadata(
     }
 }
 
+pub async fn get_cached_balance(
+    cache: &Option<Arc<dyn KeyValueStorage<BalanceResponseBody>>>,
+    address: &str,
+) -> Option<BalanceResponseBody> {
+    let cache = cache.as_ref()?;
+    cache
+        .get(&address_balance_cache_key(address))
+        .await
+        .unwrap_or(None)
+}
+
+pub async fn set_cached_balance(
+    cache: &Option<Arc<dyn KeyValueStorage<BalanceResponseBody>>>,
+    address: &str,
+    item: &BalanceResponseBody,
+) {
+    if let Some(cache) = cache {
+        cache
+            .set(
+                &address_balance_cache_key(address),
+                item,
+                Some(BALANCE_CACHE_TTL),
+            )
+            .await
+            .unwrap_or_else(|e| error!("Failed to set balance cache: {}", e));
+    }
+}
+
 pub async fn handler(
     state: State<Arc<AppState>>,
     query: Query<BalanceQueryParams>,
@@ -136,6 +169,13 @@ async fn handler_internal(
         return Ok(Json(BalanceResponseBody { balances: vec![] }).into_response());
     }
 
+    // Get the cached balance and return it if found except if force_update is needed
+    if query.force_update.is_none() {
+        if let Some(cached_balance) = get_cached_balance(&state.balance_cache, &address).await {
+            return Ok(Json(cached_balance).into_response());
+        }
+    }
+
     // If the namespace is not provided, then default to the Ethereum namespace
     let namespace = query
         .chain_id
@@ -342,5 +382,16 @@ async fn handler_internal(
         }
     }
 
+    // Spawn a background task to update the balance cache without blocking
+    {
+        tokio::spawn({
+            let address_key = address.clone();
+            let response = response.clone();
+            async move {
+                set_cached_balance(&state.balance_cache, &address_key, &response).await;
+            }
+        });
+    }
+
     Ok(Json(response).into_response())
 }

src/lib.rs

@@ -2,7 +2,9 @@ use {
     crate::{
         env::Config,
         handlers::{
-            balance::TokenMetadataCacheItem, identity::IdentityResponse, rate_limit_middleware,
+            balance::{BalanceResponseBody, TokenMetadataCacheItem},
+            identity::IdentityResponse,
+            rate_limit_middleware,
         },
         metrics::Metrics,
         project::Registry,
@@ -138,6 +140,12 @@ pub async fn bootstrap(config: Config) -> RpcResult<()> {
         .map(|addr| redis::Redis::new(&addr, config.storage.redis_max_connections))
         .transpose()?
         .map(|r| Arc::new(r) as Arc<dyn KeyValueStorage<TokenMetadataCacheItem> + 'static>);
+    let balance_cache = config
+        .storage
+        .project_data_redis_addr()
+        .map(|addr| redis::Redis::new(&addr, config.storage.redis_max_connections))
+        .transpose()?
+        .map(|r| Arc::new(r) as Arc<dyn KeyValueStorage<BalanceResponseBody> + 'static>);
 
     let providers = init_providers(&config.providers);
 
@@ -196,6 +204,7 @@ pub async fn bootstrap(config: Config) -> RpcResult<()> {
         irn_client,
         identity_cache,
         token_metadata_cache,
+        balance_cache,
     );
 
     let port = state.config.server.port;

src/state.rs

@@ -3,12 +3,14 @@ use {
         analytics::RPCAnalytics,
         env::Config,
         error::RpcError,
-        handlers::{balance::TokenMetadataCacheItem, identity::IdentityResponse},
+        handlers::{
+            balance::{BalanceResponseBody, TokenMetadataCacheItem},
+            identity::IdentityResponse,
+        },
         metrics::Metrics,
         project::Registry,
         providers::ProviderRepository,
-        storage::irn::Irn,
-        storage::KeyValueStorage,
+        storage::{irn::Irn, KeyValueStorage},
         utils::{build::CompileInfo, rate_limit::RateLimit},
     },
     cerberus::project::ProjectDataWithQuota,
@@ -37,6 +39,7 @@ pub struct AppState {
     // Redis caching
     pub identity_cache: Option<Arc<dyn KeyValueStorage<IdentityResponse>>>,
     pub token_metadata_cache: Option<Arc<dyn KeyValueStorage<TokenMetadataCacheItem>>>,
+    pub balance_cache: Option<Arc<dyn KeyValueStorage<BalanceResponseBody>>>,
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -52,6 +55,7 @@ pub fn new_state(
     irn: Option<Irn>,
     identity_cache: Option<Arc<dyn KeyValueStorage<IdentityResponse>>>,
     token_metadata_cache: Option<Arc<dyn KeyValueStorage<TokenMetadataCacheItem>>>,
+    balance_cache: Option<Arc<dyn KeyValueStorage<BalanceResponseBody>>>,
 ) -> AppState {
     AppState {
         config,
@@ -67,6 +71,7 @@ pub fn new_state(
         irn,
         identity_cache,
         token_metadata_cache,
+        balance_cache,
     }
 }