From 1c6064cd42a9fd7ae41dde796e6460d82b76ec5d Mon Sep 17 00:00:00 2001 From: Denis Fadeev Date: Sun, 14 Jul 2024 17:01:34 +0300 Subject: [PATCH] refactor(getBalances): Use multicall3 to reduce numer of RPC calls --- packages/client/src/getBalances.ts | 204 ++++++++++++++++++--------- packages/tasks/src/account.ts | 2 +- packages/tasks/src/balances.ts | 2 +- packages/tasks/src/bitcoinAddress.ts | 9 +- 4 files changed, 142 insertions(+), 75 deletions(-) diff --git a/packages/client/src/getBalances.ts b/packages/client/src/getBalances.ts index 995e40a2..03ca9b19 100644 --- a/packages/client/src/getBalances.ts +++ b/packages/client/src/getBalances.ts @@ -20,6 +20,29 @@ export interface TokenBalance { zrc20?: string; } +const MULTICALL3_ABI = [ + { + inputs: [ + { + components: [ + { internalType: "address", name: "target", type: "address" }, + { internalType: "bytes", name: "callData", type: "bytes" }, + ], + internalType: "struct IMulticall3.Call[]", + name: "calls", + type: "tuple[]", + }, + ], + name: "aggregate", + outputs: [ + { internalType: "uint256", name: "blockNumber", type: "uint256" }, + { internalType: "bytes[]", name: "returnData", type: "bytes[]" }, + ], + stateMutability: "view", + type: "function", + }, +]; + /** * Get token balances of all tokens on all chains connected to ZetaChain. * @@ -105,80 +128,123 @@ export const getBalances = async function ( }) .filter((token: any) => token.chain_name); - const balances = await Promise.all( - tokens.map(async (token: any) => { - const isGas = token.coin_type === "Gas"; - const isBitcoin = ["btc_testnet", "btc_mainnet"].includes( - token.chain_name + const multicallAddress = "0xca11bde05977b3631167028862be2a173976ca11"; + + const multicallContexts: Record = {}; + + tokens.forEach((token: any) => { + if (token.coin_type === "ERC20" || token.coin_type === "ZRC20") { + if (!multicallContexts[token.chain_name]) { + multicallContexts[token.chain_name] = []; + } + multicallContexts[token.chain_name].push({ + target: token.contract, + callData: new ethers.utils.Interface( + token.coin_type === "ERC20" ? ERC20_ABI.abi : ZRC20.abi + ).encodeFunctionData("balanceOf", [evmAddress]), + }); + } + }); + + const balances: TokenBalance[] = []; + + await Promise.all( + Object.keys(multicallContexts).map(async (chainName) => { + const rpc = await this.getEndpoint("evm", chainName); + const provider = new ethers.providers.StaticJsonRpcProvider(rpc); + const multicallContract = new ethers.Contract( + multicallAddress, + MULTICALL3_ABI, + provider ); - const isERC = token.coin_type === "ERC20"; - const isZRC = token.coin_type === "ZRC20"; - if (isGas && !isBitcoin) { - let rpc; - try { - rpc = await this.getEndpoint("evm", token.chain_name); - } catch (e) { - return token; - } - const provider = new ethers.providers.StaticJsonRpcProvider(rpc); - return provider.getBalance(evmAddress).then((balance) => { - return { ...token, balance: formatUnits(balance, token.decimals) }; - }); - } else if (isGas && isBitcoin && btcAddress) { - let API; - try { - API = this.getEndpoint("esplora", token.chain_name); - } catch (e) { - return { ...token, balance: undefined }; - } - return fetch(`${API}/address/${btcAddress}`).then(async (response) => { - const r = await response.json(); - const { funded_txo_sum, spent_txo_sum } = r.chain_stats; - const balance = ( - (funded_txo_sum - spent_txo_sum) / - 100000000 - ).toString(); - return { ...token, balance }; - }); - } else if (isERC) { - let rpc; - try { - rpc = await this.getEndpoint("evm", token.chain_name); - } catch (e) { - return token; - } - const provider = new ethers.providers.StaticJsonRpcProvider(rpc); - const contract = new ethers.Contract( - token.contract, - ERC20_ABI.abi, - provider - ); - const decimals = await contract.decimals(); - return contract.balanceOf(evmAddress).then((balance: string) => { - return { - ...token, - balance: formatUnits(balance, decimals), - decimals, - }; - }); - } else if (isZRC) { - const rpc = await this.getEndpoint("evm", token.chain_name); - const provider = new ethers.providers.StaticJsonRpcProvider(rpc); - const contract = new ethers.Contract( - token.contract, - ZRC20.abi, - provider + + const calls = multicallContexts[chainName]; + + try { + const { returnData } = await multicallContract.callStatic.aggregate( + calls ); - return contract.balanceOf(evmAddress).then((balance: string) => { - return { - ...token, - balance: formatUnits(balance, token.decimals), - }; + + returnData.forEach((data: any, index: number) => { + const token = tokens.find( + (t) => + t.chain_name === chainName && + (t.coin_type === "ERC20" || t.coin_type === "ZRC20") && + t.contract === calls[index].target + ); + if (token) { + const balance = ethers.utils.defaultAbiCoder.decode( + ["uint256"], + data + )[0]; + const formattedBalance = formatUnits(balance, token.decimals); + balances.push({ ...token, balance: formattedBalance }); + } }); - } else { - return Promise.resolve(token); + } catch (error) { + console.error(`Multicall failed for ${chainName}:`, error); + // Fallback to individual calls if multicall fails + for (const token of tokens.filter( + (t) => + t.chain_name === chainName && + (t.coin_type === "ERC20" || t.coin_type === "ZRC20") + )) { + try { + const contract = new ethers.Contract( + token.contract, + token.coin_type === "ERC20" ? ERC20_ABI.abi : ZRC20.abi, + provider + ); + const balance = await contract.balanceOf(evmAddress); + const formattedBalance = formatUnits(balance, token.decimals); + balances.push({ ...token, balance: formattedBalance }); + } catch (err) { + console.error( + `Failed to get balance for ${token.symbol} on ${chainName}:`, + err + ); + } + } } }) ); + + await Promise.all( + tokens + .filter( + (token) => + token.coin_type === "Gas" && + !["btc_testnet", "btc_mainnet"].includes(token.chain_name) + ) + .map(async (token) => { + const rpc = await this.getEndpoint("evm", token.chain_name); + const provider = new ethers.providers.StaticJsonRpcProvider(rpc); + const balance = await provider.getBalance(evmAddress); + const formattedBalance = formatUnits(balance, token.decimals); + balances.push({ ...token, balance: formattedBalance }); + }) + ); + + await Promise.all( + tokens + .filter( + (token) => + token.coin_type === "Gas" && + ["btc_testnet", "btc_mainnet"].includes(token.chain_name) && + btcAddress + ) + .map(async (token) => { + const API = this.getEndpoint("esplora", token.chain_name); + const response = await fetch(`${API}/address/${btcAddress}`); + const r = await response.json(); + const { funded_txo_sum, spent_txo_sum } = r.chain_stats; + const balance = ( + (funded_txo_sum - spent_txo_sum) / + 100000000 + ).toString(); + balances.push({ ...token, balance }); + }) + ); + return balances; }; diff --git a/packages/tasks/src/account.ts b/packages/tasks/src/account.ts index b9fb86c4..3469ca72 100644 --- a/packages/tasks/src/account.ts +++ b/packages/tasks/src/account.ts @@ -70,7 +70,7 @@ export const main = async (args: any, hre: HardhatRuntimeEnvironment) => { 🔑 Private key: ${pk}`); mnemonic && console.log(`🔐 Mnemonic phrase: ${mnemonic.phrase}`); console.log(`😃 EVM address: ${address} -😃 Bitcoin address: ${bitcoinAddress(pk)} +😃 Bitcoin address: ${bitcoinAddress(pk, "testnet")} 😃 Bech32 address: ${hexToBech32Address(address, "zeta")} `); diff --git a/packages/tasks/src/balances.ts b/packages/tasks/src/balances.ts index 53257dd7..eb605542 100644 --- a/packages/tasks/src/balances.ts +++ b/packages/tasks/src/balances.ts @@ -47,7 +47,7 @@ const main = async (args: any, hre: HardhatRuntimeEnvironment) => { evmAddress = args.address; } else if (pk) { evmAddress = new ethers.Wallet(pk).address; - btcAddress = bitcoinAddress(pk); + btcAddress = bitcoinAddress(pk, args.mainnet ? "mainnet" : "testnet"); } else { spinner.stop(); console.error(walletError + balancesError); diff --git a/packages/tasks/src/bitcoinAddress.ts b/packages/tasks/src/bitcoinAddress.ts index 0f21f5a1..9ddd53e2 100644 --- a/packages/tasks/src/bitcoinAddress.ts +++ b/packages/tasks/src/bitcoinAddress.ts @@ -2,15 +2,16 @@ import * as bitcoin from "bitcoinjs-lib"; import ECPairFactory from "ecpair"; import * as ecc from "tiny-secp256k1"; -export const bitcoinAddress = (pk: string) => { - const TESTNET = bitcoin.networks.testnet; +export const bitcoinAddress = (pk: string, network: "testnet" | "mainnet") => { + const bitcoinNetwork = + network === "testnet" ? bitcoin.networks.testnet : bitcoin.networks.bitcoin; const ECPair = ECPairFactory(ecc); const key = ECPair.fromPrivateKey(Buffer.from(pk, "hex"), { - network: TESTNET, + network: bitcoinNetwork, }); const { address } = bitcoin.payments.p2wpkh({ - network: TESTNET, + network: bitcoinNetwork, pubkey: key.publicKey, }); if (!address) throw new Error("Unable to generate bitcoin address");