fix(kernel): optimize AccountInfoUnit's speed (#881)

common/changes/@yuants/kernel/2024-10-31-21-10.json

@@ -0,0 +1,10 @@
+{
+  "changes": [
+    {
+      "packageName": "@yuants/kernel",
+      "comment": "optimize AccountInfoUnit's speed",
+      "type": "patch"
+    }
+  ],
+  "packageName": "@yuants/kernel"
+}

libraries/kernel/etc/kernel.api.md

@@ -666,6 +666,7 @@ export class ProductLoadingUnit extends BasicUnit {
 
 // @public
 export class QuoteDataUnit extends BasicUnit {
+    dirtyProductIds: Set<string>;
     // (undocumented)
     dump(): {
         mapProductIdToQuote: Record<string, Record<string, IQuote>>;

libraries/kernel/src/units/AccountInfoUnit.ts

@@ -1,4 +1,4 @@
-import { IAccountInfo, IAccountMoney, IPosition, createEmptyAccountInfo } from '@yuants/data-model';
+import { IAccountInfo, IPosition, createEmptyAccountInfo } from '@yuants/data-model';
 import { roundToStep } from '@yuants/utils';
 import { Subscription } from 'rxjs';
 import { Kernel } from '../kernel';
@@ -36,7 +36,11 @@ export class AccountInfoUnit extends BasicUnit {
 
   private mapAccountIdToBalance: Record<string, number> = {};
 
-  private mapAccountIdToPositionIdToPosition: Record<string, Record<string, IPosition>> = {};
+  private mapAccountIdToPositionIdToPosition: Record<string, Record<string, IPosition>> = {}; // 复用引用，增强性能
+  private mapAccountIdToPositions: Record<string, IPosition[]> = {}; // 复用引用，增强性能
+  private mapAccountIdToProductIdToPositions: Record<string, Record<string, IPosition[]>> = {}; // 复用引用，增强性能
+  private mapPositionIdToUsed: Record<string, number | undefined> = {};
+  private mapPositionIdToFloatingProfit: Record<string, number | undefined> = {};
 
   mapAccountIdToAccountInfo: Map<string, IAccountInfo> = new Map();
 
@@ -47,7 +51,10 @@ export class AccountInfoUnit extends BasicUnit {
     }
     const newAccountInfo = createEmptyAccountInfo(account_id, currency, leverage, initial_balance);
     this.mapAccountIdToBalance[account_id] = newAccountInfo.money.balance;
+    this.mapAccountIdToPositions[account_id] = newAccountInfo.positions;
     this.mapAccountIdToAccountInfo.set(account_id, newAccountInfo);
+    this.mapAccountIdToPositionIdToPosition[account_id] = {};
+    this.mapAccountIdToProductIdToPositions[account_id] = {};
     return newAccountInfo;
   }
 
@@ -57,19 +64,45 @@ export class AccountInfoUnit extends BasicUnit {
     product_id: string,
     direction: string,
   ): IPosition => {
-    return ((this.mapAccountIdToPositionIdToPosition[account_id] ??= {})[position_id] ??= {
-      position_id,
-      product_id,
-      direction,
-      position_price: NaN,
-      volume: 0,
-      closable_price: NaN,
-      floating_profit: 0,
-      free_volume: 0,
-      valuation: 0,
-    });
+    const mapPositionIdToPosition = this.mapAccountIdToPositionIdToPosition[account_id];
+    if (!mapPositionIdToPosition[position_id]) {
+      const initPosition = {
+        position_id,
+        product_id,
+        direction,
+        position_price: NaN,
+        volume: 0,
+        closable_price: NaN,
+        floating_profit: 0,
+        free_volume: 0,
+        valuation: 0,
+      };
+      mapPositionIdToPosition[position_id] = initPosition;
+      this.mapAccountIdToPositions[account_id].push(initPosition);
+      (this.mapAccountIdToProductIdToPositions[account_id][product_id] ??= []).push(initPosition);
+    }
+    return mapPositionIdToPosition[position_id];
   };
 
+  private removePosition(account_id: string, position_id: string) {
+    delete this.mapAccountIdToPositionIdToPosition[account_id][position_id];
+    const positions = this.mapAccountIdToPositions[account_id];
+    const idx = positions.findIndex((position) => position.position_id === position_id);
+    if (idx < 0) return;
+    positions.splice(idx, 1);
+    const product_id = positions[idx].product_id;
+    const productPositions = this.mapAccountIdToProductIdToPositions[account_id][product_id];
+    const idx2 = productPositions.findIndex((position) => position.position_id === position_id);
+    if (idx2 < 0) return;
+    productPositions.splice(idx2, 1);
+    const theAccountInfo = this.mapAccountIdToAccountInfo.get(account_id);
+    if (!theAccountInfo) return;
+    theAccountInfo.money.used -= this.mapPositionIdToUsed[position_id] || 0;
+    delete this.mapPositionIdToUsed[position_id];
+    theAccountInfo.money.profit -= this.mapPositionIdToFloatingProfit[position_id] || 0;
+    delete this.mapPositionIdToFloatingProfit[position_id];
+  }
+
   updateAccountInfo(accountId: string) {
     const theAccountInfo = this.mapAccountIdToAccountInfo.get(accountId);
 
@@ -78,13 +111,14 @@ export class AccountInfoUnit extends BasicUnit {
     }
 
     // 根据订单更新头寸
+    const dirtyProductIds = new Set<string>(this.quoteDataUnit.dirtyProductIds);
     // ISSUE: 假设订单一旦成交即全部成交
     for (let idx = this.orderIdx; idx < this.historyOrderUnit.historyOrders.length; idx++) {
       const order = this.historyOrderUnit.historyOrders[idx];
       if (order.account_id !== accountId) {
         continue;
       }
-
+      dirtyProductIds.add(order.product_id);
       if (order.profit_correction) {
         this.mapAccountIdToBalance[accountId] =
           (this.mapAccountIdToBalance[accountId] || 0) + order.profit_correction;
@@ -108,13 +142,9 @@ export class AccountInfoUnit extends BasicUnit {
               (thePosition.volume / thePosition.position_price + order.volume / order.traded_price!)
             : (thePosition.position_price * thePosition.volume + order.traded_price! * order.volume) /
               nextVolume;
-          const position: IPosition = {
-            ...thePosition,
-            volume: nextVolume,
-            free_volume: nextVolume,
-            position_price: nextPositionPrice,
-          };
-          this.mapAccountIdToPositionIdToPosition[accountId][order.position_id!] = position;
+          thePosition.volume = nextVolume;
+          thePosition.free_volume = nextVolume;
+          thePosition.position_price = nextPositionPrice;
         }
       } else {
         // 平仓
@@ -129,7 +159,7 @@ export class AccountInfoUnit extends BasicUnit {
           const nextVolume = roundToStep(thePosition.volume - tradedVolume, theProduct.volume_step ?? 1);
           // 如果头寸已经平仓完了，就删除头寸
           if (nextVolume === 0) {
-            delete this.mapAccountIdToPositionIdToPosition[accountId][order.position_id!];
+            this.removePosition(accountId, order.position_id!);
           }
           thePosition.volume = nextVolume;
           thePosition.free_volume = nextVolume;
@@ -150,72 +180,58 @@ export class AccountInfoUnit extends BasicUnit {
     }
 
     // 检查因为报价变化导致的头寸变化
-    const positions = Object.values(this.mapAccountIdToPositionIdToPosition[accountId] ?? {})
-      .filter((pos) => pos.volume > 0) // 过滤掉空的头寸
-      .map((position): IPosition => {
+    const mapProductIdToPositions = (this.mapAccountIdToProductIdToPositions[accountId] ??= {});
+    for (const product_id of dirtyProductIds) {
+      const influencedPositions = mapProductIdToPositions[product_id];
+      if (!influencedPositions) continue;
+      for (const position of influencedPositions) {
+        if (!(position.volume > 0)) continue; // 过滤掉空的头寸
         const product_id = position.product_id;
         const quote = this.quoteDataUnit.getQuote(accountId, product_id);
         const product = this.productDataUnit.getProduct(accountId, product_id);
         if (product && quote) {
           const closable_price = position.direction === 'LONG' ? quote.bid : quote.ask;
-          const floating_profit = getProfit(
-            product,
-            position.position_price,
-            closable_price,
-            position.volume,
-            position.direction!,
-            theAccountInfo.money.currency,
-            (product_id) => this.quoteDataUnit.getQuote(accountId, product_id),
-          );
-          const nextPosition = {
-            ...position,
-            closable_price,
-            floating_profit,
-          };
-          this.mapAccountIdToPositionIdToPosition[accountId][nextPosition.position_id] = nextPosition;
-          return nextPosition;
+          const floating_profit =
+            getProfit(
+              product,
+              position.position_price,
+              closable_price,
+              position.volume,
+              position.direction!,
+              theAccountInfo.money.currency,
+              (product_id) => this.quoteDataUnit.getQuote(accountId, product_id),
+            ) || 0;
+          position.closable_price = closable_price;
+          position.floating_profit = floating_profit;
+          theAccountInfo.money.profit +=
+            floating_profit - (this.mapPositionIdToFloatingProfit[position.position_id] || 0);
+          this.mapPositionIdToFloatingProfit[position.position_id] = floating_profit;
+        }
+        // 维护账户保证金
+        if (product) {
+          const used =
+            getMargin(
+              product,
+              position.position_price,
+              position.volume,
+              position.direction!,
+              theAccountInfo.money.currency,
+              (product_id) => this.quoteDataUnit.getQuote(accountId, product_id),
+            ) / (theAccountInfo.money.leverage ?? 1) || 0;
+          theAccountInfo.money.used += used - (this.mapPositionIdToUsed[position.position_id] || 0);
+          this.mapPositionIdToUsed[position.position_id] = used;
         }
-        return position;
-      });
-    // 维护账户保证金
-    const used = positions.reduce((acc, cur) => {
-      const product = this.productDataUnit.getProduct(accountId, cur.product_id);
-      if (!product) {
-        return acc;
       }
-      return (
-        acc +
-        getMargin(
-          product,
-          cur.position_price,
-          cur.volume,
-          cur.direction!,
-          theAccountInfo.money.currency,
-          (product_id) => this.quoteDataUnit.getQuote(accountId, product_id),
-        ) /
-          (theAccountInfo.money.leverage ?? 1)
-      );
-    }, 0);
+    }
+
     // 维护账户
     const balance = this.mapAccountIdToBalance[accountId] || 0;
-    const profit = positions.reduce((acc, cur) => acc + cur.floating_profit, 0);
-    const equity = balance + profit;
-    const free = equity - used;
-    const money: IAccountMoney = {
-      ...theAccountInfo.money,
-      equity,
-      balance,
-      profit,
-      used,
-      free,
-    };
-    this.mapAccountIdToAccountInfo.set(accountId, {
-      ...theAccountInfo,
-      updated_at: this.kernel.currentTimestamp,
-      money: money,
-      currencies: [money],
-      positions,
-    });
+    const equity = balance + theAccountInfo.money.profit;
+    const money = theAccountInfo.money;
+    money.equity = equity;
+    money.balance = balance;
+    money.free = equity - theAccountInfo.money.used;
+    theAccountInfo.updated_at = this.kernel.currentTimestamp;
   }
 
   onEvent(): void | Promise<void> {

libraries/kernel/src/units/QuoteDataUnit.ts

@@ -18,6 +18,11 @@ export class QuoteDataUnit extends BasicUnit {
   private productDataUnit: ProductDataUnit | undefined;
   private mapProductIdToQuote: Record<string, Record<string, IQuote>> = {};
   private mapDatasourceIdMapProductIdToAccountIds: Record<string, Record<string, Set<string>>> = {};
+  private currentEventId: number = -1;
+  /**
+   * 本事件内发生变化的产品ID
+   */
+  dirtyProductIds: Set<string> = new Set();
 
   onInit(): void | Promise<void> {
     this.adrUnit = this.kernel.findUnit(AccountDatasourceRelationUnit);
@@ -38,13 +43,19 @@ export class QuoteDataUnit extends BasicUnit {
   }
 
   private _updateQuote(quote: IQuote) {
+    if (this.currentEventId !== this.kernel.currentEventId) {
+      this.dirtyProductIds.clear();
+      this.currentEventId = this.kernel.currentEventId;
+    }
+
     (this.mapProductIdToQuote[quote.datasource_id] ??= {})[quote.product_id] = quote;
     (this.mapProductIdToQuote[''] ??= {})[quote.product_id] = quote;
     this.mapDatasourceIdMapProductIdToAccountIds[quote.datasource_id]?.[quote.product_id]?.forEach(
       (account_id) => {
         (this.mapProductIdToQuote[account_id] ??= {})[quote.product_id] = quote;
       },
     );
+    this.dirtyProductIds.add(quote.product_id);
   }
 
   updateQuote(datasource_id: string, product_id: string, ask: number, bid: number) {