diff --git a/apps/api/src/match/match-retry-policy.spec.ts b/apps/api/src/match/match-retry-policy.spec.ts new file mode 100644 index 0000000..7ebdb70 --- /dev/null +++ b/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"); + }); +}); diff --git a/apps/api/src/match/match-retry-policy.ts b/apps/api/src/match/match-retry-policy.ts new file mode 100644 index 0000000..e036a3f --- /dev/null +++ b/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"; +} diff --git a/apps/api/src/match/match.e2e-spec.ts b/apps/api/src/match/match.e2e-spec.ts index 34619e6..2103381 100644 --- a/apps/api/src/match/match.e2e-spec.ts +++ b/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) => { diff --git a/apps/api/src/match/match.service.ts b/apps/api/src/match/match.service.ts index acf5747..807ac8c 100644 --- a/apps/api/src/match/match.service.ts +++ b/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>` - 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( diff --git a/prisma/migrations/0009_bottle_pick_candidate_index/migration.sql b/prisma/migrations/0009_bottle_pick_candidate_index/migration.sql new file mode 100644 index 0000000..435d6cf --- /dev/null +++ b/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; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index 2cf4699..bd6cf30 100644 --- a/prisma/schema.prisma +++ b/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) diff --git a/tests/integration/database.spec.ts b/tests/integration/database.spec.ts index bcbe5be..94c3b91 100644 --- a/tests/integration/database.spec.ts +++ b/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>` + 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>` + 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; + 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 }>