@exodus/solana-api 3.25.4 → 3.26.0

package/src/tx-log/ws-monitor.js

@@ -0,0 +1,390 @@
+import delay from 'delay'
+import assert from 'minimalistic-assert'
+
+import { DEFAULT_POOL_ADDRESS } from '../account-state.js'
+import { SolanaClarityMonitor } from './clarity-monitor.js'
+
+const DEFAULT_REMOTE_CONFIG = {
+  clarityUrl: [],
+  rpcUrl: [],
+  ws: [],
+  staking: { enabled: true, pool: DEFAULT_POOL_ADDRESS },
+}
+
+export class SolanaWebsocketMonitor extends SolanaClarityMonitor {
+  constructor({ wsApi, ...args }) {
+    assert(wsApi, 'wsApi is required')
+    super(args)
+    this.wsApi = wsApi
+    this.tokenAccountsByOwner = Object.create(null)
+    this.batch = Object.create(null)
+
+    this.addHook('before-start', (...args) => this.beforeStart(...args))
+    this.addHook('before-stop', (...args) => this.beforeStop(...args))
+  }
+
+  async beforeStart() {
+    this.assets = await this.aci.getAssetsForNetwork({ baseAssetName: this.asset.name })
+    this.api.setTokens(this.assets)
+
+    const walletAccounts = await this.aci.getWalletAccounts({ assetName: this.asset.name })
+    await Promise.all(
+      walletAccounts.map((walletAccount) => this.#subscribeWalletAddresses({ walletAccount }))
+    )
+  }
+
+  async beforeStop() {
+    const walletAccounts = await this.aci.getWalletAccounts({ assetName: this.asset.name })
+    await Promise.all(
+      walletAccounts.map((walletAccount) => this.#unsubscribeWalletAddresses({ walletAccount }))
+    )
+    const isBatchActive = walletAccounts.some((walletAccount) => !!this.batch[walletAccount])
+    if (isBatchActive) await delay(2000) // wait for any batch opened to close
+  }
+
+  async #subscribeWalletAddresses({ walletAccount }) {
+    const address = await this.aci.getReceiveAddress({
+      assetName: this.asset.name,
+      walletAccount,
+      useCache: true,
+    })
+
+    const { accounts: tokenAccountsByOwner } = await this.api.getTokensBalancesAndAccounts({
+      address,
+    })
+    this.tokenAccountsByOwner[walletAccount] = tokenAccountsByOwner
+    const tokensAddresses = tokenAccountsByOwner.map((acc) => acc.tokenAccountAddress)
+
+    return this.wsApi.watchAddress({
+      address,
+      tokensAddresses,
+      onMessage: async (json) => {
+        // parse event, then update state or tx-log
+        await this.#handleMessage({ address, walletAccount, data: json })
+      },
+    })
+  }
+
+  async #unsubscribeWalletAddresses({ walletAccount }) {
+    const address = await this.aci.getReceiveAddress({
+      assetName: this.asset.name,
+      walletAccount,
+      useCache: true,
+    })
+    return this.wsApi.unwatchAddress({ address })
+  }
+
+  setServer(config = {}) {
+    const { ws } = { ...DEFAULT_REMOTE_CONFIG, ...config }
+    this.wsApi.setWsEndpoint(ws[0])
+    super.setServer(config)
+  }
+
+  async tick({ walletAccount, refresh }) {
+    // we tick using Clarity only on startup, then we rely only on WS events and we periodically check for new tokens only.
+    if (refresh || this.tickCount[walletAccount] === 0) {
+      return super.tick({ walletAccount, refresh }) // Clarity refresh or first tick
+    }
+
+    const assetName = this.asset.name
+    this.assets = await this.aci.getAssetsForNetwork({ baseAssetName: assetName })
+    this.api.setTokens(this.assets)
+    const address = await this.aci.getReceiveAddress({
+      assetName,
+      walletAccount,
+      useCache: true,
+    })
+
+    // we only update balance of the tokens we just subcribed to:
+    const accountState = await this.aci.getAccountState({
+      assetName: this.asset.name,
+      walletAccount,
+    })
+
+    // we call this periodically to detect new token accounts created (there's not a WS event for this in Helius yet, we need programSubscribe)
+    const { accounts: tokenAccounts, balances: splBalances } =
+      await this.api.getTokensBalancesAndAccounts({
+        address,
+      })
+    this.tokenAccountsByOwner[walletAccount] = tokenAccounts
+
+    const unknownTokensList = await this.emitUnknownTokensEvent({
+      tokenAccounts,
+    })
+
+    // subscribe to new tokenAccounts
+    for (const mintAddress of unknownTokensList) {
+      await this.wsApi.accountSubscribe({ owner: address, account: mintAddress })
+      const tokenName = this.api.tokens.get(mintAddress)?.name
+      if (!tokenName) {
+        console.log(`Unknown token mint address: ${mintAddress}`)
+        continue
+      }
+
+      // update only token balances for known tokens
+      const amount = splBalances[mintAddress]
+      const newData = {
+        tokenBalances: {
+          ...accountState.tokenBalances,
+          [tokenName]: this.assets[tokenName].currency.baseUnit(amount),
+        },
+      }
+      await this.#updateStateBatch({ newData, walletAccount })
+    }
+
+    // await this.updateState({ account, walletAccount, staking }) // we could update tokenBalances but we gotta test for race-conditions
+  }
+
+  async #handleMessage({ address, walletAccount, data }) {
+    const accountState = await this.aci.getAccountState({
+      assetName: this.asset.name,
+      walletAccount,
+    })
+    const tokenAccountsByOwner = this.tokenAccountsByOwner[walletAccount]
+
+    /*
+      1.	A new event arrives.
+      2.	Open a 2-second batch window.
+      3.	Add balance updates and transactions to the batch.
+      4.	After 2 seconds, close the window and execute the batch.
+    */
+
+    if (['accountNotification', 'transactionNotification'].includes(data?.method)) {
+      this.#ensureBatch(walletAccount)
+    }
+
+    switch (data?.method) {
+      case 'accountNotification':
+        // balance changed events for known tokens or SOL address
+
+        const { amount, tokenMintAddress } = this.wsApi.parseAccountNotification({
+          address,
+          walletAccount,
+          tokenAccountsByOwner,
+          result: data.params.result,
+        })
+
+        // update account state balance for SOL or Token
+        if (tokenMintAddress) {
+          // token balance changed
+          const tokenName = this.api.tokens.get(tokenMintAddress)?.name
+          if (!tokenName) {
+            console.log(`Unknown token mint address: ${tokenMintAddress}`)
+            return
+          }
+
+          const newData = {
+            // ...accountState, // we don't wanna update SOL "balance"
+            tokenBalances: {
+              ...accountState.tokenBalances,
+              [tokenName]: this.assets[tokenName].currency.baseUnit(amount),
+            },
+          }
+          await this.#updateStateBatch({ newData, walletAccount })
+        } else {
+          // SOL balance changed
+
+          // Wait up to 1.5s looking for an unconfirmed staking tx in the txLog.
+          // If we find an unconfirmed staking tx, we don't update balance.
+          // we can update staking balances in this handler because we cannot distinguish staking balance update from other updates
+          // moreover to compute a right "balance" we need an updated stakingInfo. That we don't have here.
+          const watchDelay = 150 // ms
+          const watchTimeout = 1500 // ms (1.5s)
+          const start = Date.now()
+
+          while (Date.now() - start < watchTimeout) {
+            const baseAssetTxLog = await this.aci.getTxLog({
+              assetName: this.asset.name,
+              walletAccount,
+            })
+            for (const tx of baseAssetTxLog) {
+              if (tx.pending && tx.data?.staking) {
+                console.log('found unconfirmed SOL staking tx. Skipping balance update.')
+                return // no balance update, the transactionNotification flow handle staking updates
+              }
+            }
+
+            await new Promise((resolve) => setTimeout(resolve, watchDelay))
+          }
+
+          // otherwise regular SOL update:
+          const stakingInfo = accountState.stakingInfo // NB. we cannot call this.getStakingInfo(...) since it's not in sync with the ws event! we must wait a lot of seconds.
+          const balance = this.#computeTotalBalance({ amount, address, stakingInfo, walletAccount })
+
+          const newData = {
+            // ...accountState, // we don't wanna update "tokenBalances" with old values (after sending an SPL token we have 2 events one updating SOL balance and one updating tokenBalances)
+            balance,
+          }
+
+          await this.#updateStateBatch({ newData, walletAccount })
+        }
+
+        return
+      case 'transactionNotification':
+        // update tx-log with new txs and state with cursor
+
+        // NB. if we receive a tx with a new "token", never received before,
+        // this.api.tokens will be populated with new tokens only after we ran a monitor tick
+        // since the unknown-tokens event is captured and processed in other places.
+        // hence what happens is we skip the txLog update for that new token.
+
+        const { logItemsByAsset, cursorState = {} } = this.wsApi.parseTransactionNotification({
+          address,
+          walletAccount,
+          baseAsset: this.asset,
+          assets: this.assets,
+          tokens: this.api.tokens,
+          tokenAccountsByOwner,
+          result: data.params.result, // raw tx
+        })
+
+        if (Object.keys(logItemsByAsset).length === 0) return // cannot parse tx
+
+        const clearedLogItems = await this.markStaleTransactions({ walletAccount, logItemsByAsset })
+
+        // if there are notifications about staking txs, we need to update balance and stakingInfo
+        // check if we're dealing with a staking tx
+        let stakingTx
+        clearedLogItems?.solana?.forEach((tx) => {
+          if (tx.data.staking) stakingTx = tx
+        })
+
+        if (stakingTx) stakingTx.confirmations = 0 // HACK: mark it as unconfirmed till we call and update the stakingInfo
+
+        // update txLog
+        await this.#updateHistoryBatch({
+          walletAccount,
+          logItemsByAsset: clearedLogItems,
+          refresh: false,
+        })
+
+        const newData = {
+          ...cursorState,
+        }
+
+        if (stakingTx) {
+          // for staking the balance is not updated by the balance handler
+          // staking operations won't spend or modify the "total" wallet balance.
+
+          // we update stakingInfo (before fetching the new one)
+          const temporaryStakingInfo = { ...accountState.stakingInfo }
+          // eslint-disable-next-line sonarjs/no-nested-switch
+          switch (stakingTx.data.staking?.method) {
+            case 'delegate':
+              const stakeAmount = this.asset.currency.baseUnit(stakingTx.data.staking?.stake || '0')
+              temporaryStakingInfo.activating = temporaryStakingInfo.activating.add(
+                stakeAmount.abs()
+              )
+              break
+            case 'withdraw':
+              temporaryStakingInfo.withdrawable = this.asset.currency.ZERO
+              break
+            case 'undelegate':
+              temporaryStakingInfo.pending = temporaryStakingInfo.pending
+                .add(temporaryStakingInfo.locked)
+                .add(temporaryStakingInfo.activating)
+              temporaryStakingInfo.locked = this.asset.currency.ZERO
+              temporaryStakingInfo.activating = this.asset.currency.ZERO
+              break
+          }
+
+          newData.stakingInfo = temporaryStakingInfo
+
+          // we update only stakingInfo, so that in this 12sec window, the spendable balance is valid. In the meanwhile we can fetch the complete one from RPC and update it in background (this one has the "staking accounts" references).
+          await this.#updateStateBatch({ newData, walletAccount })
+
+          await delay(12_000) // we introduce a delay to make sure the getStakingInfo RPC returns updated data
+          // because we NEED at a certain point to have an updated stakingInfo in the state to know what's the new solana stake account address.
+          const stakingInfo = await this.getStakingInfo({ address, accountState, walletAccount })
+          const balance = this.#computeTotalBalance({
+            amount: this.asset.currency.baseUnit(await this.api.getBalance(address)), // we needed the delay otherwise will never be updated right away.
+            address,
+            stakingInfo,
+            walletAccount,
+          })
+
+          stakingTx.confirmations = 1 // mark tx as confirmed
+          await this.#updateHistoryBatch({
+            walletAccount,
+            logItemsByAsset: clearedLogItems,
+            refresh: false,
+          })
+          await this.#updateStateBatch({ newData: { balance, stakingInfo }, walletAccount })
+
+          return
+        }
+
+        // NB. we don't update any balances here to avoid race conditions with the accountNotification event above
+        await this.#updateStateBatch({ newData, walletAccount })
+
+        return
+      default:
+        if (data?.result && typeof data.result === 'number') {
+          // subscription confirmation, skip
+          return
+        }
+
+        if (data?.result && data.result?.['solana-core']) {
+          // ping response, skip
+          return
+        }
+
+        console.log(`Unknown SOL WS method: ${data?.method}`)
+        break
+    }
+  }
+
+  #ensureBatch(walletAccount) {
+    // open a new operations batch window if it's not opened already
+    if (this.batch[walletAccount]) return this.batch[walletAccount]
+    this.batch[walletAccount] = this.aci.createOperationsBatch()
+    setTimeout(async () => {
+      const batch = this.batch[walletAccount]
+      this.batch[walletAccount] = null
+      await this.aci.executeOperationsBatch(batch)
+    }, 2000)
+    return this.batch[walletAccount]
+  }
+
+  async #updateStateBatch({ newData, walletAccount }) {
+    const assetName = this.asset.name
+    const batch = this.#ensureBatch(walletAccount)
+    this.updateAccountStateBatch({ assetName, walletAccount, newData, batch })
+  }
+
+  async #updateHistoryBatch({ walletAccount, logItemsByAsset, refresh }) {
+    const batch = this.#ensureBatch(walletAccount)
+    this.updateTxLogByAssetBatch({ logItemsByAsset, walletAccount, refresh, batch })
+  }
+
+  #computeTotalBalance({ amount, address, stakingInfo, walletAccount }) {
+    const solBalance = this.asset.currency.baseUnit(amount)
+
+    const { stakedBalance, activatingBalance, withdrawableBalance, pendingBalance } =
+      this.#getStakingBalances({
+        address,
+        stakingInfo,
+        walletAccount,
+      })
+
+    return solBalance
+      .add(stakedBalance)
+      .add(activatingBalance)
+      .add(withdrawableBalance)
+      .add(pendingBalance)
+  }
+
+  #getStakingBalances({ address, stakingInfo, walletAccount }) {
+    const stakedBalance = this.asset.currency.baseUnit(stakingInfo.locked)
+    const activatingBalance = this.asset.currency.baseUnit(stakingInfo.activating)
+    const withdrawableBalance = this.asset.currency.baseUnit(stakingInfo.withdrawable)
+    const pendingBalance = this.asset.currency.baseUnit(stakingInfo.pending)
+
+    return {
+      stakedBalance, // staked
+      activatingBalance, // staking
+      withdrawableBalance, // unstaked
+      pendingBalance, // unstaking
+    }
+  }
+}