jkjoy/plp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
root
2026-09-15 21:13:28 +08:00
parent 0c5cbd75d9
commit c94d7b1af0
7 changed files with 425 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/src/match/match-retry-policy.spec.ts
+77
View File
@@ -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
+116
View File
@@ -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
+53 -1
View File
@@ -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
+107 -48
View File
@@ -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 }>>`
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 } },
},
},
},
},
},
});
const previous = await this.findPrevious(tx, pickerId, requestId);
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(
ErrorCode.BOTTLE_POOL_EMPTY,
"Bottle pool empty",
HttpStatus.NOT_FOUND,
);
throw new CandidateBatchContended();
},
{
isolationLevel: "ReadCommitted",
maxWait: 10_000,
timeout: 10_000,
...limits,
},
);
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) {
return (
error instanceof Prisma.PrismaClientKnownRequestError &&
(error.code === "P2028" ||
error.code === "P2034" ||
error.code === "P2002")
);
private findPrevious(
client: Prisma.TransactionClient | PrismaService,
pickerId: string,
requestId: string,
) {
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
+5
View File
@@ -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
+1
View File
@@ -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
+66
View File
@@ -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 }>