Subscribe to websocket for account changes

package.json

@@ -13,6 +13,7 @@
     "bip39": "^3.0.2",
     "bn.js": "^5.1.2",
     "buffer-layout": "^1.2.0",
+    "immutable-tuple": "^0.4.10",
     "react": "^16.13.1",
     "react-dom": "^16.13.1",
     "react-scripts": "3.4.1",

src/WalletPage.js

@@ -6,8 +6,9 @@ import { useWallet } from './utils/wallet';
 import { Account, LAMPORTS_PER_SOL } from '@solana/web3.js';
 import { createAndInitializeMint } from './utils/tokens';
 import Grid from '@material-ui/core/Grid';
-import { useIsProdNetwork } from './utils/connection';
+import { refreshAccountInfo, useIsProdNetwork } from './utils/connection';
 import { useUpdateTokenName } from './utils/tokens/names';
+import { sleep } from './utils/utils';
 
 export default function WalletPage() {
   const isProdNetwork = useIsProdNetwork();
@@ -35,11 +36,17 @@ function DevnetButtons() {
       <Button
         variant="contained"
         color="primary"
-        onClick={() => {
-          wallet.connection
-            .requestAirdrop(wallet.account.publicKey, LAMPORTS_PER_SOL)
-            .then(console.log)
-            .catch(console.warn);
+        onClick={async () => {
+          try {
+            await wallet.connection.requestAirdrop(
+              wallet.account.publicKey,
+              LAMPORTS_PER_SOL,
+            );
+            await sleep(1000);
+            refreshAccountInfo(wallet.connection, wallet.account.publicKey);
+          } catch (e) {
+            console.warn(e);
+          }
         }}
       >
         Request Airdrop
@@ -51,8 +58,8 @@ function DevnetButtons() {
           let mint = new Account();
           updateTokenName(
             mint.publicKey,
-            `Test Token ${mint.publicKey.toBase58().slice(40)}`,
-            `TEST${mint.publicKey.toBase58().slice(40)}`,
+            `Test Token ${mint.publicKey.toBase58().slice(0, 4)}`,
+            `TEST${mint.publicKey.toBase58().slice(0, 4)}`,
           );
           createAndInitializeMint({
             connection: wallet.connection,

src/components/BalancesList.js

@@ -27,6 +27,7 @@ import IconButton from '@material-ui/core/IconButton';
 import Tooltip from '@material-ui/core/Tooltip';
 import AddTokenDialog from './AddTokenDialog';
 import SendDialog from './SendDialog';
+import { refreshAccountInfo } from '../utils/connection';
 
 const balanceFormat = new Intl.NumberFormat(undefined, {
   minimumFractionDigits: 4,
@@ -35,9 +36,9 @@ const balanceFormat = new Intl.NumberFormat(undefined, {
 });
 
 export default function BalancesList() {
+  const wallet = useWallet();
   const accountCount = useWalletAccountCount();
   const [showAddTokenDialog, setShowAddTokenDialog] = useState(false);
-  const [refreshKey, setRefreshKey] = useState(0);
 
   return (
     <Paper>
@@ -53,7 +54,15 @@ export default function BalancesList() {
           </Tooltip>
           <Tooltip title="Refresh" arrow>
             <IconButton
-              onClick={() => setRefreshKey((i) => i + 1)}
+              onClick={() =>
+                [...Array(accountCount + 5).keys()].map((i) =>
+                  refreshAccountInfo(
+                    wallet.connection,
+                    wallet.getAccount(i).publicKey,
+                    true,
+                  ),
+                )
+              }
               style={{ marginRight: -12 }}
             >
               <RefreshIcon />
@@ -63,7 +72,7 @@ export default function BalancesList() {
       </AppBar>
       <List disablePadding>
         {[...Array(accountCount + 5).keys()].map((i) => (
-          <BalanceListItem key={i + ' ' + refreshKey} index={i} />
+          <BalanceListItem key={i} index={i} />
         ))}
       </List>
       <AddTokenDialog

src/utils/connection.js

@@ -1,6 +1,8 @@
-import React, { useContext, useMemo } from 'react';
+import React, { useContext, useEffect, useMemo } from 'react';
 import { clusterApiUrl, Connection } from '@solana/web3.js';
 import { useLocalStorageState } from './utils';
+import { refreshCache, useAsyncData } from './fetch-loop';
+import tuple from 'immutable-tuple';
 
 const ConnectionContext = React.createContext(null);
 
@@ -10,7 +12,9 @@ export function ConnectionProvider({ children }) {
     clusterApiUrl('devnet'),
   );
 
-  const connection = useMemo(() => new Connection(endpoint), [endpoint]);
+  const connection = useMemo(() => new Connection(endpoint, 'single'), [
+    endpoint,
+  ]);
 
   return (
     <ConnectionContext.Provider value={{ endpoint, setEndpoint, connection }}>
@@ -32,3 +36,27 @@ export function useIsProdNetwork() {
   const endpoint = useContext(ConnectionContext).endpoint;
   return endpoint === clusterApiUrl('mainnet-beta');
 }
+
+export function useAccountInfo(publicKey) {
+  const connection = useConnection();
+  const cacheKey = tuple(connection, publicKey?.toBase58());
+  const [accountInfo, loaded] = useAsyncData(
+    async () => (publicKey ? connection.getAccountInfo(publicKey) : null),
+    cacheKey,
+  );
+  useEffect(() => {
+    if (!publicKey) {
+      return () => {};
+    }
+    const id = connection.onAccountChange(publicKey, () =>
+      refreshCache(cacheKey),
+    );
+    return () => connection.removeAccountChangeListener(id);
+  }, [connection, publicKey?.toBase58(), cacheKey]);
+  return [accountInfo, loaded];
+}
+
+export function refreshAccountInfo(connection, publicKey, clearCache = false) {
+  const cacheKey = tuple(connection, publicKey.toBase58());
+  refreshCache(cacheKey, clearCache);
+}

src/utils/fetch-loop.js

@@ -168,11 +168,15 @@ export function useAsyncData(
   return [data, loaded];
 }
 
-export function refreshCache(cacheKey) {
-  return globalLoops.loops.get(cacheKey)?.refresh();
-}
-
-export function clearCache(cacheKey) {
-  globalCache.delete(cacheKey);
-  return globalLoops.loops.get(cacheKey)?.notifyListeners();
+export function refreshCache(cacheKey, clearCache = false) {
+  if (clearCache) {
+    globalCache.delete(cacheKey);
+  }
+  const loop = globalLoops.loops.get(cacheKey);
+  if (loop) {
+    loop.refresh();
+    if (clearCache) {
+      loop.notifyListeners();
+    }
+  }
 }

src/utils/wallet.js

@@ -1,8 +1,12 @@
-import React, { useContext, useEffect, useMemo, useState } from 'react';
+import React, { useContext, useEffect, useMemo } from 'react';
 import * as bip32 from 'bip32';
 import { Account } from '@solana/web3.js';
 import nacl from 'tweetnacl';
-import { useConnection } from './connection';
+import {
+  refreshAccountInfo,
+  useAccountInfo,
+  useConnection,
+} from './connection';
 import { createAndInitializeTokenAccount, transferTokens } from './tokens';
 import { TOKEN_PROGRAM_ID } from './tokens/instructions';
 import {
@@ -11,7 +15,7 @@ import {
   parseTokenAccountData,
 } from './tokens/data';
 import EventEmitter from 'events';
-import { sleep, useListener, useLocalStorageState } from './utils';
+import { useListener, useLocalStorageState } from './utils';
 import { useTokenName } from './tokens/names';
 
 export class Wallet {
@@ -36,46 +40,6 @@ export class Wallet {
     return Wallet.getAccountFromSeed(this.seed, this.walletIndex, index);
   };
 
-  getAccountBalance = async (index) => {
-    let publicKey = this.getAccount(index).publicKey;
-    let info = await this.connection.getAccountInfo(publicKey, 'single');
-
-    if (info && this.accountCount < index + 1) {
-      this.accountCount = index + 1;
-      this.emitter.emit('accountCountChange');
-    }
-
-    if (info?.owner.equals(TOKEN_PROGRAM_ID)) {
-      let { mint, owner, amount } = parseTokenAccountData(info.data);
-      if (!owner.equals(this.account.publicKey)) {
-        console.warn(
-          'token account %s not owned by wallet',
-          publicKey.toBase58(),
-        );
-        return null;
-      }
-      let mintInfo = await this.connection.getAccountInfo(mint, 'single');
-      let { decimals } = parseMintData(mintInfo.data);
-      return {
-        amount,
-        decimals,
-        mint,
-        tokenName: null,
-        tokenSymbol: null,
-        initialized: true,
-      };
-    }
-
-    return {
-      amount: info?.lamports ?? 0,
-      decimals: 9,
-      mint: null,
-      tokenName: 'Solana',
-      tokenSymbol: 'SOL',
-      initialized: false,
-    };
-  };
-
   createTokenAccount = async (tokenAddress) => {
     let index = this.accountCount;
     await createAndInitializeTokenAccount({
@@ -85,7 +49,7 @@ export class Wallet {
       newAccount: this.getAccount(index),
     });
     ++this.accountCount;
-    // TODO: clear cache
+    refreshAccountInfo(this.connection, this.getAccount(index).publicKey, true);
     this.emitter.emit('accountCountChange');
   };
 
@@ -144,35 +108,47 @@ export function useWalletAccountCount() {
   return wallet.accountCount;
 }
 
-export function useBalanceInfo(tokenIndex) {
+export function useBalanceInfo(index) {
   let wallet = useWallet();
-  let [info, setInfo] = useState(null);
-  useEffect(() => {
-    let cancelled = false;
-    (async () => {
-      let errors = 0;
-      while (!cancelled) {
-        try {
-          let info = await wallet.getAccountBalance(tokenIndex);
-          setInfo(info);
-          errors = 0;
-          await sleep(60000 * (1 + Math.random()));
-        } catch (e) {
-          ++errors;
-          await sleep(
-            1000 * Math.min(Math.pow(2, errors), 30) * (1 + Math.random()),
-          );
-        }
-      }
-    })();
-    return () => (cancelled = true);
-  }, [wallet, tokenIndex]);
-  let { name, symbol } = useTokenName(info?.mint);
+  let publicKey = wallet.getAccount(index).publicKey;
+  let [accountInfo, accountInfoLoaded] = useAccountInfo(publicKey);
+  let { mint, owner, amount } = accountInfo?.owner.equals(TOKEN_PROGRAM_ID)
+    ? parseTokenAccountData(accountInfo.data)
+    : {};
+  let [mintInfo, mintInfoLoaded] = useAccountInfo(mint);
+  let { name, symbol } = useTokenName(mint);
 
-  if (info?.mint && !info?.tokenSymbol) {
-    info = { ...info, tokenName: name, tokenSymbol: symbol };
+  useEffect(() => {
+    if (accountInfo && wallet.accountCount < index + 1) {
+      wallet.accountCount = index + 1;
+      wallet.emitter.emit('accountCountChange');
+    }
+  }, [wallet, accountInfo, index]);
+
+  if (accountInfoLoaded && mint && mintInfoLoaded) {
+    let { decimals } = parseMintData(mintInfo.data);
+    return {
+      amount,
+      decimals,
+      mint,
+      owner,
+      tokenName: name,
+      tokenSymbol: symbol,
+      initialized: true,
+    };
+  } else if (accountInfoLoaded && !mint) {
+    return {
+      amount: accountInfo?.lamports ?? 0,
+      decimals: 9,
+      mint: null,
+      owner: publicKey,
+      tokenName: 'Solana',
+      tokenSymbol: 'SOL',
+      initialized: false,
+    };
+  } else {
+    return null;
   }
-  return info;
 }
 
 export function useWalletSelector() {

yarn.lock

@@ -5698,6 +5698,11 @@ immer@1.10.0:
   resolved "https://registry.yarnpkg.com/immer/-/immer-1.10.0.tgz#bad67605ba9c810275d91e1c2a47d4582e98286d"
   integrity sha512-O3sR1/opvCDGLEVcvrGTMtLac8GJ5IwZC4puPrLuRj3l7ICKvkmA0vGuU9OW8mV9WIBRnaxp5GJh9IEAaNOoYg==
 
+immutable-tuple@^0.4.10:
+  version "0.4.10"
+  resolved "https://registry.yarnpkg.com/immutable-tuple/-/immutable-tuple-0.4.10.tgz#e0b1625384f514084a7a84b749a3bb26e9179929"
+  integrity sha512-45jheDbc3Kr5Cw8EtDD+4woGRUV0utIrJBZT8XH0TPZRfm8tzT0/sLGGzyyCCFqFMG5Pv5Igf3WY/arn6+8V9Q==
+
 import-cwd@^2.0.0:
   version "2.1.0"
   resolved "https://registry.yarnpkg.com/import-cwd/-/import-cwd-2.1.0.tgz#aa6cf36e722761285cb371ec6519f53e2435b0a9"