Просмотр исходного кода

perf(statistics): batch total cost queries on quotas/users page

The previous fix removed the date filter but the page still crashed
PostgreSQL shared memory. Root cause: N+M concurrent SUM(cost_usd)
queries (one per user + one per key) via Promise.all.

Add sumUserTotalCostBatch and sumKeyTotalCostBatchByIds that use
GROUP BY to aggregate all users/keys in a single SQL query each.
Rewrite quotas/users page to use 2 batch queries instead of N+M
individual ones.
ding113 14 часов назад
Родитель
Сommit
455a332997
2 измененных файлов с 118 добавлено и 59 удалено
  1. 49 57
      src/app/[locale]/dashboard/quotas/users/page.tsx
  2. 69 2
      src/repository/statistics.ts

+ 49 - 57
src/app/[locale]/dashboard/quotas/users/page.tsx

@@ -6,7 +6,7 @@ import { QuotaToolbar } from "@/components/quota/quota-toolbar";
 import { Alert, AlertDescription } from "@/components/ui/alert";
 import { Link, redirect } from "@/i18n/routing";
 import { getSession } from "@/lib/auth";
-import { sumKeyTotalCostById, sumUserTotalCost } from "@/repository/statistics";
+import { sumKeyTotalCostBatchByIds, sumUserTotalCostBatch } from "@/repository/statistics";
 import { getSystemSettings } from "@/repository/system-config";
 import { UsersQuotaSkeleton } from "../_components/users-quota-skeleton";
 import type { UserKeyWithUsage, UserQuotaWithUsage } from "./_components/types";
@@ -15,71 +15,63 @@ import { UsersQuotaClient } from "./_components/users-quota-client";
 // Force dynamic rendering (this page needs real-time data and auth)
 export const dynamic = "force-dynamic";
 
