fix: 限定捞瓶查询与重试成本

apps/api/src/match/match-retry-policy.spec.ts

@@ -0,0 +1,77 @@
+import { Prisma } from "@prisma/client";
+import { describe, expect, it } from "vitest";
+import {
+  CandidateBatchContended,
+  classifyPickError,
+  readMatchRetryConfig,
+  transactionLimits,
+} from "./match-retry-policy.js";
+
+const prismaError = (code: string, target?: string[]) =>
+  new Prisma.PrismaClientKnownRequestError("test", {
+    code,
+    clientVersion: "test",
+    ...(target ? { meta: { target } } : {}),
+  });
+
+describe("match pick retry policy", () => {
+  it("uses bounded defaults and caps each transaction by the overall deadline", () => {
+    const config = readMatchRetryConfig({});
+    expect(config).toEqual({
+      budgetMs: 5_000,
+      maxAttempts: 64,
+      transactionMaxWaitMs: 1_000,
+      transactionTimeoutMs: 2_000,
+    });
+    expect(transactionLimits(config, 14_500, 14_750)).toEqual({
+      maxWait: 125,
+      timeout: 125,
+    });
+    expect(transactionLimits(config, 15_000, 15_000)).toBeNull();
+  });
+
+  it.each([
+    ["MATCH_PICK_RETRY_BUDGET_MS", "0"],
+    ["MATCH_PICK_RETRY_BUDGET_MS", "30001"],
+    ["MATCH_PICK_MAX_ATTEMPTS", "101"],
+    ["MATCH_PICK_TRANSACTION_MAX_WAIT_MS", "10001"],
+    ["MATCH_PICK_TRANSACTION_TIMEOUT_MS", "nope"],
+  ])("rejects invalid or excessive %s", (name, value) => {
+    expect(() => readMatchRetryConfig({ [name]: value })).toThrow(name);
+  });
+
+  it("classifies exact pick races without hiding unrelated unique violations", () => {
+    expect(classifyPickError(prismaError("P2034"))).toBe("RETRY");
+    expect(classifyPickError(new CandidateBatchContended())).toBe("RETRY");
+    expect(
+      classifyPickError(prismaError("P2002", ["picker_id", "request_id"])),
+    ).toBe("READ_WINNER");
+    expect(
+      classifyPickError(
+        prismaError("P2002", ["bottle_pick_requests_picker_id_request_id_key"]),
+      ),
+    ).toBe("READ_WINNER");
+    expect(
+      classifyPickError(prismaError("P2002", ["bottle_id", "picker_id"])),
+    ).toBe("RETRY");
+    expect(
+      classifyPickError(
+        prismaError("P2002", ["bottle_pick_leases_one_active_per_bottle"]),
+      ),
+    ).toBe("RETRY");
+    expect(classifyPickError(prismaError("P2002", ["bottle_id"]))).toBe(
+      "RETRY",
+    );
+    expect(classifyPickError(prismaError("P2002", ["active_lease_id"]))).toBe(
+      "RETRY",
+    );
+    expect(
+      classifyPickError(prismaError("P2002", ["unrelated_bug_column"])),
+    ).toBe("THROW");
+    expect(
+      classifyPickError(prismaError("P2002", ["customer_lease_notes_key"])),
+    ).toBe("THROW");
+    expect(classifyPickError(prismaError("P2028"))).toBe("SERVICE_UNAVAILABLE");
+    expect(classifyPickError(new Error("programming bug"))).toBe("THROW");
+  });
+});

apps/api/src/match/match-retry-policy.ts

