feat: 实现公平捞瓶和领取租约

This commit is contained in:
root
2026-09-15 13:56:56 +08:00
parent 969a10d664
commit 69fc8d51ce
23 changed files with 1221 additions and 20 deletions
+40
View File
@@ -0,0 +1,40 @@
import type { PrismaClient } from "@prisma/client";
export class LeaseReaper {
constructor(private readonly prisma: PrismaClient) {}
async runOnce(): Promise<boolean> {
return this.prisma.$transaction(async (tx) => {
const rows = await tx.$queryRaw<Array<{ id: string; bottleId: string }>>`
SELECT "id", "bottle_id" AS "bottleId"
FROM "bottle_pick_leases"
WHERE "status" = 'ACTIVE' AND "expires_at" <= now()
ORDER BY "expires_at", "id"
FOR UPDATE SKIP LOCKED LIMIT 1`;
const lease = rows[0];
if (!lease) return false;
const expired = await tx.bottlePickLease.updateMany({
where: {
id: lease.id,
status: "ACTIVE",
expiresAt: { lte: new Date() },
},
data: { status: "EXPIRED", endedAt: new Date() },
});
if (expired.count !== 1) return true;
await tx.bottle.updateMany({
where: {
id: lease.bottleId,
poolStatus: "LEASED",
activeLeaseId: lease.id,
},
data: {
poolStatus: "IN_POOL",
activeLeaseId: null,
version: { increment: 1 },
},
});
return true;
});
}
}
+125
View File
@@ -0,0 +1,125 @@
/* eslint-disable @typescript-eslint/no-unsafe-assignment */
import { PrismaClient } from "@prisma/client";
import { randomUUID } from "node:crypto";
import { afterAll, beforeAll, beforeEach, describe, expect, it } from "vitest";
import { assertSafeTestDatabaseUrl } from "../../../prisma/database-safety.js";
import { LeaseReaper } from "./lease-reaper.processor.js";
const prisma = new PrismaClient();
describe("lease reaper with real PostgreSQL", () => {
beforeAll(async () => {
assertSafeTestDatabaseUrl(process.env.DATABASE_URL ?? "");
await prisma.$connect();
});
beforeEach(async () => {
await prisma.$executeRawUnsafe(`TRUNCATE TABLE "accounts" CASCADE`);
});
afterAll(() => prisma.$disconnect());
async function fixture(expiresAt: Date) {
const [author, picker] = await Promise.all([
prisma.account.create({
data: { phoneCiphertext: Buffer.from("a"), phoneHmac: randomUUID() },
}),
prisma.account.create({
data: { phoneCiphertext: Buffer.from("p"), phoneHmac: randomUUID() },
}),
]);
const bottle = await prisma.bottle.create({
data: {
authorId: author.id,
clientRequestId: randomUUID(),
contentText: "secret body",
reviewStatus: "APPROVED",
poolStatus: "IN_POOL",
},
});
const leaseId = randomUUID();
await prisma.$transaction(async (tx) => {
await tx.bottle.update({
where: { id: bottle.id },
data: {
poolStatus: "LEASED",
activeLeaseId: leaseId,
version: { increment: 1 },
},
});
await tx.bottlePickLease.create({
data: {
id: leaseId,
bottleId: bottle.id,
pickerId: picker.id,
leaseTokenHash: randomUUID(),
leaseTokenCiphertext: Buffer.from("ciphertext"),
expiresAt,
},
});
});
return { bottle, leaseId };
}
it("atomically expires a due active lease and returns its current bottle to the pool", async () => {
const { bottle, leaseId } = await fixture(new Date(Date.now() - 1000));
expect(await new LeaseReaper(prisma).runOnce()).toBe(true);
expect(
await prisma.bottlePickLease.findUniqueOrThrow({
where: { id: leaseId },
}),
).toMatchObject({ status: "EXPIRED", endedAt: expect.any(Date) });
expect(
await prisma.bottle.findUniqueOrThrow({ where: { id: bottle.id } }),
).toMatchObject({ poolStatus: "IN_POOL", activeLeaseId: null, version: 3 });
});
it("does not reap a future lease", async () => {
const { bottle, leaseId } = await fixture(new Date(Date.now() + 60_000));
expect(await new LeaseReaper(prisma).runOnce()).toBe(false);
expect(
(
await prisma.bottlePickLease.findUniqueOrThrow({
where: { id: leaseId },
})
).status,
).toBe("ACTIVE");
expect(
(await prisma.bottle.findUniqueOrThrow({ where: { id: bottle.id } }))
.activeLeaseId,
).toBe(leaseId);
});
it("never releases a newer lease when processing stale old ownership", async () => {
const { bottle, leaseId: oldLeaseId } = await fixture(
new Date(Date.now() - 1000),
);
const picker = await prisma.account.create({
data: { phoneCiphertext: Buffer.from("n"), phoneHmac: randomUUID() },
});
const newLeaseId = randomUUID();
await prisma.$transaction(async (tx) => {
await tx.bottlePickLease.update({
where: { id: oldLeaseId },
data: { status: "EXPIRED", endedAt: new Date() },
});
await tx.bottle.update({
where: { id: bottle.id },
data: { activeLeaseId: newLeaseId },
});
await tx.bottlePickLease.create({
data: {
id: newLeaseId,
bottleId: bottle.id,
pickerId: picker.id,
leaseTokenHash: randomUUID(),
leaseTokenCiphertext: Buffer.from("ciphertext"),
expiresAt: new Date(Date.now() + 60_000),
},
});
});
expect(await new LeaseReaper(prisma).runOnce()).toBe(false);
expect(
(await prisma.bottle.findUniqueOrThrow({ where: { id: bottle.id } }))
.activeLeaseId,
).toBe(newLeaseId);
});
});
+16 -1
View File
@@ -1,7 +1,22 @@
import { describe, expect, it, vi } from "vitest";
import { handleTopLevelError, logTopLevelError, runWorker } from "./main.js";
import {
createCombinedWorker,
handleTopLevelError,
logTopLevelError,
runWorker,
} from "./main.js";
describe("worker lifecycle", () => {
it("runs both reaping and moderation every iteration without starvation", async () => {
const reaper = { runOnce: vi.fn().mockResolvedValue(true) };
const moderation = { runOnce: vi.fn().mockResolvedValue(false) };
const worker = createCombinedWorker(reaper, moderation);
await expect(worker.runOnce()).resolves.toBe(true);
expect(reaper.runOnce).toHaveBeenCalledTimes(1);
expect(moderation.runOnce).toHaveBeenCalledTimes(1);
});
it("waits for the current run before disconnecting on SIGTERM", async () => {
let finish!: () => void;
const runOnce = vi.fn(
+19 -1
View File
@@ -1,6 +1,7 @@
import { pathToFileURL } from "node:url";
import { PrismaClient } from "@prisma/client";
import { ModerationWorker } from "./moderation-worker.js";
import { LeaseReaper } from "./lease-reaper.processor.js";
type SignalSource = {
once(signal: "SIGTERM" | "SIGINT", listener: () => void): unknown;
@@ -16,6 +17,21 @@ type RunWorkerOptions = {
signals?: SignalSource;
};
type Worker = { runOnce(): Promise<boolean> };
export function createCombinedWorker(
reaper: Worker,
moderation: Worker,
): Worker {
return {
async runOnce() {
const reaped = await reaper.runOnce();
const moderated = await moderation.runOnce();
return reaped || moderated;
},
};
}
export async function runWorker(options: RunWorkerOptions) {
const signals = options.signals ?? process;
const sleep =
@@ -47,7 +63,9 @@ export async function runWorker(options: RunWorkerOptions) {
export async function main() {
const prisma = new PrismaClient();
const worker = new ModerationWorker(prisma);
const moderation = new ModerationWorker(prisma);
const reaper = new LeaseReaper(prisma);
const worker = createCombinedWorker(reaper, moderation);
await runWorker({
worker,
connect: () => prisma.$connect(),