use batch account balance in cron

across-protocol · Sep 12, 2024 · 1539f8b · 1539f8b
1 parent d62a6f9
commit 1539f8b
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 27 deletions.
diff --git a/api/_cache.ts b/api/_cache.ts
@@ -87,8 +87,10 @@ export function makeCacheGetterAndSetter<T>(
     get: async () => {
       return getCachedValue(key, ttl, fetcher, parser);
     },
-    set: async () => {
-      const value = await fetcher();
+    set: async (value?: T) => {
+      if (!value) {
+        value = await fetcher();
+      }
       await redisCache.set(key, value, ttl);
     },
   };

diff --git a/api/_utils.ts b/api/_utils.ts
@@ -1089,11 +1089,11 @@ export const getCachedTokenBalance = async (
   account: string,
   token: string
 ): Promise<BigNumber> => {
-  const balance = await latestBalanceCache(
-    Number(chainId),
-    token,
-    account
-  ).get();
+  const balance = await latestBalanceCache({
+    chainId: Number(chainId),
+    tokenAddress: token,
+    address: account,
+  }).get();
   return balance;
 };
 
@@ -1904,11 +1904,12 @@ export function getCachedLatestBlock(chainId: number) {
   );
 }
 
-export function latestBalanceCache(
-  chainId: number,
-  tokenAddress: string,
-  address: string
-) {
+export function latestBalanceCache(params: {
+  chainId: number;
+  tokenAddress: string;
+  address: string;
+}) {
+  const { chainId, tokenAddress, address } = params;
   const ttlPerChain = {
     default: 60,
     [CHAIN_IDs.MAINNET]: 60,

diff --git a/api/account-balance.ts b/api/account-balance.ts
@@ -31,11 +31,11 @@ const handler = async (
 
     let { token, account, chainId } = query;
 
-    const balance = await latestBalanceCache(
-      Number(chainId),
-      token,
-      account
-    ).get();
+    const balance = await latestBalanceCache({
+      chainId: Number(chainId),
+      tokenAddress: token,
+      address: account,
+    }).get();
     const result = {
       balance: balance.toString(),
       account: account,

diff --git a/api/cron-cache-balances.ts b/api/cron-cache-balances.ts
@@ -1,9 +1,10 @@
 import { VercelResponse } from "@vercel/node";
-import { ethers } from "ethers";
+import { BigNumber, ethers } from "ethers";
 import { TypedVercelRequest } from "./_types";
 
 import {
   HUB_POOL_CHAIN_ID,
+  getCachedTokenBalances,
   getLogger,
   handleErrorCondition,
   latestBalanceCache,
@@ -52,23 +53,40 @@ const handler = async (
       return;
     }
 
+    const allRelayers = [...fullRelayers, ...transferRestrictedRelayers];
     for (const chain of mainnetChains) {
-      for (const token of chain.inputTokens) {
-        const setTokenBalance = async (relayer: string) => {
-          await latestBalanceCache(chain.chainId, relayer, token.address).set();
-        };
-        await Promise.all([
-          Promise.all(fullRelayers.map(setTokenBalance)),
-          Promise.all(transferRestrictedRelayers.map(setTokenBalance)),
-        ]);
-      }
+      const batchResult = await getCachedTokenBalances(
+        chain.chainId,
+        allRelayers,
+        [
+          ethers.constants.AddressZero,
+          ...chain.outputTokens.map((token) => token.address),
+        ]
+      );
+
+      await Promise.all(
+        allRelayers.map((relayer) => {
+          return Promise.all(
+            chain.outputTokens.map((token) => {
+              return latestBalanceCache({
+                chainId: chain.chainId,
+                address: relayer,
+                tokenAddress: token.address,
+              }).set(
+                BigNumber.from(batchResult.balances[relayer][token.address])
+              );
+            })
+          );
+        })
+      );
     }
 
     logger.debug({
       at: "CronCacheBalances",
       message: "Finished",
     });
     response.status(200);
+    response.send("OK");
   } catch (error: unknown) {
     return handleErrorCondition("cron-cache-balances", response, logger, error);
   }