@@ -0,0 +1,116 @@
+import { Prisma } from "@prisma/client";
+
+export type MatchRetryConfig = {
+  budgetMs: number;
+  maxAttempts: number;
+  transactionMaxWaitMs: number;
+  transactionTimeoutMs: number;
+};
+
+export class CandidateBatchContended extends Error {
+  constructor() {
+    super("match pick candidate batch contended");
+    this.name = "CandidateBatchContended";
+  }
+}
+
+const boundedInt = (
+  environment: NodeJS.ProcessEnv,
+  name: string,
+  fallback: number,
+  maximum: number,
+): number => {
+  const raw = environment[name];
+  const value = raw === undefined ? fallback : Number(raw);
+  if (!Number.isInteger(value) || value <= 0 || value > maximum)
+    throw new Error(
+      `${name} must be a positive integer no greater than ${maximum}`,
+    );
+  return value;
+};
+
+export function readMatchRetryConfig(
+  environment: NodeJS.ProcessEnv,
+): MatchRetryConfig {
+  return {
+    budgetMs: boundedInt(
+      environment,
+      "MATCH_PICK_RETRY_BUDGET_MS",
+      5_000,
+      30_000,
+    ),
+    maxAttempts: boundedInt(environment, "MATCH_PICK_MAX_ATTEMPTS", 64, 100),
+    transactionMaxWaitMs: boundedInt(
+      environment,
+      "MATCH_PICK_TRANSACTION_MAX_WAIT_MS",
+      1_000,
+      10_000,
+    ),
+    transactionTimeoutMs: boundedInt(
+      environment,
+      "MATCH_PICK_TRANSACTION_TIMEOUT_MS",
+      2_000,
+      10_000,
+    ),
+  };
+}
+
+export function transactionLimits(
+  config: MatchRetryConfig,
+  nowMs: number,
+  deadlineMs: number,
+): { maxWait: number; timeout: number } | null {
+  const remaining = Math.floor(deadlineMs - nowMs);
+  if (remaining <= 1) return null;
+  // Prisma's queue wait and transaction execution timeout are sequential. Split
+  // the final sliver so their sum can never exceed the overall deadline.
+  const maxWait = Math.min(
+    config.transactionMaxWaitMs,
+    Math.floor(remaining / 2),
+  );
+  const timeout = Math.min(config.transactionTimeoutMs, remaining - maxWait);
+  return maxWait > 0 && timeout > 0 ? { maxWait, timeout } : null;
+}
+
+export type PickErrorAction =
+  "RETRY" | "READ_WINNER" | "SERVICE_UNAVAILABLE" | "THROW";
+
+const uniqueTarget = (target: unknown): string[] =>
+  (Array.isArray(target) ? target : [target])
+    .filter((field): field is string => typeof field === "string")
+    .map((field) => field.toLowerCase());
+
+const isPickRequestRace = (target: unknown): boolean => {
+  const fields = uniqueTarget(target);
+  return (
+    fields.includes("bottle_pick_requests_picker_id_request_id_key") ||
+    (fields.length === 2 &&
+      fields.includes("picker_id") &&
+      fields.includes("request_id"))
+  );
+};
+
+const isCandidateRace = (target: unknown): boolean => {
+  const fields = uniqueTarget(target);
+  return (
+    fields.includes("bottle_pick_leases_one_active_per_bottle") ||
+    fields.includes("bottle_pick_history_bottle_id_picker_id_key") ||
+    (fields.length === 1 &&
+      (fields.includes("bottle_id") || fields.includes("active_lease_id"))) ||
+    (fields.length === 2 &&
+      fields.includes("bottle_id") &&
+      fields.includes("picker_id"))
+  );
+};
+
+export function classifyPickError(error: unknown): PickErrorAction {
+  if (error instanceof CandidateBatchContended) return "RETRY";
+  if (!(error instanceof Prisma.PrismaClientKnownRequestError)) return "THROW";
+  if (error.code === "P2034") return "RETRY";
+  if (error.code === "P2028") return "SERVICE_UNAVAILABLE";
+  if (error.code === "P2002") {
+    if (isPickRequestRace(error.meta?.target)) return "READ_WINNER";
+    if (isCandidateRace(error.meta?.target)) return "RETRY";
+  }
+  return "THROW";
+}

apps/api/src/match/match.e2e-spec.ts

