fix: 改进候选随机性与并发配额稳定性

2026-09-15 15:14:01 +08:00
parent 69fc8d51ce
commit 0c5cbd75d9
5 changed files with 384 additions and 69 deletions
@@ -1,3 +1,4 @@
+import { randomUUID } from "node:crypto";
 import { Injectable } from "@nestjs/common";
 import type { Prisma } from "@prisma/client";

@@ -14,14 +15,20 @@ export type Candidate = {

@Injectable()
 export class CandidateRepository {
-  async findBatch(tx: Prisma.TransactionClient, pickerId: string, limit = 32) {
-    return tx.$queryRaw<Candidate[]>`
+  async findBatch(
+    tx: Prisma.TransactionClient,
+    pickerId: string,
+    limit = 32,
+    pivot: string = randomUUID(),
+  ) {
+    const after = await tx.$queryRaw<Candidate[]>`
      SELECT b."id", b."version", b."author_id" AS "authorId",
             b."content_text" AS "contentText", p."public_id" AS "publicId",
             p."nickname", p."avatar_color" AS "avatarColor", p."bio"
      FROM "bottles" b
      JOIN "anonymous_profiles" p ON p."account_id" = b."author_id"
-      WHERE b."author_id" <> ${pickerId}::uuid
+      WHERE b."id" >= ${pivot}::uuid
+        AND b."author_id" <> ${pickerId}::uuid
        AND b."review_status" = 'APPROVED'
        AND b."pool_status" = 'IN_POOL'
        AND b."active_lease_id" IS NULL
@@ -35,7 +42,33 @@ export class CandidateRepository {
          WHERE (x."blocker_id" = ${pickerId}::uuid AND x."blocked_id" = b."author_id")
             OR (x."blocker_id" = b."author_id" AND x."blocked_id" = ${pickerId}::uuid)
        )
-      ORDER BY b."created_at", b."id"
+      ORDER BY b."id"
      LIMIT ${limit}`;
+    const remaining = limit - after.length;
+    if (remaining <= 0) return after;
+    const before = await tx.$queryRaw<Candidate[]>`
+      SELECT b."id", b."version", b."author_id" AS "authorId",
+             b."content_text" AS "contentText", p."public_id" AS "publicId",
+             p."nickname", p."avatar_color" AS "avatarColor", p."bio"
+      FROM "bottles" b
+      JOIN "anonymous_profiles" p ON p."account_id" = b."author_id"
+      WHERE b."id" < ${pivot}::uuid
+        AND b."author_id" <> ${pickerId}::uuid
+        AND b."review_status" = 'APPROVED'
+        AND b."pool_status" = 'IN_POOL'
+        AND b."active_lease_id" IS NULL
+        AND p."review_status" = 'APPROVED'
+        AND NOT EXISTS (
+          SELECT 1 FROM "bottle_pick_history" h
+          WHERE h."bottle_id" = b."id" AND h."picker_id" = ${pickerId}::uuid
+        )
+        AND NOT EXISTS (
+          SELECT 1 FROM "blocks" x
+          WHERE (x."blocker_id" = ${pickerId}::uuid AND x."blocked_id" = b."author_id")
+             OR (x."blocker_id" = b."author_id" AND x."blocked_id" = ${pickerId}::uuid)
+        )
+      ORDER BY b."id"
+      LIMIT ${remaining}`;
+    return [...after, ...before];
  }
 }
@@ -10,6 +10,7 @@ import { assertSafeTestDatabaseUrl } from "../../../../prisma/database-safety.js
 import { AppModule } from "../app.module.js";
 import { TokenService } from "../auth/token.service.js";
 import { configureApp } from "../main.js";
+import { CandidateRepository } from "./candidate.repository.js";

 process.env.NODE_ENV = "test";
 process.env.WEB_ORIGIN = "http://localhost:3000";
@@ -122,6 +123,49 @@ describe("match API with real PostgreSQL", () => {
      .set("Authorization", authorization)
      .send({ leaseId, token });

+  it("requires authentication for both match endpoints", async () => {
+    await request(app.getHttpServer())
+      .post("/api/v1/bottles/pick")
+      .set("Idempotency-Key", randomUUID())
+      .send({})
+      .expect(401);
+    await request(app.getHttpServer())
+      .post(`/api/v1/bottles/${randomUUID()}/return`)
+      .send({ leaseId: randomUUID(), token: "x".repeat(43) })
+      .expect(401);
+  });
+
+  it("uses a bounded UUID keyset pivot so every region beyond the oldest 32 is reachable", async () => {
+    const author = await actor("author-device");
+    const picker = await actor("picker-device");
+    const ids = Array.from({ length: 96 }, (_, index) => {
+      const prefix = index.toString(16).padStart(8, "0");
+      return `${prefix}-0000-4000-8000-000000000000`;
+    });
+    await prisma.bottle.createMany({
+      data: ids.map((id, index) => ({
+        id,
+        authorId: author.id,
+        clientRequestId: `pivot-${index}`,
+        contentText: `pivot-${index}`,
+        reviewStatus: "APPROVED" as const,
+        poolStatus: "IN_POOL" as const,
+        approvedAt: new Date(),
+      })),
+    });
+    const repository = app.get(CandidateRepository);
+    const front = await repository.findBatch(prisma, picker.id, 32, ids[0]);
+    const middle = await repository.findBatch(prisma, picker.id, 32, ids[40]);
+    const back = await repository.findBatch(prisma, picker.id, 32, ids[88]);
+
+    expect(front).toHaveLength(32);
+    expect(middle[0]?.id).toBe(ids[40]);
+    expect(back.map(({ id }) => id)).toContain(ids[95]);
+    expect(back.map(({ id }) => id)).toContain(ids[0]);
+    expect(new Set(back.map(({ id }) => id)).size).toBe(back.length);
+    expect(middle.map(({ id }) => id)).not.toEqual(front.map(({ id }) => id));
+  });
+
  it("picks an approved in-pool bottle and exposes only an anonymous public snapshot", async () => {
    const author = await actor("author-device");
    const picker = await actor("picker-device");
@@ -191,6 +235,34 @@ describe("match API with real PostgreSQL", () => {
        poolStatus: "IN_POOL",
      },
    });
+    const closed = await bottle(validAuthor.id, "closed");
+    await prisma.bottle.update({
+      where: { id: closed.id },
+      data: { poolStatus: "CLOSED" },
+    });
+    const leased = await bottle(validAuthor.id, "leased");
+    await prisma.bottle.update({
+      where: { id: leased.id },
+      data: { poolStatus: "LEASED" },
+    });
+    const activeLease = await bottle(validAuthor.id, "active lease id");
+    const leaseId = randomUUID();
+    await prisma.$transaction(async (tx) => {
+      await tx.bottle.update({
+        where: { id: activeLease.id },
+        data: { activeLeaseId: leaseId },
+      });
+      await tx.bottlePickLease.create({
+        data: {
+          id: leaseId,
+          bottleId: activeLease.id,
+          pickerId: picker.id,
+          leaseTokenHash: randomUUID(),
+          leaseTokenCiphertext: Buffer.from("ciphertext"),
+          expiresAt: new Date(Date.now() + 60_000),
+        },
+      });
+    });
    const picked = await pick(picker.authorization).expect(201);
    expect(picked.body.data.bottle.id).toBe(valid.id);
    expect(picked.body.data.bottle.id).not.toBe(own.id);
@@ -251,6 +323,37 @@ describe("match API with real PostgreSQL", () => {
    expect((await prisma.dailyUsage.findFirstOrThrow()).bottlesPicked).toBe(1);
  });

+  it.each(["RETURNED", "EXPIRED"] as const)(
+    "returns the original pick response when the idempotent lease is %s",
+    async (status) => {
+      const author = await actor("author-device");
+      const picker = await actor("picker-device");
+      const source = await bottle(author.id);
+      const key = randomUUID();
+      const first = await pick(picker.authorization, key).expect(201);
+      if (status === "RETURNED") {
+        await returnBottle(
+          picker.authorization,
+          source.id,
+          first.body.data.lease.id,
+          first.body.data.lease.token,
+        ).expect(201);
+      } else {
+        await prisma.bottlePickLease.update({
+          where: { id: first.body.data.lease.id },
+          data: { status: "EXPIRED", endedAt: new Date() },
+        });
+      }
+
+      const retry = await pick(picker.authorization, key).expect(201);
+      expect(retry.body.data).toEqual(first.body.data);
+      expect(await prisma.bottlePickLease.count()).toBe(1);
+      expect((await prisma.dailyUsage.findFirstOrThrow()).bottlesPicked).toBe(
+        1,
+      );
+    },
+  );
+
  it("issues independent random opaque tokens for different requests", async () => {
    const author = await actor("author-device");
    const picker = await actor("picker-device");
@@ -322,9 +425,16 @@ describe("match API with real PostgreSQL", () => {
  it("enforces twenty successful picks atomically and resets on another UTC+8 day", async () => {
    const author = await actor("author-device");
    const picker = await actor("picker-device");
-    await Promise.all(
-      Array.from({ length: 22 }, (_, i) => bottle(author.id, `bottle-${i}`)),
-    );
+    await prisma.bottle.createMany({
+      data: Array.from({ length: 22 }, (_, i) => ({
+        authorId: author.id,
+        clientRequestId: randomUUID(),
+        contentText: `bottle-${i}`,
+        reviewStatus: "APPROVED",
+        poolStatus: "IN_POOL",
+        approvedAt: new Date(),
+      })),
+    });
    const results = await Promise.all(
      Array.from({ length: 21 }, () => pick(picker.authorization)),
    );
@@ -339,7 +449,7 @@ describe("match API with real PostgreSQL", () => {
      data: { usageDate: new Date("2020-01-01T00:00:00.000Z") },
    });
    await pick(picker.authorization).expect(201);
-  });
+  }, 15_000);

  it.each(["SUSPENSION", "BAN"] as const)(
    "rejects active %s sanction",
@@ -1,7 +1,7 @@
-import { randomBytes, randomUUID } from "node:crypto";
+import { randomBytes, randomInt, randomUUID } from "node:crypto";
 import { HttpStatus, Inject, Injectable } from "@nestjs/common";
 import { ErrorCode } from "@drift/contracts";
-import type { Prisma } from "@prisma/client";
+import { 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";
@@ -56,69 +56,101 @@ export class MatchService {
  ) {}

  async pick(pickerId: string, requestId: string) {
-    return this.prisma.$transaction(
-      async (tx) => {
-        await tx.$executeRaw`SELECT pg_advisory_xact_lock(hashtextextended(${`pick:${pickerId}`}, 0))`;
-        const previous = await tx.bottlePickRequest.findUnique({
-          where: { pickerId_requestId: { pickerId, requestId } },
-          include: {
-            lease: {
+    for (let attempt = 0; attempt < 100; attempt += 1) {
+      try {
+        const result = await this.prisma.$transaction(
+          async (tx) => {
+            const [lock] = await tx.$queryRaw<Array<{ locked: boolean }>>`
+              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: {
-                bottle: {
-                  include: { author: { include: { anonymousProfile: true } } },
+                lease: {
+                  include: {
+                    bottle: {
+                      include: {
+                        author: { include: { anonymousProfile: true } },
+                      },
+                    },
+                  },
                },
              },
-            },
-          },
-        });
-        if (previous?.lease) {
-          let token: string;
-          try {
-            token = decryptLeaseToken(
-              previous.lease.leaseTokenCiphertext,
-              previous.lease.id,
-            );
-          } catch {
-            throw new DomainException(
-              ErrorCode.SERVICE_UNAVAILABLE,
-              "Bottle pick temporarily unavailable",
-              HttpStatus.SERVICE_UNAVAILABLE,
-            );
-          }
-          return this.response(previous.lease, token);
-        }
-        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();
+            });
+            if (previous?.lease) {
+              let token: string;
+              try {
+                token = decryptLeaseToken(
+                  previous.lease.leaseTokenCiphertext,
+                  previous.lease.id,
+                );
+              } catch {
+                throw new DomainException(
+                  ErrorCode.SERVICE_UNAVAILABLE,
+                  "Bottle pick temporarily unavailable",
+                  HttpStatus.SERVICE_UNAVAILABLE,
+                );
+              }
+              return this.response(previous.lease, token);
+            }
+            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);
-        while (batch.length) {
-          const index = Math.floor(Math.random() * batch.length);
-          const candidate = batch.splice(index, 1)[0]!;
-          const result = await this.claim(
-            tx,
-            pickerId,
-            requestId,
-            usageDate,
-            candidate,
-          );
-          if (result) return result;
-        }
-        throw new DomainException(
-          ErrorCode.BOTTLE_POOL_EMPTY,
-          "Bottle pool empty",
-          HttpStatus.NOT_FOUND,
+            const batch = await this.candidates.findBatch(tx, pickerId);
+            while (batch.length) {
+              const index = Math.floor(Math.random() * batch.length);
+              const candidate = batch.splice(index, 1)[0]!;
+              const result = await this.claim(
+                tx,
+                pickerId,
+                requestId,
+                usageDate,
+                candidate,
+              );
+              if (result) return result;
+            }
+            throw new DomainException(
+              ErrorCode.BOTTLE_POOL_EMPTY,
+              "Bottle pool empty",
+              HttpStatus.NOT_FOUND,
+            );
+          },
+          {
+            isolationLevel: "ReadCommitted",
+            maxWait: 10_000,
+            timeout: 10_000,
+          },
        );
