autodetect balances for vaults which are created before the backend started

src/dexorder/event_handler.py

@@ -1,5 +1,4 @@
 import asyncio
-import functools
 import logging
 from uuid import UUID
 
@@ -14,7 +13,7 @@ from dexorder.transaction import submit_transaction_request
 from dexorder.pools import new_pool_prices, pool_prices, get_uniswap_data
 from dexorder.contract.dexorder import vault_address, VaultContract
 from dexorder.contract import ERC20
-from dexorder.vault_blockdata import vault_owners, vault_balances
+from dexorder.vault_blockdata import vault_owners, vault_balances, adjust_balance
 from dexorder.database.model.block import current_block
 from dexorder.database.model.transaction import TransactionJob
 from dexorder.base.orderlib import SwapOrderState
@@ -118,22 +117,6 @@ async def handle_order_cancel_all(event: EventData):
                 log.warning(f'DexorderCanceled IGNORED due to missing order {vault} {orderIndex}')
 
 
-def balance_adjuster(vault, token_address, amount):
-    async def _adjust(vaddr, taddr, amt, old_balances):
-        result = dict(old_balances)  # copy
-        try:
-            old = old_balances[vaddr]
-            new_amt = old + amt
-            if new_amt < 0:
-                log.warning(f'NEGATIVE BALANCE for vault {current_chain.get()} {vault} token {taddr}  {old} {amt:+} = {new_amt}')
-                new_amt = 0
-        except KeyError:
-            new_amt = await ERC20(taddr).balanceOf(vaddr)
-        result[taddr] = new_amt
-        return result
-    return functools.partial(_adjust, vault, token_address, amount)
-
-
 async def handle_transfer(transfer: EventData):
     # todo handle native transfers incl gas for token transfers
     log.debug(f'Transfer {transfer}')
@@ -144,12 +127,12 @@ async def handle_transfer(transfer: EventData):
         log.debug(f'deposit {to_address} {amount}')
         vault = to_address
         token_address = transfer['address']
-        await vault_balances.async_modify(vault, balance_adjuster(vault, token_address, +amount), default={})
+        await adjust_balance(vault, token_address, amount)
     if from_address in vault_owners and to_address != from_address:
         log.debug(f'withdraw {to_address} {amount}')
         vault = from_address
         token_address = transfer['address']
-        await vault_balances.async_modify(vault, balance_adjuster(vault, token_address, amount), default={})
+        await adjust_balance(vault, token_address, amount)
     # if to_address not in vault_owners and from_address not in vault_owners:
     #     vaults = vault_owners.keys()
     #     log.debug(f'vaults: {list(vaults)}')
@@ -221,23 +204,29 @@ async def activate_price_triggers():
         await maywait(t(None))
 
 
-def process_active_tranches():
+async def process_active_tranches():
     for tk, proof in active_tranches.items():
         old_req = execution_requests.get(tk)
         height = current_block.get().height
         if old_req is None or old_req.height <= height: # '<=' is used so proof is updated with more recent values
-            if has_funds(tk):
+            if await has_funds(tk):
                 log.info(f'execution request for {tk}')
                 execution_requests[tk] = ExecutionRequest(height, proof)
             else: log.debug(f'underfunded tranche {tk}')
 
 
-def has_funds(tk: TrancheKey):
+async def has_funds(tk: TrancheKey):
     # log.debug(f'has funds? {tk.vault}')
     order = Order.of(tk)
     minimum = order.status.order.minFillAmount if order.amount_is_input else 0
     balances = vault_balances.get(tk.vault, {})
-    token_balance = balances.get(order.status.order.tokenIn, 0)
+    token_addr = order.status.order.tokenIn
+    token_balance = balances.get(token_addr)
+    if token_balance is None:
+        # unknown balance
+        token_balance = balances[token_addr] = await ERC20(token_addr).balanceOf(tk.vault)
+        log.debug(f'queried token balance {token_addr}.balanceOf({tk.vault}) = {token_balance}')
+        await adjust_balance(tk.vault, token_addr, token_balance)
     # log.debug(f'minimum {minimum} balances {balances}')
     return token_balance > minimum
 

src/dexorder/vault_blockdata.py

@@ -1,7 +1,9 @@
+import functools
 import logging
 
 from dexorder.base.chain import current_chain
 from dexorder.blockstate import BlockDict
+from dexorder.contract import ERC20
 from dexorder.util import json
 
 log = logging.getLogger(__name__)
@@ -27,3 +29,22 @@ vault_balances: BlockDict[str, dict[str, int]] = BlockDict(
     pub=pub_vault_balances
 )
 
+
+async def adjust_balance(vault, token_address, amount):
+    await vault_balances.async_modify(vault, balance_adjuster(vault, token_address, amount), default={})
+
+
+def balance_adjuster(vault, token_address, amount):
+    async def _adjust(vaddr, taddr, amt, old_balances):
+        result = dict(old_balances)  # copy
+        try:
+            old = old_balances[vaddr]
+            new_amt = old + amt
+            if new_amt < 0:
+                log.warning(f'NEGATIVE BALANCE for vault {current_chain.get()} {vault} token {taddr}  {old} {amt:+} = {new_amt}')
+                new_amt = 0
+        except KeyError:
+            new_amt = await ERC20(taddr).balanceOf(vaddr)
+        result[taddr] = new_amt
+        return result
+    return functools.partial(_adjust, vault, token_address, amount)