feat(balance): implementing tokens metadata caching between Zerion and Dune providers (#906)

* feat(balance): adding tokens metadata caching between Zerion and Dune

* fix: decreasing the cache to 1 day

* fix: setting the cache async

* fix: don't fall to eip155:1 in Zerion when the chainId not found in converter

Author: geekbrother

src/handlers/balance.rs

@@ -4,6 +4,7 @@ use {
         analytics::{BalanceLookupInfo, MessageSource},
         error::RpcError,
         state::AppState,
+        storage::KeyValueStorage,
         utils::{crypto, network},
     },
     axum::{
@@ -14,13 +15,17 @@ use {
     ethers::{abi::Address, types::H160},
     hyper::HeaderMap,
     serde::{Deserialize, Serialize},
-    std::{net::SocketAddr, sync::Arc},
+    std::{net::SocketAddr, sync::Arc, time::Duration},
     tap::TapFallible,
     tracing::log::{debug, error},
     wc::future::FutureExt,
 };
 
+// Empty address for the contract address mimicking the Ethereum native token
+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
 
 #[derive(Debug, Deserialize, Clone)]
 #[serde(rename_all = "camelCase")]
@@ -61,6 +66,46 @@ pub struct BalanceQuantity {
     pub numeric: String,
 }
 
+#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct TokenMetadataCacheItem {
+    pub name: String,
+    pub symbol: String,
+    pub icon_url: String,
+}
+
+fn token_metadata_cache_key(caip10_token_address: &str) -> String {
+    format!("token_metadata/{}", caip10_token_address)
+}
+
+pub async fn get_cached_metadata(
+    cache: &Option<Arc<dyn KeyValueStorage<TokenMetadataCacheItem>>>,
+    caip10_token_address: &str,
+) -> Option<TokenMetadataCacheItem> {
+    let cache = cache.as_ref()?;
+    cache
+        .get(&token_metadata_cache_key(caip10_token_address))
+        .await
+        .unwrap_or(None)
+}
+
+pub async fn set_cached_metadata(
+    cache: &Option<Arc<dyn KeyValueStorage<TokenMetadataCacheItem>>>,
+    caip10_token_address: &str,
+    item: &TokenMetadataCacheItem,
+) {
+    if let Some(cache) = cache {
+        cache
+            .set(
+                &token_metadata_cache_key(caip10_token_address),
+                item,
+                Some(METADATA_CACHE_TTL),
+            )
+            .await
+            .unwrap_or_else(|e| error!("Failed to set metadata cache: {}", e));
+    }
+}
+
 pub async fn handler(
     state: State<Arc<AppState>>,
     query: Query<BalanceQueryParams>,
@@ -113,7 +158,12 @@ async fn handler_internal(
     let mut balance_response = None;
     for provider in providers.iter() {
         let provider_response = provider
-            .get_balance(address.clone(), query.clone().0, state.metrics.clone())
+            .get_balance(
+                address.clone(),
+                query.clone().0,
+                &state.token_metadata_cache,
+                state.metrics.clone(),
+            )
             .await
             .tap_err(|e| {
                 error!("Failed to call balance with {}", e);
@@ -170,7 +220,6 @@ async fn handler_internal(
         if namespace != crypto::CaipNamespaces::Eip155 {
             return Err(RpcError::UnsupportedNamespace(namespace));
         }
-        const H160_EMPTY_ADDRESS: H160 = H160::repeat_byte(0xee);
         let rpc_project_id = state
             .config
             .server

src/lib.rs

@@ -1,7 +1,9 @@
 use {
     crate::{
         env::Config,
-        handlers::{identity::IdentityResponse, rate_limit_middleware},
+        handlers::{
+            balance::TokenMetadataCacheItem, identity::IdentityResponse, rate_limit_middleware,
+        },
         metrics::Metrics,
         project::Registry,
         providers::ProvidersConfig,
@@ -130,6 +132,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<IdentityResponse> + 'static>);
+    let token_metadata_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<TokenMetadataCacheItem> + 'static>);
 
     let providers = init_providers(&config.providers);
 
@@ -187,11 +195,12 @@ pub async fn bootstrap(config: Config) -> RpcResult<()> {
         providers,
         metrics.clone(),
         registry,
-        identity_cache,
         analytics,
         http_client,
         rate_limiting,
         irn_client,
+        identity_cache,
+        token_metadata_cache,
     );
 
     let port = state.config.server.port;

src/providers/dune.rs

@@ -3,11 +3,15 @@ use {
     crate::{
         env::DuneConfig,
         error::{RpcError, RpcResult},
-        handlers::balance::{BalanceQueryParams, BalanceResponseBody},
+        handlers::balance::{
+            get_cached_metadata, set_cached_metadata, BalanceQueryParams, BalanceResponseBody,
+            TokenMetadataCacheItem, H160_EMPTY_ADDRESS,
+        },
         providers::{
             balance::{BalanceItem, BalanceQuantity},
             ProviderKind,
         },
+        storage::KeyValueStorage,
         utils::crypto,
         Metrics,
     },
@@ -153,6 +157,7 @@ impl BalanceProvider for DuneProvider {
         &self,
         address: String,
         params: BalanceQueryParams,
+        metadata_cache: &Option<Arc<dyn KeyValueStorage<TokenMetadataCacheItem>>>,
         metrics: Arc<Metrics>,
     ) -> RpcResult<BalanceResponseBody> {
         let namespace = params
@@ -175,76 +180,124 @@ impl BalanceProvider for DuneProvider {
             }
         };
 
-        let balances_vec = balance_response
-            .balances
-            .into_iter()
-            .filter_map(|mut f| {
-                // Skip the asset if there are no symbol, decimals, since this
-                // is likely a spam token
-                let symbol = f.symbol.take()?;
-                let price_usd = f.price_usd.take()?;
-                let decimals = f.decimals.take()?;
-                let caip2_chain_id = match f.chain_id {
-                    Some(cid) => format!("{}:{}", namespace, cid),
-                    None => match namespace {
-                        // Using defaul Mainnet chain ids if not provided since
-                        // Dune doesn't provide balances for testnets
-                        crypto::CaipNamespaces::Eip155 => format!("{}:{}", namespace, "1"),
-                        crypto::CaipNamespaces::Solana => {
-                            format!("{}:{}", namespace, "5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp")
+        let mut balances_vec = Vec::new();
+        for f in balance_response.balances {
+            // Skip if missing required fields as a possible spam token
+            let (Some(symbol), Some(price_usd), Some(decimals)) =
+                (f.symbol, f.price_usd, f.decimals)
+            else {
+                continue;
+            };
+
+            // Build a CAIP-2 chain ID
+            let caip2_chain_id = match f.chain_id {
+                Some(cid) => format!("{}:{}", namespace, cid),
+                None => match namespace {
+                    // Use default Mainnet chain IDs if not provided
+                    crypto::CaipNamespaces::Eip155 => format!("{}:1", namespace),
+                    crypto::CaipNamespaces::Solana => {
+                        format!("{}:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp", namespace)
+                    }
+                },
+            };
+
+            // Determine name
+            let name = if f.address == "native" {
+                f.chain.clone()
+            } else {
+                symbol.clone()
+            };
+
+            // Determine icon URL
+            let icon_url = if f.address == "native" {
+                NATIVE_TOKEN_ICONS.get(&symbol).unwrap_or(&"").to_string()
+            } else {
+                // If there's no token_metadata or no logo, skip
+                match &f.token_metadata {
+                    Some(m) => m.logo.clone(),
+                    None => continue,
+                }
+            };
+
+            // Build the CAIP-10 address
+            let caip10_token_address_strict = if f.address == "native" {
+                match namespace {
+                    crypto::CaipNamespaces::Eip155 => {
+                        format!("{}:{}", caip2_chain_id, H160_EMPTY_ADDRESS)
+                    }
+                    crypto::CaipNamespaces::Solana => {
+                        format!("{}:{}", caip2_chain_id, crypto::SOLANA_NATIVE_TOKEN_ADDRESS)
+                    }
+                }
+            } else {
+                format!("{}:{}", caip2_chain_id, f.address)
+            };
+
+            // Get token metadata from the cache or update it
+            let token_metadata =
+                match get_cached_metadata(metadata_cache, &caip10_token_address_strict).await {
+                    Some(cached) => cached,
+                    None => {
+                        let new_item = TokenMetadataCacheItem {
+                            name: name.clone(),
+                            symbol: symbol.clone(),
+                            icon_url: icon_url.clone(),
+                        };
+                        // Spawn a background task to set the cache without blocking
+                        {
+                            let metadata_cache = metadata_cache.clone();
+                            let address_key = caip10_token_address_strict.clone();
+                            let new_item_to_store = new_item.clone();
+                            tokio::spawn(async move {
+                                set_cached_metadata(
+                                    &metadata_cache,
+                                    &address_key,
+                                    &new_item_to_store,
+                                )
+                                .await;
+                            });
                         }
-                    },
+                        new_item
+                    }
                 };
-                Some(BalanceItem {
-                    name: {
-                        if f.address == "native" {
-                            f.chain
-                        } else {
-                            symbol.clone()
-                        }
-                    },
-                    symbol: symbol.clone(),
-                    chain_id: Some(caip2_chain_id.clone()),
-                    address: {
-                        // Return None if the address is native for the native token
-                        if f.address == "native" {
-                            // UI expecting `None`` for the Eip155 and Solana's native
-                            // token address for Solana
-                            match namespace {
-                                crypto::CaipNamespaces::Eip155 => None,
-                                crypto::CaipNamespaces::Solana => {
-                                    Some(crypto::SOLANA_NATIVE_TOKEN_ADDRESS.to_string())
-                                }
+
+            // Construct the final BalanceItem
+            let balance_item = BalanceItem {
+                name: token_metadata.name,
+                symbol: token_metadata.symbol,
+                chain_id: Some(caip2_chain_id.clone()),
+                address: {
+                    // Return None if the address is native (for EIP-155).
+                    // For Solana’s “native” token, we return the Solana native token address.
+                    if f.address == "native" {
+                        match namespace {
+                            crypto::CaipNamespaces::Eip155 => None,
+                            crypto::CaipNamespaces::Solana => {
+                                Some(crypto::SOLANA_NATIVE_TOKEN_ADDRESS.to_string())
                             }
-                        } else {
-                            Some(format!("{}:{}", caip2_chain_id, f.address.clone()))
-                        }
-                    },
-                    value: f.value_usd,
-                    price: price_usd,
-                    quantity: BalanceQuantity {
-                        decimals: decimals.to_string(),
-                        numeric: crypto::format_token_amount(
-                            U256::from_dec_str(&f.amount).unwrap_or_default(),
-                            decimals,
-                        ),
-                    },
-                    icon_url: {
-                        if f.address == "native" {
-                            NATIVE_TOKEN_ICONS.get(&symbol).unwrap_or(&"").to_string()
-                        } else {
-                            f.token_metadata?.logo
                         }
-                    },
-                })
-            })
-            .collect::<Vec<_>>();
+                    } else {
+                        Some(format!("{}:{}", caip2_chain_id, f.address))
+                    }
+                },
+                value: f.value_usd,
+                price: price_usd,
+                quantity: BalanceQuantity {
+                    decimals: decimals.to_string(),
+                    numeric: crypto::format_token_amount(
+                        U256::from_dec_str(&f.amount).unwrap_or_default(),
+                        decimals,
+                    ),
+                },
+                icon_url: token_metadata.icon_url,
+            };
 
-        let response = BalanceResponseBody {
-            balances: balances_vec,
-        };
+            balances_vec.push(balance_item);
+        }
 
-        Ok(response)
+        Ok(BalanceResponseBody {
+            balances: balances_vec,
+        })
     }
 }
 

src/providers/mod.rs

@@ -4,7 +4,7 @@ use {
         env::{BalanceProviderConfig, ProviderConfig},
         error::{RpcError, RpcResult},
         handlers::{
-            balance::{self, BalanceQueryParams, BalanceResponseBody},
+            balance::{self, BalanceQueryParams, BalanceResponseBody, TokenMetadataCacheItem},
             convert::{
                 allowance::{AllowanceQueryParams, AllowanceResponseBody},
                 approve::{ConvertApproveQueryParams, ConvertApproveResponseBody},
@@ -22,6 +22,7 @@ use {
             portfolio::{PortfolioQueryParams, PortfolioResponseBody},
             RpcQueryParams, SupportedCurrencies,
         },
+        storage::KeyValueStorage,
         utils::crypto::CaipNamespaces,
         Metrics,
     },
@@ -863,6 +864,7 @@ pub trait BalanceProvider: Send + Sync {
         &self,
         address: String,
         params: BalanceQueryParams,
+        metadata_cache: &Option<Arc<dyn KeyValueStorage<TokenMetadataCacheItem>>>,
         metrics: Arc<Metrics>,
     ) -> RpcResult<BalanceResponseBody>;
 }