-      },
-      { isolationLevel: "ReadCommitted" },
+        if (result) return result;
+      } catch (error) {
+        if (!this.isContention(error)) throw error;
+      }
+      await new Promise((resolve) => setTimeout(resolve, randomInt(5, 21)));
+    }
+    throw new DomainException(
+      ErrorCode.SERVICE_UNAVAILABLE,
+      "Bottle pick temporarily unavailable",
+      HttpStatus.SERVICE_UNAVAILABLE,
+    );
+  }
+
+  private isContention(error: unknown) {
+    return (
+      error instanceof Prisma.PrismaClientKnownRequestError &&
+      (error.code === "P2028" ||
+        error.code === "P2034" ||
+        error.code === "P2002")
    );
  }

@@ -9,8 +9,11 @@ WITH expired AS (
  RETURNING "bottle_id"
 )
 UPDATE "bottles" b
-SET "pool_status" = 'IN_POOL', "version" = b."version" + 1
+SET "pool_status" = 'IN_POOL',
+    "active_lease_id" = NULL,
+    "version" = b."version" + 1
 WHERE b."pool_status" = 'LEASED'
+  AND b."active_lease_id" IS NULL
  AND b."id" IN (SELECT "bottle_id" FROM expired);
 UPDATE "bottle_pick_leases"
 SET "lease_token_ciphertext" = decode(repeat('00', 48), 'hex')
