AccountBalanceManager for Solana, can monitor the ATAs of various token balances in an account.

AccountBalanceManager.ts

import { Account, Mint, TOKEN_PROGRAM_ID, getAssociatedTokenAddress, unpackAccount, unpackMint } from "@solana/spl-token";
import { AccountChangeCallback, AccountInfo, Connection, Keypair, PublicKey } from "@solana/web3.js";
import { WellKnownTokenMint } from "../../const";
import { ToDecimal } from "../../utils";
import { BigNumber } from "bignumber.js";
import { TradeQueueItem, TradeExecutor } from "./trade";
import BalanceQueue from "../queues/balance";

export type BalanceChangeCallback = (queue: BalanceQueue<TradeQueueItem, TradeExecutor>) => Promise<TradeExecutor>;
export type BalanceChangeParams = [callback: BalanceChangeCallback, args: BalanceQueue<TradeQueueItem, TradeExecutor>]

type AtaAccountChangeCallback = (ata: PublicKey, ...args: Parameters<AccountChangeCallback>) => Promise<void>;

export default class AccountBalanceManager {
    connection: Connection;
    ataSubscriptionIds: Map<string, number>;
    public tokenAccounts: Map<string, Account | null>;
    public tokenMints: Map<string, Mint | null>;
    lamportsSubscriptionId: number;
    public lamportsAccount: AccountInfo<Buffer> | undefined;
    public readonly lamportsDecimals: number = 9;
    keyPair: Keypair;

    private tokenBalanceChangeCallbacks: Map<string, BalanceChangeParams> = new Map<string, BalanceChangeParams>();

    constructor(connection: Connection, keyPair: Keypair) {
        this.keyPair = keyPair;
        this.connection = connection;
        this.ataSubscriptionIds = new Map<string, number>();
        this.tokenAccounts = new Map<string, Account | null>();
        this.tokenMints = new Map<string, Mint | null>();
    }

    public addCallbackForTokenMintBalanceChange = async (
        queue: BalanceQueue<TradeQueueItem, TradeExecutor>,
        callback: BalanceChangeCallback
    ) => {
        const peeked = queue.peek();
        const mint = peeked.quote.inputMint;
        this.tokenBalanceChangeCallbacks.set(mint, [callback, queue]);
        return await this.addTokenProgramAccountSubscription(new PublicKey(mint));
    }

    public sufficientFundsForOrder = async (
        inputMint: string,
        inAmount: string,
    ): Promise<boolean> => {
        if (this.tokenBalanceChangeCallbacks.has(inputMint)) {
            return false;
        }
        const { amount: balance } = this.tokenAccounts.get(inputMint);
        if (balance === undefined) {
            console.warn(`No balance found for ${inputMint}`)
            return false;
        }
        const balanceBN = new BigNumber(`${balance}`)
        const inAmountBN = new BigNumber(inAmount);
        return balanceBN.minus(inAmountBN).isGreaterThanOrEqualTo(inAmountBN.multipliedBy(2));
    };

    public addMultipleTokenProgramAccountSubscriptions = async (tokenMints: PublicKey[]) => await Promise.all(tokenMints.map(async mint => this.addTokenProgramAccountSubscription(mint)));

    /**
     * Add a subscription for the account associated with our keypair and the token mint
     * @param tokenMint
     * @returns
     */
    public addTokenProgramAccountSubscription = async (tokenMint: PublicKey) => {
        if (this.ataSubscriptionIds.has(tokenMint.toString())) {
            return Promise.resolve();
        }
        this.tokenAccounts.set(tokenMint.toString(), null);
        this.tokenMints.set(tokenMint.toString(), null);
        const ata = await getAssociatedTokenAddress(tokenMint, this.keyPair.publicKey, false);
        const subscriptionId = this.connection.onAccountChange(ata,
            async (accountInfo, context) => await this.onTokenBalanceChange(ata, accountInfo, context)
        );
        this.ataSubscriptionIds.set(tokenMint.toString(), subscriptionId);
        console.log(`Subscribed to token mint: ${tokenMint.toString()}`)
        return Promise.resolve();
    }

