Skip to content

Commit

Permalink
Token balance display fix (#1982)
Browse files Browse the repository at this point in the history
* Token balance display fix

Previously the token balances would only be shown after the TokensModal
was opened to fetch and set the wallet's balances. The useGetTokenBalance
hook is now used to fetch the balances on demand (if not already cached).

* TokensModal uses hook, renamed to useGetTokenBalances

* style thing

* small refactor
  • Loading branch information
kev1n-peters committed Apr 26, 2024
1 parent 18587e1 commit 88d1106
Show file tree
Hide file tree
Showing 6 changed files with 192 additions and 170 deletions.
166 changes: 25 additions & 141 deletions wormhole-connect/src/components/TokensModal.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -4,28 +4,14 @@ import Button from '@mui/material/Button';
import IconButton from '@mui/material/IconButton';
import OpenInNewIcon from '@mui/icons-material/OpenInNew';
import { useTheme } from '@mui/material/styles';
import { ChainName, TokenId } from '@wormhole-foundation/wormhole-connect-sdk';
import { ChainName } from '@wormhole-foundation/wormhole-connect-sdk';
import { AVAILABLE_MARKETS_URL } from 'config/constants';
import config from 'config';
import { TokenConfig } from 'config/types';
import { BigNumber } from 'ethers';
import TokenIcon from 'icons/TokenIcons';
import React, {
ChangeEvent,
useCallback,
useEffect,
useMemo,
useState,
} from 'react';
import { useDispatch, useSelector } from 'react-redux';
import React, { ChangeEvent, useEffect, useMemo, useState } from 'react';
import { useSelector } from 'react-redux';
import { RootState } from 'store';
import {
Balances,
ChainBalances,
accessChainBalances,
formatBalance,
setBalances,
} from 'store/transferInput';
import { makeStyles } from 'tss-react/mui';
import { displayAddress, sortTokens } from 'utils';
import { isGatewayChain } from 'utils/cosmos';
Expand All @@ -39,6 +25,8 @@ import Tabs from './Tabs';
import { CCTPManual_CHAINS } from '../routes/cctpManual';
import { isTBTCCanonicalChain } from 'routes/tbtc';
import { CHAIN_ID_ETH } from '@certusone/wormhole-sdk/lib/esm/utils';
import useGetTokenBalances from 'hooks/useGetTokenBalances';
import { Balances } from 'store/transferInput';

const useStyles = makeStyles()((theme: any) => ({
tokensContainer: {
Expand Down Expand Up @@ -149,7 +137,7 @@ const displayNativeChain = (token: TokenConfig): string => {

type DisplayTokensProps = {
tokens: TokenConfig[];
balances: any;
balances: Balances;
walletAddress: string | undefined;
chain: any;
selectToken: (tokenKey: string) => void;
Expand All @@ -176,8 +164,7 @@ function DisplayTokens(props: DisplayTokensProps) {

const showCircularProgress = (token: string): boolean => {
if (!chain || !walletAddress) return false;
if (!balances) return true;
if (balances && balances[token] !== null) return true;
if (balances[token]?.balance !== null) return true;
return false;
};

Expand Down Expand Up @@ -211,8 +198,8 @@ function DisplayTokens(props: DisplayTokensProps) {
<div className={classes.tokenRowRight}>
<div className={classes.tokenRowBalanceText}>Balance</div>
<div className={classes.tokenRowBalance}>
{balances && balances[token.key] && walletAddress ? (
<div>{balances[token.key]}</div>
{balances[token.key]?.balance && walletAddress ? (
<div>{balances[token.key].balance}</div>
) : showCircularProgress(token.key) ? (
<CircularProgress size={14} />
) : (
Expand Down Expand Up @@ -293,16 +280,12 @@ function isGatewayNativeToken(token: TokenConfig) {

function TokensModal(props: Props) {
const theme = useTheme();
const dispatch = useDispatch();
const { open, chain, walletAddress, type } = props;
const mobile = useMediaQuery(theme.breakpoints.down('sm'));
const [loading, setLoading] = useState(false);
const [balancesLoaded, setBalancesLoaded] = useState(false);
const [tokens, setTokens] = useState<TokenConfig[]>([]);
const [search, setSearch] = useState('');

const {
balances,
supportedSourceTokens,
supportedDestTokens,
allSupportedDestTokens: allSupportedDestTokensBase,
Expand All @@ -320,9 +303,14 @@ function TokensModal(props: Props) {
return supported.filter((t) => !isGatewayNativeToken(t));
}, [type, supportedSourceTokens, supportedDestTokens]);

const chainBalancesCache: ChainBalances | undefined = useMemo(() => {
return accessChainBalances(balances, walletAddress, chain);
}, [chain, balances, walletAddress]);
const queryTokens =
type === 'dest' ? allSupportedDestTokens : supportedTokens;

const { isFetching, balances } = useGetTokenBalances(
walletAddress || '',
chain,
queryTokens,
);

// search tokens
const handleSearch = (
Expand Down Expand Up @@ -376,110 +364,6 @@ function TokensModal(props: Props) {
closeTokensModal();
};

useEffect(() => {
setBalancesLoaded(false);
}, [chain, walletAddress]);

const getBalances = useCallback(async () => {
if (!walletAddress || !chain) return;
const fiveMinutesAgo = Date.now() - 60 * 1000 * 5;
if (
chainBalancesCache &&
chainBalancesCache.balances &&
chainBalancesCache.lastUpdated! > fiveMinutesAgo
) {
setBalancesLoaded(true);
return;
}

const queryTokens =
type === 'dest' ? allSupportedDestTokens : supportedTokens;
const nativeQueryToken = queryTokens.find(
(t) => !t.tokenId && t.nativeChain === chain,
);
const queryTokensWithIds = queryTokens.filter((t) => !!t.tokenId); // pre-filter so indexes line up
const tokenIds = queryTokensWithIds.reduce<TokenId[]>(
(tIds, t) => (t.tokenId ? [...tIds, t.tokenId] : tIds),
[],
);
let balances: Balances = {};
if (nativeQueryToken) {
let nativeBalance: BigNumber | null = null;
try {
nativeBalance = await config.wh.getNativeBalance(walletAddress, chain);
balances = {
...balances,
...formatBalance(chain, nativeQueryToken, nativeBalance),
};
} catch (e) {
console.warn('Failed to fetch native balance', e);
}
}
try {
const tokenBalances = await config.wh.getTokenBalances(
walletAddress,
tokenIds,
chain,
);
balances = tokenIds.reduce<Balances>(
(balances, tId, idx) => ({
...balances,
...formatBalance(chain, queryTokensWithIds[idx], tokenBalances[idx]),
}),
balances,
);
} catch (e) {
console.warn('Failed to fetch balances', e);
}

dispatch(
setBalances({
address: walletAddress,
chain,
balances,
}),
);
}, [
walletAddress,
chain,
dispatch,
type,
supportedTokens,
chainBalancesCache,
allSupportedDestTokens,
]);

// fetch token balances and set in store
useEffect(() => {
let active = true;
if (!walletAddress || !chain) {
setTokens(supportedTokens);
return;
}

if (!balancesLoaded) {
setLoading(true);
getBalances().finally(() => {
if (active) {
setLoading(false);
setBalancesLoaded(true);
}
});
}
return () => {
active = false;
};
}, [
walletAddress,
supportedTokens,
chain,
dispatch,
getBalances,
type,
open,
balancesLoaded,
]);

useEffect(() => {
// get tokens that exist on the chain and have a balance greater than 0
const filtered = supportedTokens.filter((t) => {
Expand All @@ -506,26 +390,26 @@ function TokensModal(props: Props) {
}

if (type === 'dest') return true;
if (!chainBalancesCache) return true;
const b = chainBalancesCache.balances[t.key];
const isNonzeroBalance = b !== null && b !== '0';
if (!balances[t.key]) return true;
const { balance } = balances[t.key];
const isNonzeroBalance = balance !== null && balance !== '0';
return isNonzeroBalance;
});
setTokens(filtered);
}, [chainBalancesCache, chain, supportedTokens, type]);
}, [balances, chain, supportedTokens, type]);

const tabs = [
{
label: 'Available Tokens',
panel: (
<DisplayTokens
tokens={displayedTokens}
balances={chainBalancesCache?.balances}
balances={balances}
walletAddress={walletAddress}
chain={chain}
selectToken={selectToken}
moreTokens={handleMoreTokens}
loading={loading}
loading={isFetching}
search={search}
/>
),
Expand All @@ -535,12 +419,12 @@ function TokensModal(props: Props) {
panel: (
<DisplayTokens
tokens={type === 'dest' ? allSupportedDestTokens : supportedTokens}
balances={chainBalancesCache?.balances}
balances={balances}
walletAddress={walletAddress}
chain={chain}
selectToken={selectToken}
moreTokens={handleMoreTokens}
loading={loading}
loading={isFetching}
search={search}
/>
),
Expand Down
127 changes: 127 additions & 0 deletions wormhole-connect/src/hooks/useGetTokenBalances.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
import { useDispatch, useSelector } from 'react-redux';
import { RootState } from 'store';
import { ChainName, TokenId } from '@wormhole-foundation/wormhole-connect-sdk';
import { useEffect, useState } from 'react';
import {
accessBalance,
Balances,
formatBalance,
updateBalances,
} from 'store/transferInput';
import config from 'config';
import { TokenConfig } from 'config/types';

const useGetTokenBalances = (
walletAddress: string,
chain: ChainName | undefined,
tokens: TokenConfig[],
): { isFetching: boolean; balances: Balances } => {
const [isFetching, setIsFetching] = useState(false);
const [balances, setBalances] = useState<Balances>({});
const cachedBalances = useSelector(
(state: RootState) => state.transferInput.balances,
);
const dispatch = useDispatch();

useEffect(() => {
setIsFetching(true);
setBalances({});
if (
!walletAddress ||
!chain ||
!config.chains[chain] ||
tokens.length === 0
) {
setIsFetching(false);
return;
}
const chainConfig = config.chains[chain];
if (!chainConfig) {
setIsFetching(false);
return;
}

let isActive = true;

const getBalances = async () => {
const balances: Balances = {};
type TokenConfigWithId = TokenConfig & { tokenId: TokenId };
const needsUpdate: TokenConfigWithId[] = [];
const now = Date.now();
const fiveMinutesAgo = now - 5 * 60 * 1000;
let updateCache = false;
for (const token of tokens) {
const cachedBalance = accessBalance(
cachedBalances,
walletAddress,
chain,
token.key,
);
if (cachedBalance && cachedBalance.lastUpdated > fiveMinutesAgo) {
balances[token.key] = cachedBalance;
} else {
if (token.key === chainConfig.gasToken) {
try {
const balance = await config.wh.getNativeBalance(
walletAddress,
chain,
token.key,
);
balances[token.key] = {
balance: formatBalance(chain, token, balance),
lastUpdated: now,
};
updateCache = true;
} catch (e) {
console.error('Failed to get native balance', e);
}
} else if (token.tokenId) {
needsUpdate.push(token as TokenConfigWithId);
}
}
}
if (needsUpdate.length > 0) {
try {
const result = await config.wh.getTokenBalances(
walletAddress,
needsUpdate.map((t) => t.tokenId),
chain,
);
result.forEach((balance, i) => {
const token = needsUpdate[i];
balances[token.key] = {
balance: formatBalance(chain, token, balance),
lastUpdated: now,
};
});
updateCache = true;
} catch (e) {
console.error('Failed to get token balances', e);
}
}
if (isActive) {
setIsFetching(false);
setBalances(balances);
if (updateCache) {
dispatch(
updateBalances({
address: walletAddress,
chain,
balances,
}),
);
}
}
};

getBalances();

return () => {
isActive = false;
};
}, [cachedBalances, walletAddress, chain, tokens]);

return { isFetching, balances };
};

export default useGetTokenBalances;
Loading

0 comments on commit 88d1106

Please sign in to comment.