@@ -0,0 +1,137 @@
+import { readFile } from "node:fs/promises";
+import { randomUUID } from "node:crypto";
+import { PrismaClient } from "@prisma/client";
+import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import {
+  assertSafeTestDatabaseUrl,
+  resolveTestDatabaseUrl,
+} from "../../prisma/database-safety";
+
+const databaseUrl = resolveTestDatabaseUrl();
+const prisma = new PrismaClient({ datasources: { db: { url: databaseUrl } } });
+
+async function executeStatements(
+  tx: Parameters<Parameters<typeof prisma.$transaction>[0]>[0],
+  sql: string,
+) {
+  for (const statement of sql
+    .split(";")
+    .map((part) => part.trim())
+    .filter(Boolean))
+    await tx.$executeRawUnsafe(statement);
+}
+
+describe("0008 legacy lease upgrade", () => {
+  beforeAll(async () => {
+    assertSafeTestDatabaseUrl(databaseUrl);
+    await prisma.$connect();
+  });
+  afterAll(async () => prisma.$disconnect());
+
+  it("expires legacy ACTIVE leases, releases their bottles, and preserves terminal rows", async () => {
+    const schema = `migration_0008_${randomUUID().replaceAll("-", "")}`;
+    const authorId = randomUUID();
+    const pickerId = randomUUID();
+    const activeBottleId = randomUUID();
+    const returnedBottleId = randomUUID();
+    const activeLeaseId = randomUUID();
+    const returnedLeaseId = randomUUID();
+    const migration = await readFile(
+      new URL(
+        "../../prisma/migrations/0008_bottle_pick_leases/migration.sql",
+        import.meta.url,
+      ),
+      "utf8",
+    );
+
+    try {
+      await prisma.$transaction(async (tx) => {
+        await tx.$executeRawUnsafe(`CREATE SCHEMA "${schema}"`);
+        await tx.$executeRawUnsafe(`SET LOCAL search_path TO "${schema}"`);
+        await executeStatements(
+          tx,
+          `
+          CREATE TYPE "BottlePoolStatus" AS ENUM ('IN_POOL', 'LEASED', 'CONSUMED', 'REMOVED', 'CLOSED');
+          CREATE TYPE "BottlePickLeaseStatus" AS ENUM ('ACTIVE', 'RETURNED', 'EXPIRED', 'CONSUMED');
+          CREATE TABLE "accounts" ("id" UUID PRIMARY KEY);
+          CREATE TABLE "bottles" (
+            "id" UUID PRIMARY KEY,
+            "pool_status" "BottlePoolStatus" NOT NULL,
+            "version" INTEGER NOT NULL
+          );
+          CREATE TABLE "bottle_pick_leases" (
+            "id" UUID PRIMARY KEY,
+            "bottle_id" UUID NOT NULL REFERENCES "bottles"("id"),
+            "picker_id" UUID NOT NULL REFERENCES "accounts"("id"),
+            "lease_token_hash" VARCHAR(255) NOT NULL UNIQUE,
+            "status" "BottlePickLeaseStatus" NOT NULL,
+            "ended_at" TIMESTAMPTZ(3),
+            "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP
+          );
+          INSERT INTO "accounts" ("id") VALUES ('${authorId}'), ('${pickerId}');
+          INSERT INTO "bottles" ("id", "pool_status", "version") VALUES
+            ('${activeBottleId}', 'LEASED', 4),
+            ('${returnedBottleId}', 'CLOSED', 7);
+          INSERT INTO "bottle_pick_leases"
+            ("id", "bottle_id", "picker_id", "lease_token_hash", "status", "ended_at") VALUES
+            ('${activeLeaseId}', '${activeBottleId}', '${pickerId}', 'active-hash', 'ACTIVE', NULL),
+            ('${returnedLeaseId}', '${returnedBottleId}', '${pickerId}', 'returned-hash', 'RETURNED', CURRENT_TIMESTAMP);
+        `,
+        );
+        await executeStatements(tx, migration);
+
+        const leases = await tx.$queryRawUnsafe<
+          Array<{
+            id: string;
+            status: string;
+            endedAt: Date | null;
+            ciphertextHex: string;
+          }>
+        >(`
+          SELECT "id", "status", "ended_at" AS "endedAt",
+                 encode("lease_token_ciphertext", 'hex') AS "ciphertextHex"
+          FROM "bottle_pick_leases" ORDER BY "id"
+        `);
+        const bottles = await tx.$queryRawUnsafe<
+          Array<{ id: string; poolStatus: string; version: number }>
+        >(`
+          SELECT "id", "pool_status" AS "poolStatus", "version"
+          FROM "bottles" ORDER BY "id"
+        `);
+        const ciphertextColumn = await tx.$queryRawUnsafe<
+          Array<{ isNullable: string }>
+        >(`
+          SELECT is_nullable AS "isNullable" FROM information_schema.columns
+          WHERE table_schema = '${schema}' AND table_name = 'bottle_pick_leases'
+            AND column_name = 'lease_token_ciphertext'
+        `);
+
+        expect(leases.find(({ id }) => id === activeLeaseId)).toMatchObject({
+          status: "EXPIRED",
+          endedAt: expect.any(Date) as Date,
+          ciphertextHex: "00".repeat(48),
+        });
+        expect(leases.find(({ id }) => id === returnedLeaseId)).toMatchObject({
+          status: "RETURNED",
+          endedAt: expect.any(Date) as Date,
+          ciphertextHex: "00".repeat(48),
+        });
+        expect(bottles.find(({ id }) => id === activeBottleId)).toMatchObject({
+          poolStatus: "IN_POOL",
+          version: 5,
+        });
+        expect(bottles.find(({ id }) => id === returnedBottleId)).toMatchObject(
+          {
+            poolStatus: "CLOSED",
+            version: 7,
+          },
+        );
+        expect(ciphertextColumn).toEqual([{ isNullable: "NO" }]);
+      });
+    } finally {
+      await prisma.$executeRawUnsafe(
+        `DROP SCHEMA IF EXISTS "${schema}" CASCADE`,
+      );
+    }
+  });
+});