diff --git a/components/Balances.tsx b/components/Balances.tsx index daea26f..41b09a3 100644 --- a/components/Balances.tsx +++ b/components/Balances.tsx @@ -12,13 +12,13 @@ const Balances = () => { return quotient.toNumber() + remainder.toNumber() / divisor.toNumber() } - function calculateBalance(a) { - const mint = mints[a.account.mint.toBase58()] - return mint ? fixedPointToNumber(a.account.amount, mint.decimals) : 0 + function calculateBalance(acc) { + const mint = mints[acc.mint.toBase58()] + return mint ? fixedPointToNumber(acc.amount, mint.decimals) : 0 } - const displayedBalances = tokenAccounts - .map((a) => ({ id: a.publicKey.toBase58(), balance: calculateBalance(a) })) + const displayedBalances = Object.entries(tokenAccounts) + .map(([id, acc]) => ({ id, balance: calculateBalance(acc) })) .sort((a, b) => (a.id > b.id ? 1 : -1)) return ( diff --git a/hooks/useWallet.tsx b/hooks/useWallet.tsx index 6948ffb..83f44c2 100644 --- a/hooks/useWallet.tsx +++ b/hooks/useWallet.tsx @@ -127,9 +127,8 @@ export default function useWallet() { }, [wallet, setWalletStore]) useInterval(async () => { - await actions.fetchWalletTokenAccounts() await actions.fetchWalletMints() - }, 20 * SECONDS) + }, 120 * SECONDS) return { connected, wallet } } diff --git a/stores/useWalletStore.tsx b/stores/useWalletStore.tsx index 6c29f91..3381e35 100644 --- a/stores/useWalletStore.tsx +++ b/stores/useWalletStore.tsx @@ -6,7 +6,7 @@ import { EndpointInfo, WalletAdapter } from '../@types/types' import { getOwnedTokenAccounts, getMint, - ProgramAccount, + subscribeToTokenAccount, TokenAccount, MintAccount, } from '../utils/tokens' @@ -39,8 +39,9 @@ interface WalletStore extends State { } current: WalletAdapter | undefined providerUrl: string - tokenAccounts: ProgramAccount[] + tokenAccounts: { [pubkey: string]: TokenAccount } mints: { [pubkey: string]: MintAccount } + subscriptions: { [pubkey: string]: () => void } set: (x: any) => void actions: any } @@ -55,28 +56,43 @@ const useWalletStore = create((set, get) => ({ }, current: null, providerUrl: null, - tokenAccounts: [], + tokenAccounts: {}, mints: {}, + subscriptions: {}, actions: { async fetchWalletTokenAccounts() { const connection = get().connection.current const connected = get().connected + const subscriptions = get().subscriptions const wallet = get().current const walletOwner = wallet?.publicKey const set = get().set if (connected && walletOwner) { - const ownedTokenAccounts = await getOwnedTokenAccounts( + const ownedAccountsResult = await getOwnedTokenAccounts( connection, walletOwner ) + const newTokenAccounts: { [pubkey: string]: TokenAccount } = {} + ownedAccountsResult.forEach((r) => { + newTokenAccounts[r.publicKey.toBase58()] = r.account + }) + + // cancel all subscriptions + Object.values(subscriptions).forEach((s) => s()) + set((state) => { - state.tokenAccounts = ownedTokenAccounts + state.subscriptions = {} + state.tokenAccounts = newTokenAccounts + }) + + ownedAccountsResult.forEach((r) => { + this.subscribeToTokenAccount(r.publicKey) }) } else { set((state) => { - state.tokenAccounts = [] + state.tokenAccounts = {} }) } }, @@ -87,14 +103,14 @@ const useWalletStore = create((set, get) => ({ const set = get().set if (connected) { - const fetchMints = tokenAccounts.map((a) => - getMint(connection, a.account.mint) + const fetchMints = Object.values(tokenAccounts).map((a) => + getMint(connection, a.mint) ) const mintResults = await Promise.all(fetchMints) const newMints: { [pubkey: string]: MintAccount } = {} mintResults.forEach( - (m) => (newMints[m.publicKey.toBase58()] = m.account) + (r) => (newMints[r.publicKey.toBase58()] = r.account) ) set((state) => { @@ -106,6 +122,23 @@ const useWalletStore = create((set, get) => ({ }) } }, + async subscribeToTokenAccount(pubkey) { + const connection = get().connection.websocket + const connected = get().connected + const set = get().set + + if (connected) { + const sub = subscribeToTokenAccount(connection, pubkey, (r) => { + set((s) => { + s.tokenAccounts[pubkey.toBase58()] = r.account + }) + }) + + set((s) => { + s.subscriptions[pubkey.toBase58()] = sub + }) + } + }, }, set: (fn) => set(produce(fn)), })) diff --git a/utils/tokens.tsx b/utils/tokens.tsx index ab40271..5c9b6c2 100644 --- a/utils/tokens.tsx +++ b/utils/tokens.tsx @@ -9,7 +9,7 @@ import { export type TokenAccount = AccountInfo export type MintAccount = MintInfo -export type ProgramAccount = { +export type AccountResponse = { publicKey: PublicKey account: T } @@ -17,7 +17,7 @@ export type ProgramAccount = { export async function getOwnedTokenAccounts( connection: Connection, publicKey: PublicKey -): Promise[]> { +): Promise[]> { const results = await connection.getTokenAccountsByOwner(publicKey, { programId: TOKEN_PROGRAM_ID, }) @@ -32,7 +32,7 @@ export async function getOwnedTokenAccounts( export async function getMint( connection: Connection, publicKey: PublicKey -): Promise> { +): Promise> { const result = await connection.getAccountInfo(publicKey) const data = Buffer.from(result.data) const account = parseMintAccountData(data) @@ -42,6 +42,21 @@ export async function getMint( } } +export function subscribeToTokenAccount( + connection: Connection, + publicKey: PublicKey, + cb: (r: AccountResponse) => void +): () => void { + const id = connection.onAccountChange(publicKey, (info) => { + const data = Buffer.from(info.data) + const account = parseTokenAccountData(publicKey, data) + cb({ publicKey, account }) + }) + return () => { + connection.removeAccountChangeListener(id) + } +} + // copied from @solana/spl-token const TOKEN_PROGRAM_ID = new PublicKey(