Fix tracker query balance performance

CATProtocol · Nov 8, 2024 · 1aa0706 · 1aa0706
1 parent cffb1b0
commit 1aa0706
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 58 deletions.
diff --git a/packages/tracker/src/routes/address/address.service.ts b/packages/tracker/src/routes/address/address.service.ts
@@ -1,6 +1,5 @@
 import { Injectable } from '@nestjs/common';
 import { TokenService } from '../token/token.service';
-import { xOnlyPubKeyToAddress } from '../../common/utils';
 import { CommonService } from '../../services/common/common.service';
 import { TokenTypeScope } from '../../common/types';
 
@@ -19,29 +18,17 @@ export class AddressService {
     return this.getBalances(ownerAddrOrPkh, TokenTypeScope.NonFungible);
   }
 
-  private async getBalances(ownerAddrOrPkh: string, scope: TokenTypeScope) {
+  private async getBalances(
+    ownerAddrOrPkh: string,
+    scope: TokenTypeScope.Fungible | TokenTypeScope.NonFungible,
+  ) {
     const lastProcessedHeight =
       await this.commonService.getLastProcessedBlockHeight();
-    const utxos = await this.tokenService.queryTokenUtxosByOwnerAddress(
+    const balances = await this.tokenService.queryTokenBalancesByOwnerAddress(
       lastProcessedHeight,
       ownerAddrOrPkh,
+      scope,
     );
-    const tokenBalances = await this.tokenService.groupTokenBalances(utxos);
-    const balances = [];
-    for (const tokenPubKey in tokenBalances) {
-      const tokenAddr = xOnlyPubKeyToAddress(tokenPubKey);
-      const tokenInfo =
-        await this.tokenService.getTokenInfoByTokenIdOrTokenAddress(
-          tokenAddr,
-          scope,
-        );
-      if (tokenInfo) {
-        balances.push({
-          tokenId: tokenInfo.tokenId,
-          confirmed: tokenBalances[tokenPubKey].toString(),
-        });
-      }
-    }
     return { balances, trackerBlockHeight: lastProcessedHeight };
   }
 }
diff --git a/packages/tracker/src/routes/collection/collection.controller.ts b/packages/tracker/src/routes/collection/collection.controller.ts
@@ -265,11 +265,13 @@ export class CollectionController {
         TokenTypeScope.NonFungible,
         ownerAddrOrPkh,
       );
-      return okResponse({
-        collectionId: balance.tokenId,
-        confirmed: balance.confirmed,
-        trackerBlockHeight: balance.trackerBlockHeight,
-      });
+      return okResponse(
+        balance && {
+          collectionId: balance.tokenId,
+          confirmed: balance.confirmed,
+          trackerBlockHeight: balance.trackerBlockHeight,
+        },
+      );
     } catch (e) {
       return errorResponse(e);
     }

diff --git a/packages/tracker/src/routes/token/token.service.ts b/packages/tracker/src/routes/token/token.service.ts
@@ -159,7 +159,7 @@ export class TokenService {
 
   async getTokenBalanceByOwnerAddress(
     tokenIdOrTokenAddr: string,
-    scope: TokenTypeScope,
+    scope: TokenTypeScope.Fungible | TokenTypeScope.NonFungible,
     ownerAddrOrPkh: string,
   ) {
     const lastProcessedHeight =
@@ -168,22 +168,18 @@ export class TokenService {
       tokenIdOrTokenAddr,
       scope,
     );
-    let utxos = [];
-    if (tokenInfo) {
-      utxos = await this.queryTokenUtxosByOwnerAddress(
-        lastProcessedHeight,
-        ownerAddrOrPkh,
-        tokenInfo,
-      );
-    }
-    let confirmed = '0';
-    if (tokenInfo?.tokenPubKey) {
-      const tokenBalances = await this.groupTokenBalances(utxos);
-      confirmed = tokenBalances[tokenInfo.tokenPubKey]?.toString() || '0';
+    if (!tokenInfo) {
+      return null;
     }
+    const balances = await this.queryTokenBalancesByOwnerAddress(
+      lastProcessedHeight,
+      ownerAddrOrPkh,
+      scope,
+      tokenInfo,
+    );
     return {
-      tokenId: tokenInfo?.tokenId || null,
-      confirmed,
+      tokenId: tokenInfo.tokenId,
+      confirmed: balances.length === 1 ? balances[0].confirmed.toString() : '0',
       trackerBlockHeight: lastProcessedHeight,
     };
   }
@@ -222,6 +218,49 @@ export class TokenService {
     });
   }
 
+  async queryTokenBalancesByOwnerAddress(
+    lastProcessedHeight: number,
+    ownerAddrOrPkh: string,
+    scope: TokenTypeScope.Fungible | TokenTypeScope.NonFungible,
+    tokenInfo: TokenInfoEntity | null = null,
+  ) {
+    const ownerPubKeyHash =
+      ownerAddrOrPkh.length === Constants.PUBKEY_HASH_BYTES * 2
+        ? ownerAddrOrPkh
+        : ownerAddressToPubKeyHash(ownerAddrOrPkh);
+    if (
+      lastProcessedHeight === null ||
+      (tokenInfo && !tokenInfo.tokenPubKey) ||
+      !ownerPubKeyHash
+    ) {
+      return [];
+    }
+    const query = this.txOutRepository
+      .createQueryBuilder('t1')
+      .select('t2.token_id', 'tokenId')
+      .innerJoin(TokenInfoEntity, 't2', 't1.xonly_pubkey = t2.token_pubkey')
+      .where('t1.spend_txid IS NULL')
+      .andWhere('t1.owner_pkh = :ownerPkh', { ownerPkh: ownerPubKeyHash })
+      .groupBy('t2.token_id');
+    if (scope === TokenTypeScope.Fungible) {
+      query
+        .addSelect('SUM(t1.token_amount)', 'confirmed')
+        .andWhere('t2.decimals >= 0');
+    } else {
+      query.addSelect('COUNT(1)', 'confirmed').andWhere('t2.decimals < 0');
+    }
+    if (tokenInfo) {
+      query.andWhere('t1.xonly_pubkey = :tokenPubKey', {
+        tokenPubKey: tokenInfo.tokenPubKey,
+      });
+    }
+    const results = await query.getRawMany();
+    return results.map((r) => ({
+      tokenId: r.tokenId,
+      confirmed: r.confirmed,
+    }));
+  }
+
   async queryStateHashes(txid: string) {
     let cached = TokenService.stateHashesCache.get(txid);
     if (!cached) {
@@ -279,25 +318,6 @@ export class TokenService {
     return renderedUtxos;
   }
 
-  /**
-   * @param utxos utxos with the same owner address
-   * @returns token balances grouped by xOnlyPubKey
-   */
-  async groupTokenBalances(utxos: TxOutEntity[]) {
-    const balances = {};
-    for (const utxo of utxos) {
-      const tokenInfo = await this.getTokenInfoByTokenPubKey(
-        utxo.xOnlyPubKey,
-        TokenTypeScope.All,
-      );
-      if (tokenInfo) {
-        const acc = tokenInfo.decimals >= 0 ? BigInt(utxo.tokenAmount) : 1n;
-        balances[utxo.xOnlyPubKey] = (balances[utxo.xOnlyPubKey] || 0n) + acc;
-      }
-    }
-    return balances;
-  }
-
   async getTokenMintCount(
     tokenIdOrTokenAddr: string,
     scope: TokenTypeScope.Fungible | TokenTypeScope.NonFungible,