@@ -285,7 +285,9 @@ describe("match API with real PostgreSQL", () => {
     const results = await Promise.all(
       Array.from({ length: 8 }, () => pick(picker.authorization, key)),
     );
-    expect(results.every((x) => x.status === 201)).toBe(true);
+    expect(results.map((x) => [x.status, x.body.code])).toEqual(
+      Array.from({ length: 8 }, () => [201, "OK"]),
+    );
     expect(new Set(results.map((x) => x.body.data.lease.id))).toHaveLength(1);
     expect(new Set(results.map((x) => x.body.data.lease.token))).toHaveLength(
       1,
@@ -451,6 +453,56 @@ describe("match API with real PostgreSQL", () => {
     await pick(picker.authorization).expect(201);
   }, 15_000);
 
+  it("sustains twenty rounds of 21 concurrent distinct requests without picker-level serialization", async () => {
+    for (let round = 0; round < 20; round += 1) {
+      await prisma.$executeRawUnsafe(`TRUNCATE TABLE "accounts" CASCADE`);
+      const author = await actor(`stress-author-${round}`);
+      const picker = await actor(`stress-picker-${round}`);
+      await prisma.bottle.createMany({
+        data: Array.from({ length: 96 }, (_, index) => ({
+          authorId: author.id,
+          clientRequestId: `stress-${round}-${index}`,
+          contentText: `stress-${round}-${index}`,
+          reviewStatus: "APPROVED" as const,
+          poolStatus: "IN_POOL" as const,
+          approvedAt: new Date(),
+        })),
+      });
+
+      const results = await Promise.all(
+        Array.from({ length: 21 }, (_, index) =>
+          pick(picker.authorization, `stress-key-${round}-${index}`),
+        ),
+      );
+      expect(results.filter(({ status }) => status === 201)).toHaveLength(20);
+      expect(
+        results.filter(
+          ({ status, body }) =>
+            status === 429 && body.code === "BOTTLE_DAILY_LIMIT",
+        ),
+      ).toHaveLength(1);
+      expect(results.filter(({ status }) => status >= 500)).toHaveLength(0);
+      expect(
+        await prisma.dailyUsage.findFirstOrThrow({
+          where: { accountId: picker.id },
+        }),
+      ).toMatchObject({ bottlesPicked: 20 });
+      expect(
+        await prisma.bottlePickLease.count({ where: { pickerId: picker.id } }),
+      ).toBe(20);
+      expect(
+        await prisma.bottlePickHistory.count({
+          where: { pickerId: picker.id },
+        }),
+      ).toBe(20);
+      expect(
+        await prisma.bottlePickRequest.count({
+          where: { pickerId: picker.id },
+        }),
+      ).toBe(20);
+    }
+  }, 120_000);
+
   it.each(["SUSPENSION", "BAN"] as const)(
     "rejects active %s sanction",
     async (type) => {

apps/api/src/match/match.service.ts

@@ -1,7 +1,7 @@
 import { randomBytes, randomInt, randomUUID } from "node:crypto";
 import { HttpStatus, Inject, Injectable } from "@nestjs/common";
 import { ErrorCode } from "@drift/contracts";
-import { Prisma } from "@prisma/client";
+import type { Prisma } from "@prisma/client";
 import { DomainException } from "../common/domain.exception.js";
 import { PrismaService } from "../database/prisma.service.js";
 import { utc8UsageDate } from "../bottle/usage-date.js";
@@ -11,6 +11,12 @@ import {
   leaseHmac,
 } from "../auth/auth.config.js";
 import { CandidateRepository, type Candidate } from "./candidate.repository.js";
+import {
+  CandidateBatchContended,
+  classifyPickError,
+  readMatchRetryConfig,
+  transactionLimits,
+} from "./match-retry-policy.js";
 
 type LeaseWithBottle = {
   id: string;
@@ -48,6 +54,7 @@ export class MatchService {
     300_000,
     "BOTTLE_LEASE_TTL_MS",
   );
+  private readonly retry = readMatchRetryConfig(process.env);
 
   constructor(
     @Inject(PrismaService) private readonly prisma: PrismaService,
@@ -56,29 +63,14 @@ export class MatchService {
   ) {}
 
   async pick(pickerId: string, requestId: string) {
-    for (let attempt = 0; attempt < 100; attempt += 1) {
+    const deadline = performance.now() + this.retry.budgetMs;
+    for (let attempt = 0; attempt < this.retry.maxAttempts; attempt += 1) {
+      const limits = transactionLimits(this.retry, performance.now(), deadline);
+      if (!limits) break;
       try {
         const result = await this.prisma.$transaction(
           async (tx) => {
-            const [lock] = await tx.$queryRaw<Array<{ locked: boolean }>>`
+            const previous = await this.findPrevious(tx, pickerId, requestId);
-              SELECT pg_try_advisory_xact_lock(
-                hashtextextended(${`pick:${pickerId}`}, 0)
-              ) AS "locked"`;
-            if (!lock?.locked) return null;
-            const previous = await tx.bottlePickRequest.findUnique({
-              where: { pickerId_requestId: { pickerId, requestId } },
-              include: {
-                lease: {
-                  include: {
-                    bottle: {
-                      include: {
-                        author: { include: { anonymousProfile: true } },
-                      },
-                    },
-                  },
-                },
-              },
-            });
             if (previous?.lease) {
               let token: string;
               try {
@@ -97,17 +89,13 @@ export class MatchService {
             }
             await this.assertEligible(tx, pickerId);
             const usageDate = utc8UsageDate(new Date());
-            const current = await tx.dailyUsage.findUnique({
-              where: {
-                accountId_usageDate: {
-                  accountId: pickerId,
-                  usageDate: new Date(`${usageDate}T00:00:00.000Z`),
-                },
-              },
-            });
-            if ((current?.bottlesPicked ?? 0) >= 20) this.limit();
-
             const batch = await this.candidates.findBatch(tx, pickerId);
+            if (!batch.length)
+              throw new DomainException(
+                ErrorCode.BOTTLE_POOL_EMPTY,
+                "Bottle pool empty",
+                HttpStatus.NOT_FOUND,
+              );
             while (batch.length) {
               const index = Math.floor(Math.random() * batch.length);
               const candidate = batch.splice(index, 1)[0]!;
@@ -120,23 +108,51 @@ export class MatchService {
               );
               if (result) return result;
             }
-            throw new DomainException(
+            throw new CandidateBatchContended();
-              ErrorCode.BOTTLE_POOL_EMPTY,
-              "Bottle pool empty",
-              HttpStatus.NOT_FOUND,
-            );
           },
           {
             isolationLevel: "ReadCommitted",
-            maxWait: 10_000,
+            ...limits,
-            timeout: 10_000,
           },
         );
-        if (result) return result;
+        return result;
       } catch (error) {
-        if (!this.isContention(error)) throw error;
+        if (error instanceof DomainException) {
+          if (error.code === ErrorCode.BOTTLE_POOL_EMPTY) {
+            const winner = await this.readWinnerUntil(
+              pickerId,
+              requestId,
+              Math.min(deadline, performance.now() + 250),
+            );
+            if (winner) return winner;
+          }
+          throw error;
+        }
+        const action = classifyPickError(error);
+        if (action === "THROW") throw error;
+        if (action === "SERVICE_UNAVAILABLE") break;
+        if (action === "READ_WINNER") {
+          try {
+            const winner = await this.readWinner(pickerId, requestId);
+            if (winner) return winner;
+          } catch (readError) {
+            if (classifyPickError(readError) === "SERVICE_UNAVAILABLE") break;
+            throw readError;
+          }
+        }
       }
-      await new Promise((resolve) => setTimeout(resolve, randomInt(5, 21)));
+      const remaining = deadline - performance.now();
+      if (remaining <= 1) break;
+      await new Promise((resolve) =>
+        setTimeout(
+          resolve,
+          Math.min(
+            randomInt(5, 21) * 2 ** Math.min(attempt, 3),
+            100,
+            remaining,
+          ),
+        ),
+      );
     }
     throw new DomainException(
       ErrorCode.SERVICE_UNAVAILABLE,
@@ -145,13 +161,56 @@ export class MatchService {
     );
   }
 
-  private isContention(error: unknown) {
+  private findPrevious(
-    return (
+    client: Prisma.TransactionClient | PrismaService,
-      error instanceof Prisma.PrismaClientKnownRequestError &&
+    pickerId: string,
-      (error.code === "P2028" ||
+    requestId: string,
-        error.code === "P2034" ||
+  ) {
-        error.code === "P2002")
+    return client.bottlePickRequest.findUnique({
-    );
+      where: { pickerId_requestId: { pickerId, requestId } },
+      include: {
+        lease: {
+          include: {
+            bottle: {
+              include: { author: { include: { anonymousProfile: true } } },
+            },
+          },
+        },
+      },
+    });
+  }
+
+  private async readWinner(pickerId: string, requestId: string) {
+    const previous = await this.findPrevious(this.prisma, pickerId, requestId);
+    if (previous?.lease) {
+      try {
+        const token = decryptLeaseToken(
+          previous.lease.leaseTokenCiphertext,
+          previous.lease.id,
+        );
+        return this.response(previous.lease, token);
+      } catch {
+        throw new DomainException(
+          ErrorCode.SERVICE_UNAVAILABLE,
+          "Bottle pick temporarily unavailable",
+          HttpStatus.SERVICE_UNAVAILABLE,
+        );
+      }
+    }
+    return null;
+  }
+
+  private async readWinnerUntil(
+    pickerId: string,
+    requestId: string,
+    deadline: number,
+  ) {
+    while (performance.now() < deadline) {
+      const winner = await this.readWinner(pickerId, requestId);
+      if (winner) return winner;
+      await new Promise((resolve) => setTimeout(resolve, 5));
+    }
+    return null;
   }
 
   private async claim(

prisma/migrations/0009_bottle_pick_candidate_index/migration.sql

@@ -0,0 +1,5 @@
+CREATE INDEX "bottles_pick_candidates_id_idx"
+ON "bottles"("id")
+WHERE "review_status" = 'APPROVED'
+  AND "pool_status" = 'IN_POOL'
+  AND "active_lease_id" IS NULL;

prisma/schema.prisma

@@ -197,6 +197,7 @@ model RefreshToken {
 }
 
 model Bottle {
+  /// Candidate partial index bottles_pick_candidates_id_idx is database-only and managed in 0009 SQL.
   id              String              @id @default(uuid()) @db.Uuid
   authorId        String              @map("author_id") @db.Uuid
   clientRequestId String              @default(uuid()) @map("client_request_id") @db.VarChar(128)

tests/integration/database.spec.ts

@@ -93,6 +93,72 @@ describe("database authority constraints", () => {
     ).toBe(true);
   });
 
+  it("uses the database-only bottle pick candidate partial index", async () => {
+    const author = await createAccount("candidate-index-author");
+    await prisma.bottle.createMany({
+      data: [
+        ...Array.from({ length: 2_000 }, (_, index) => ({
+          authorId: author.id,
+          clientRequestId: `ineligible-${index}`,
+          contentText: `ineligible-${index}`,
+          reviewStatus: "REVIEWING" as const,
+          poolStatus: "CLOSED" as const,
+        })),
+        ...Array.from({ length: 3 }, (_, index) => ({
+          authorId: author.id,
+          clientRequestId: `eligible-${index}`,
+          contentText: `eligible-${index}`,
+          reviewStatus: "APPROVED" as const,
+          poolStatus: "IN_POOL" as const,
+          approvedAt: new Date(),
+        })),
+      ],
+    });
+
+    const [index] = await prisma.$queryRaw<Array<{ indexdef: string }>>`
+      SELECT indexdef FROM pg_indexes
+      WHERE schemaname = 'public' AND tablename = 'bottles'
+        AND indexname = 'bottles_pick_candidates_id_idx'`;
+    expect(index?.indexdef).toMatch(
+      /ON public\.bottles USING btree \(id\) WHERE .*review_status = 'APPROVED'.*pool_status = 'IN_POOL'.*active_lease_id IS NULL/,
+    );
+
+    const { plan, rows } = await prisma.$transaction(async (tx) => {
+      await tx.$executeRawUnsafe("SET LOCAL enable_seqscan = off");
+      const explained = await tx.$queryRawUnsafe<
+        Array<{ "QUERY PLAN": Array<{ Plan: { "Index Name"?: string } }> }>
+      >(`EXPLAIN (FORMAT JSON)
+         SELECT id FROM bottles
+         WHERE review_status = 'APPROVED'
+           AND pool_status = 'IN_POOL'
+           AND active_lease_id IS NULL
+         ORDER BY id LIMIT 32`);
+      const selected = await tx.$queryRaw<Array<{ content_text: string }>>`
+        SELECT content_text FROM bottles
+        WHERE review_status = 'APPROVED'
+          AND pool_status = 'IN_POOL'
+          AND active_lease_id IS NULL
+        ORDER BY id LIMIT 32`;
+      return { plan: explained[0]!["QUERY PLAN"][0]!.Plan, rows: selected };
+    });
+    const indexNames = (node: unknown): string[] => {
+      if (!node || typeof node !== "object") return [];
+      const record = node as Record<string, unknown>;
+      return [
+        ...(typeof record["Index Name"] === "string"
+          ? [record["Index Name"]]
+          : []),
+        ...Object.values(record).flatMap(indexNames),
+      ];
+    };
+    expect(indexNames(plan)).toContain("bottles_pick_candidates_id_idx");
+    expect(rows.map(({ content_text }) => content_text).sort()).toEqual([
+      "eligible-0",
+      "eligible-1",
+      "eligible-2",
+    ]);
+  });
+
   it("backfills and requires the session token version snapshot", async () => {
     const columns = await prisma.$queryRaw<
       Array<{ is_nullable: string; column_default: string | null }>