    /** Add a subscription for the account associated with our keypair and the lamports account */
    public addLamportsAccountSubscription = async () => {
        if (this.lamportsSubscriptionId) {
            return;
        }
        this.lamportsSubscriptionId = this.connection.onAccountChange(this.keyPair.publicKey, this.onLamportsChange);
    }

    /**
     * Listener for lamports account balance change
     * @param accountInfo account info from the websocket
     * @param context context info from the websocket
     * @returns 
     */
    private onLamportsChange: AccountChangeCallback = async (accountInfo, context) => this.lamportsAccount = accountInfo;
    /**
     * Listener for token account balance change
     * @param ata Associated token address passed in from async local method that subscribes, since the returned data doesn't provide it
     * @param accountInfo account info from the websocket
     */
    private onTokenBalanceChange: AtaAccountChangeCallback = async (ata: PublicKey, accountInfo: AccountInfo<Buffer>) =>
        await new Promise(async (resolve, reject) => {
            const account = unpackAccount(ata, accountInfo);
            this.tokenAccounts.set(account.mint.toString(), account);
            const cb = this.tokenBalanceChangeCallbacks.get(account.mint.toString());
            if (cb) {
                const [callback, args] = cb;
                this.tokenBalanceChangeCallbacks.delete(account.mint.toString());
                await callback(args)
            }
            return resolve();
        });

    /** Return a formatted message of all balances */
    public getBalancesMsg = async (): Promise<string> => ([
        await Promise.all(
            Array.from(this.tokenAccounts.entries()).map(
                async ([mint, account]) => {
                    const token = this.tokenMints.get(mint);
                    if (token.decimals === undefined) {
                        return;
                    }
                    const amount = ToDecimal(account.amount, token.decimals)
                    if (mint == WellKnownTokenMint.Solana) {
                        return `\n${mint} (wrapped): ${amount}`;
                    }
                    return `\n${mint}: ${amount}`;
                }
            )
        ),
        this.lamportsAccount
            ? `\n${WellKnownTokenMint.Solana}: ${ToDecimal(this.lamportsAccount.lamports, this.lamportsDecimals)}`
            : ""
    ]).join("");

    /**
     * Unpack account info for a list of addresses
     * @param addresses a list of addresses to get account info for and unpack
     * @param unpacker is one of the unpack methods from web3.js
     * @returns Promise of the unpacked account info
     */
    private async getTokenProgramAccountInfo<T extends (address: PublicKey, info: AccountInfo<Buffer>, programId: PublicKey) => ReturnType<T>>(
        addresses: PublicKey[],
        unpacker: T
    ): Promise<ReturnType<typeof unpacker>[]> {
        const infos = await this.connection.getMultipleAccountsInfo(addresses);
        return await Promise.all(addresses.map(async (address, i) => unpacker(address, infos[i], TOKEN_PROGRAM_ID)));
    }

    /**
     * Run the account balance manager. Only needs to be run once to prepopulate the account balances.
     */
    public run = async () => {
        const mints = await Promise.all(Array.from(this.tokenAccounts.keys()).map(async mint => new PublicKey(mint)))
        const atas = await Promise.all(mints.map(async mint => getAssociatedTokenAddress(mint, this.keyPair.publicKey, false)));
        const [ataInfos, mintInfos] = await Promise.all([
            this.getTokenProgramAccountInfo(atas, unpackAccount),
            this.getTokenProgramAccountInfo(mints, unpackMint)
        ]);
        await Promise.all([
            ...mintInfos.map(async mint => this.tokenMints.set(mint.address.toString(), mint)),
            ...ataInfos.map(async account => this.tokenAccounts.set(account.mint.toString(), account)),
            this.lamportsAccount = await this.connection.getAccountInfo(this.keyPair.publicKey)
        ]);
    }
}