feat: uniswap-v3: track pool balance of tokens

Resolves BACK-786.

jest.config.js

@@ -3,7 +3,7 @@ module.exports = {
   testEnvironment: 'node',
   testRegex: [
     '/tests/.*\\.(test|spec)\\.(ts)$',
-    '/src/dex/.*\\.(test|spec)\\.(ts)$',
+    '/src/(dex|lib)/.*\\.(test|spec)\\.(ts)$',
   ],
   moduleFileExtensions: ['ts', 'js', 'json', 'node'],
   testTimeout: 30 * 1000,

src/dex/uniswap-v3/uniswap-v3-pool.ts

@@ -4,7 +4,10 @@ import { AbiItem } from 'web3-utils';
 import { Interface } from '@ethersproject/abi';
 import { DeepReadonly } from 'ts-essentials';
 import { Log, Logger, BlockHeader, Address } from '../../types';
-import { StatefulEventSubscriber } from '../../stateful-event-subscriber';
+import {
+  InitializeStateOptions,
+  StatefulEventSubscriber,
+} from '../../stateful-event-subscriber';
 import { IDexHelper } from '../../dex-helper/idex-helper';
 import {
   PoolState,
@@ -23,6 +26,8 @@ import {
   TICK_BITMAP_TO_USE,
 } from './constants';
 import { TickBitMap } from './contract-math/TickBitMap';
+import { ERC20EventSubscriber } from '../../lib/generics-events-subscribers/erc20-event-subscriber';
+import { getERC20Subscriber } from '../../lib/generics-events-subscribers/erc20-event-subscriber-factory';
 
 export class UniswapV3EventPool extends StatefulEventSubscriber<PoolState> {
   handlers: {
@@ -53,6 +58,9 @@ export class UniswapV3EventPool extends StatefulEventSubscriber<PoolState> {
 
   public readonly feeCodeAsString;
 
+  public token0sub: ERC20EventSubscriber;
+  public token1sub: ERC20EventSubscriber;
+
   constructor(
     readonly dexHelper: IDexHelper,
     parentName: string,
@@ -83,6 +91,9 @@ export class UniswapV3EventPool extends StatefulEventSubscriber<PoolState> {
       stateMultiAddress,
     );
 
+    this.token0sub = getERC20Subscriber(this.dexHelper, this.token0);
+    this.token1sub = getERC20Subscriber(this.dexHelper, this.token1);
+
     // Add handlers
     this.handlers['Swap'] = this.handleSwapEvent.bind(this);
     this.handlers['Burn'] = this.handleBurnEvent.bind(this);
@@ -102,7 +113,45 @@ export class UniswapV3EventPool extends StatefulEventSubscriber<PoolState> {
   }
 
   set poolAddress(address: Address) {
-    this._poolAddress = address;
+    this._poolAddress = address.toLowerCase();
+  }
+
+  async initialize(
+    blockNumber: number,
+    options?: InitializeStateOptions<PoolState>,
+  ) {
+    await super.initialize(blockNumber, options);
+    // only if the super call succeed
+
+    const initPromises = [];
+    if (!this.token0sub.isInitialized) {
+      initPromises.push(
+        this.token0sub.initialize(blockNumber, {
+          state: {},
+        }),
+      );
+    }
+
+    if (!this.token1sub.isInitialized) {
+      initPromises.push(
+        this.token1sub.initialize(blockNumber, {
+          state: {},
+        }),
+      );
+    }
+
+    await Promise.all(initPromises);
+
+    await Promise.all([
+      this.token0sub.subscribeToWalletBalanceChange(
+        this.poolAddress,
+        blockNumber,
+      ),
+      this.token1sub.subscribeToWalletBalanceChange(
+        this.poolAddress,
+        blockNumber,
+      ),
+    ]);
   }
 
   protected async processBlockLogs(
@@ -385,4 +434,12 @@ export class UniswapV3EventPool extends StatefulEventSubscriber<PoolState> {
       return acc;
     }, ticks);
   }
+
+  public getBalanceToken0(blockNumber: number) {
+    return this.token0sub.getBalance(this.poolAddress, blockNumber);
+  }
+
+  public getBalanceToken1(blockNumber: number) {
+    return this.token1sub.getBalance(this.poolAddress, blockNumber);
+  }
 }

src/dex/uniswap-v3/uniswap-v3.ts

@@ -41,6 +41,12 @@ import { DeepReadonly } from 'ts-essentials';
 import { uniswapV3Math } from './contract-math/uniswap-v3-math';
 import { Contract } from 'web3-eth-contract';
 import { AbiItem } from 'web3-utils';
+import { BalanceRequest, getBalances } from '../../lib/tokens/balancer-fetcher';
+import {
+  AssetType,
+  DEFAULT_ID_ERC20,
+  DEFAULT_ID_ERC20_AS_STRING,
+} from '../../lib/tokens/types';
 
 type PoolPairsInfo = {
   token0: Address;
@@ -252,6 +258,34 @@ export class UniswapV3
       return null;
     }
     this.logger.warn(`fallback to rpc for ${pools.length} pool(s)`);
+
+    const requests = pools.map<BalanceRequest>(
+      pool => ({
+        owner: pool.poolAddress,
+        asset: side == SwapSide.SELL ? from.address : to.address,
+        assetType: AssetType.ERC20,
+        ids: [
+          {
+            id: DEFAULT_ID_ERC20,
+            spenders: [],
+          },
+        ],
+      }),
+      [],
+    );
+
+    const balances = await getBalances(this.dexHelper.multiWrapper, requests);
+
+    pools = pools.filter(
+      (pool, index) =>
+        balances[index].amounts[DEFAULT_ID_ERC20_AS_STRING] >=
+        amounts[amounts.length - 1],
+    );
+
+    if (!pools.length) {
+      return null;
+    }
+
     pools.forEach(pool => {
       this.logger.warn(
         `[${this.network}][${pool.parentName}] fallback to rpc for ${pool.name}`,
@@ -461,6 +495,29 @@ export class UniswapV3
       const result = poolsToUse.poolWithState.map((pool, i) => {
         const state = states[i];
 
+        let balance = 0n;
+        if (_srcAddress === pool.token0) {
+          if (side === SwapSide.SELL) {
+            balance = pool.getBalanceToken0(blockNumber);
+          } else {
+            balance = pool.getBalanceToken1(blockNumber);
+          }
+        } else {
+          if (side === SwapSide.SELL) {
+            balance = pool.getBalanceToken1(blockNumber);
+          } else {
+            balance = pool.getBalanceToken0(blockNumber);
+          }
+        }
+
+        const requiredAmount = amounts[amounts.length - 1];
+        if (balance < requiredAmount) {
+          this.logger.debug(
+            `pool (${pool.poolAddress}) (srcToken: ${_srcAddress}, side: ${side}) have ${balance} but we need ${requiredAmount} to use it`,
+          );
+          return null;
+        }
+
         if (state.liquidity <= 0n) {
           return null;
         }

src/lib/generics-events-subscribers/erc20-event-subscriber-factory.ts

@@ -0,0 +1,18 @@
+import { IDexHelper } from '../../dex-helper';
+import { Address } from '../../types';
+import { ERC20EventSubscriber } from './erc20-event-subscriber';
+
+const subscriberMap: Record<Address, ERC20EventSubscriber> = {};
+
+export const getERC20Subscriber = (dexHelper: IDexHelper, token: string) => {
+  token = token.toLowerCase();
+  const identifier = `${dexHelper.config.data.network}-${token}`;
+  if (identifier in subscriberMap) {
+    return subscriberMap[identifier];
+  }
+
+  const sub = new ERC20EventSubscriber(dexHelper, token);
+  subscriberMap[identifier] = sub;
+
+  return sub;
+};

src/lib/generics-events-subscribers/erc20-event-subscriber.test.ts

@@ -0,0 +1,126 @@
+import dotenv from 'dotenv';
+dotenv.config();
+
+import { ERC20StateMap } from './types';
+import { ERC20EventSubscriber } from './erc20-event-subscriber';
+import { Network } from '../../constants';
+import { Address, Token } from '../../types';
+import { DummyDexHelper } from '../../dex-helper/index';
+import { testEventSubscriber } from '../../../tests/utils-events';
+import { getBalances } from '../tokens/balancer-fetcher';
+import {
+  AssetType,
+  DEFAULT_ID_ERC20,
+  DEFAULT_ID_ERC20_AS_STRING,
+} from '../tokens/types';
+import { MultiWrapper } from '../multi-wrapper';
+
+jest.setTimeout(50 * 1000);
+
+async function fetchBalance(
+  multiWrapper: MultiWrapper,
+  token: string,
+  wallet: string,
+  blockNumber: number,
+): Promise<ERC20StateMap> {
+  const balances = await getBalances(
+    multiWrapper,
+    [
+      {
+        owner: wallet,
+        asset: token,
+        assetType: AssetType.ERC20,
+        ids: [
+          {
+            id: DEFAULT_ID_ERC20,
+            spenders: [],
+          },
+        ],
+      },
+    ],
+    blockNumber,
+  );
+
+  const state = {} as ERC20StateMap;
+
+  state[wallet] = {
+    balance: balances[0].amounts[DEFAULT_ID_ERC20_AS_STRING],
+  };
+  return state;
+}
+
+// eventName -> blockNumbers
+type EventMappings = Record<string, number[]>;
+type WalletMapping = Record<string, EventMappings>;
+
+describe('ERC20 Subscriber Mainnet', function () {
+  const network = Network.MAINNET;
+  const dexHelper = new DummyDexHelper(network);
+  const token: Token = {
+    address: '0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2',
+    decimals: 18,
+  };
+
+  let erc20sub: ERC20EventSubscriber;
+
+  // tokenAddress -> EventMappings
+  const eventsToTest: Record<Address, WalletMapping> = {
+    '0x23fcf8d02b1b515ca40ec908463626c1759c2756': {
+      '0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2': {
+        Withdrawal: [16074564],
+      },
+    },
+    '0x39074b2b4434bf3115890094e1360e36d42ecbbd': {
+      '0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2': {
+        Transfer: [16074810],
+      },
+    },
+    '0x402df14df2080c5d946a8e2fc1b4bf78cbb1e73a': {
+      '0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2': {
+        Transfer: [16074809],
+      },
+    },
+  };
+
+  beforeAll(() => {
+    erc20sub = new ERC20EventSubscriber(dexHelper, token.address);
+  });
+
+  Object.keys(eventsToTest).forEach(async walletAddress => {
+    const eventsWithTokens = eventsToTest[walletAddress];
+    Object.entries(eventsWithTokens).forEach(
+      ([tokenAddress, events]: [string, EventMappings]) => {
+        describe(`Events for ${tokenAddress}`, () => {
+          Object.entries(events).forEach(
+            ([eventName, blockNumbers]: [string, number[]]) => {
+              describe(`${eventName}`, () => {
+                blockNumbers.forEach((blockNumber: number) => {
+                  it(`State after ${blockNumber}`, async function () {
+                    await erc20sub.subscribeToWalletBalanceChange(
+                      walletAddress,
+                      blockNumber - 1,
+                    );
+                    await testEventSubscriber(
+                      erc20sub,
+                      [token.address],
+                      (_blockNumber: number) =>
+                        fetchBalance(
+                          dexHelper.multiWrapper,
+                          token.address,
+                          walletAddress,
+                          _blockNumber,
+                        ),
+                      blockNumber,
+                      `${token.address}-${walletAddress}-${blockNumber}`,
+                      dexHelper.provider,
+                    );
+                  });
+                });
+              });
+            },
+          );
+        });
+      },
+    );
+  });
+});