-// Infinity means "all time" - no date filter applied to the query
-const ALL_TIME_MAX_AGE_DAYS = Infinity;
-
 async function getUsersWithQuotas(): Promise<UserQuotaWithUsage[]> {
   const users = await getUsers();
 
-  const usersWithQuotas = await Promise.all(
-    users.map(async (user) => {
-      // Fetch quota usage and total cost in parallel
-      const [quotaResult, totalUsage] = await Promise.all([
-        getUserLimitUsage(user.id),
-        sumUserTotalCost(user.id, ALL_TIME_MAX_AGE_DAYS),
-      ]);
-
-      // Map keys with their total usage
-      const keysWithUsage: UserKeyWithUsage[] = await Promise.all(
-        user.keys.map(async (key) => {
-          const keyTotalUsage = await sumKeyTotalCostById(key.id, ALL_TIME_MAX_AGE_DAYS);
-          return {
-            id: key.id,
-            name: key.name,
-            status: key.status,
-            todayUsage: key.todayUsage,
-            totalUsage: keyTotalUsage,
-            limit5hUsd: key.limit5hUsd,
-            limitDailyUsd: key.limitDailyUsd,
-            limitWeeklyUsd: key.limitWeeklyUsd,
-            limitMonthlyUsd: key.limitMonthlyUsd,
-            limitTotalUsd: key.limitTotalUsd ?? null,
-            limitConcurrentSessions: key.limitConcurrentSessions,
-            dailyResetMode: key.dailyResetMode,
-            dailyResetTime: key.dailyResetTime,
-          };
-        })
-      );
-
-      return {
-        id: user.id,
-        name: user.name,
-        note: user.note,
-        role: user.role,
-        isEnabled: user.isEnabled,
-        expiresAt: user.expiresAt ?? null,
-        providerGroup: user.providerGroup,
-        tags: user.tags,
-        quota: quotaResult.ok ? quotaResult.data : null,
-        limit5hUsd: user.limit5hUsd ?? null,
-        limitWeeklyUsd: user.limitWeeklyUsd ?? null,
-        limitMonthlyUsd: user.limitMonthlyUsd ?? null,
-        limitTotalUsd: user.limitTotalUsd ?? null,
-        limitConcurrentSessions: user.limitConcurrentSessions ?? null,
-        totalUsage,
-        keys: keysWithUsage,
-      };
-    })
-  );
-
-  return usersWithQuotas;
+  const allUserIds = users.map((u) => u.id);
+  const allKeyIds = users.flatMap((u) => u.keys.map((k) => k.id));
+
+  // 3 queries total instead of N+M individual SUM queries
+  const [quotaResults, userCostMap, keyCostMap] = await Promise.all([
+    Promise.all(users.map((u) => getUserLimitUsage(u.id))),
+    sumUserTotalCostBatch(allUserIds),
+    sumKeyTotalCostBatchByIds(allKeyIds),
+  ]);
+
+  return users.map((user, idx) => {
+    const quotaResult = quotaResults[idx];
+
+    const keysWithUsage: UserKeyWithUsage[] = user.keys.map((key) => ({
+      id: key.id,
+      name: key.name,
+      status: key.status,
+      todayUsage: key.todayUsage,
+      totalUsage: keyCostMap.get(key.id) ?? 0,
+      limit5hUsd: key.limit5hUsd,
+      limitDailyUsd: key.limitDailyUsd,
+      limitWeeklyUsd: key.limitWeeklyUsd,
+      limitMonthlyUsd: key.limitMonthlyUsd,
+      limitTotalUsd: key.limitTotalUsd ?? null,
+      limitConcurrentSessions: key.limitConcurrentSessions,
+      dailyResetMode: key.dailyResetMode,
+      dailyResetTime: key.dailyResetTime,
+    }));
+
+    return {
+      id: user.id,
+      name: user.name,
+      note: user.note,
+      role: user.role,
+      isEnabled: user.isEnabled,
+      expiresAt: user.expiresAt ?? null,
+      providerGroup: user.providerGroup,
+      tags: user.tags,
+      quota: quotaResult.ok ? quotaResult.data : null,
+      limit5hUsd: user.limit5hUsd ?? null,
+      limitWeeklyUsd: user.limitWeeklyUsd ?? null,
+      limitMonthlyUsd: user.limitMonthlyUsd ?? null,
+      limitTotalUsd: user.limitTotalUsd ?? null,
+      limitConcurrentSessions: user.limitConcurrentSessions ?? null,
+      totalUsage: userCostMap.get(user.id) ?? 0,
+      keys: keysWithUsage,
+    };
+  });
 }
 
 export default async function UsersQuotaPage({ params }: { params: Promise<{ locale: string }> }) {
   const { locale } = await params;
   const session = await getSession();
 
-  // 权限检查:仅 admin 用户可访问
   if (!session || session.user.role !== "admin") {
     return redirect({ href: session ? "/dashboard/my-quota" : "/login", locale });
   }

+ 69 - 2
src/repository/statistics.ts

@@ -1,6 +1,6 @@
 "use server";
 
-import { and, eq, gte, isNull, lt, sql } from "drizzle-orm";
+import { and, eq, gte, inArray, isNull, lt, sql } from "drizzle-orm";
 import { db } from "@/drizzle/db";
 import { keys, messageRequest } from "@/drizzle/schema";
 import { resolveSystemTimezone } from "@/lib/utils/timezone";
@@ -812,7 +812,74 @@ export async function sumUserTotalCost(userId: number, maxAgeDays: number = 365)
 }
 
 /**
- * 查询供应商历史总消费
+ * Batch query: all-time total cost grouped by user_id (single SQL query)
+ * @param userIds - Array of user IDs
+ * @returns Map of userId -> totalCost
+ */
+export async function sumUserTotalCostBatch(userIds: number[]): Promise<Map<number, number>> {
+  const result = new Map<number, number>();
+  if (userIds.length === 0) return result;
+
+  const rows = await db
+    .select({
+      userId: messageRequest.userId,
+      total: sql<number>`COALESCE(SUM(${messageRequest.costUsd}), 0)`,
+    })
+    .from(messageRequest)
+    .where(
+      and(
+        inArray(messageRequest.userId, userIds),
+        isNull(messageRequest.deletedAt),
+        EXCLUDE_WARMUP_CONDITION
+      )
+    )
+    .groupBy(messageRequest.userId);
+
+  for (const id of userIds) {
+    result.set(id, 0);
+  }
+  for (const row of rows) {
+    result.set(row.userId, Number(row.total || 0));
+  }
+  return result;
+}
+
+/**
+ * Batch query: all-time total cost grouped by key_id (single SQL query via JOIN)
+ * @param keyIds - Array of key IDs
+ * @returns Map of keyId -> totalCost
+ */
+export async function sumKeyTotalCostBatchByIds(keyIds: number[]): Promise<Map<number, number>> {
+  const result = new Map<number, number>();
+  if (keyIds.length === 0) return result;
+
+  const rows = await db
+    .select({
+      keyId: keys.id,
+      total: sql<number>`COALESCE(SUM(${messageRequest.costUsd}), 0)`,
+    })
+    .from(keys)
+    .leftJoin(
+      messageRequest,
+      and(
+        eq(messageRequest.key, keys.key),
+        isNull(messageRequest.deletedAt),
+        EXCLUDE_WARMUP_CONDITION
+      )
+    )
+    .where(inArray(keys.id, keyIds))
+    .groupBy(keys.id);
+
+  for (const id of keyIds) {
+    result.set(id, 0);
+  }
+  for (const row of rows) {
+    result.set(row.keyId, Number(row.total || 0));
+  }
+  return result;
+}
+
+/**
  * 用于供应商总消费限额检查(limit_total_usd)。
  *
  * 重要语义: