From add6d1fadc7694197bde7520212e3d194e4931c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Mon, 28 Sep 2026 00:34:11 +0200 Subject: [PATCH] fix(restore): order the restore and the consolidation passes by one lock A restore on a freshly booted instance failed its clearing step with a deadlock (40P01) while the boot-time step consolidation was folding the same account. A consolidation day updates or mints the day's total and then soft-deletes the day's per-sample rows; the restore deletes every reading of the account in one statement. When the restore reached a day's samples before its total, each side held what the other needed next. The restore already took an advisory lock per account before its first delete. Every day transaction of the passes built on runConsolidation (step and mean consolidation, the per-sample cumulative drain, the dense intraday retention) now takes the same lock in shared mode as its first statement. A restore that starts mid-day waits for that day to commit; a day that would start while a restore runs is refused at once, and the pass leaves that account for its next run instead of folding rows the restore is about to replace. The passes still share the lock with each other. A new integration test builds the lock order that deadlocked, with a row trigger holding the consolidation on the day's total while the restore starts; before this change the consolidation was cancelled with a deadlock. A guard test requires every transaction in a pass on runConsolidation to open with the shared lock. Refs #1031 --- .../restore-lock-consolidation-guard.test.ts | 75 +++++ src/lib/export/restore-backup.ts | 20 +- src/lib/export/restore-lock.ts | 80 +++++ .../__tests__/consolidate-daily-mean.test.ts | 8 + .../dense-intraday-pulse-summary.test.ts | 8 + .../dense-intraday-resting-mint.test.ts | 8 + .../dense-intraday-retention.test.ts | 8 + .../dense-intraday-spo2-summary.test.ts | 8 + .../drain-per-sample-cumulative.test.ts | 8 + .../measurements/consolidate-daily-mean.ts | 3 + .../measurements/consolidate-legacy-steps.ts | 3 + src/lib/measurements/consolidation-base.ts | 27 +- .../measurements/dense-intraday-retention.ts | 3 + .../drain-per-sample-cumulative.ts | 3 + .../restore-consolidation-deadlock.test.ts | 277 ++++++++++++++++++ 15 files changed, 531 insertions(+), 8 deletions(-) create mode 100644 src/__tests__/restore-lock-consolidation-guard.test.ts create mode 100644 src/lib/export/restore-lock.ts create mode 100644 tests/integration/restore-consolidation-deadlock.test.ts diff --git a/src/__tests__/restore-lock-consolidation-guard.test.ts b/src/__tests__/restore-lock-consolidation-guard.test.ts new file mode 100644 index 000000000..b19fbbf7c --- /dev/null +++ b/src/__tests__/restore-lock-consolidation-guard.test.ts @@ -0,0 +1,75 @@ +/** + * Every day transaction of the consolidation family holds the account + * against a restore before it touches a reading (`restore-lock.ts`). + * + * The passes built on `runConsolidation` fold an account a day at a time, and + * each day's transaction takes row locks on the account's readings. A restore + * deletes all of them in one statement. Without a common first lock the two + * take the same rows in different orders and deadlock (#1031, pinned end to + * end by `tests/integration/restore-consolidation-deadlock.test.ts` for the + * step consolidation). This guard extends the rule to every pass on the same + * base, present and future: each file that calls `runConsolidation` must open + * every transaction it starts with `holdAccountAgainstRestore(tx, …)`. + * + * The matcher is whitespace-tolerant and the file and transaction counts are + * asserted to be non-zero, so a pattern that matches nothing fails instead of + * passing. + */ +import { readdirSync, readFileSync, statSync } from "node:fs"; +import { join, relative } from "node:path"; + +import { describe, expect, it } from "vitest"; + +const ROOT = join(__dirname, "..", "lib"); + +function sourceFiles(dir: string): string[] { + const out: string[] = []; + for (const name of readdirSync(dir)) { + const path = join(dir, name); + if (statSync(path).isDirectory()) { + if (name === "__tests__") continue; + out.push(...sourceFiles(path)); + } else if (/\.ts$/.test(name) && !/\.test\.ts$/.test(name)) { + out.push(path); + } + } + return out; +} + +/** The text of a file with its comments blanked, so prose cannot match. */ +function code(path: string): string { + return readFileSync(path, "utf8") + .replace(/\/\*[\s\S]*?\*\//g, " ") + .replace(/(^|[^:])\/\/.*$/gm, "$1"); +} + +const passes = sourceFiles(ROOT).filter((path) => + /\bawait\s+runConsolidation\b/.test(code(path)), +); + +describe("consolidation passes hold the account against a restore", () => { + it("finds the passes", () => { + // Step, mean, per-sample cumulative, dense intraday retention. + expect(passes.length).toBeGreaterThanOrEqual(4); + }); + + it.each(passes.map((path) => [relative(ROOT, path), path]))( + "%s opens every transaction with holdAccountAgainstRestore", + (_label, path) => { + const text = code(path); + const opened = [ + ...text.matchAll( + /\$transaction\(\s*async\s*\(\s*(\w+)\s*\)\s*=>\s*\{/g, + ), + ]; + expect(opened.length).toBeGreaterThan(0); + for (const match of opened) { + const tx = match[1]; + const after = text.slice(match.index! + match[0].length).trimStart(); + expect(after.startsWith(`await holdAccountAgainstRestore(${tx},`)).toBe( + true, + ); + } + }, + ); +}); diff --git a/src/lib/export/restore-backup.ts b/src/lib/export/restore-backup.ts index b1a2087b9..ff3bb1175 100644 --- a/src/lib/export/restore-backup.ts +++ b/src/lib/export/restore-backup.ts @@ -38,6 +38,7 @@ import { storedBackupRefusal, type StoredBackupRef, } from "@/lib/export/stored-backup"; +import { takeRestoreLock } from "@/lib/export/restore-lock"; import { insertMeasurementRows, type MeasurementInsertRow, @@ -683,12 +684,6 @@ export async function restoreBackup( report("clearing"); outcome = await prisma.$transaction( async (tx) => { - // One restore of an account at a time, whatever let a second one - // start: the job table admits one queued or running job per account, - // and this lock holds even if a worker that was thought gone is - // still writing. A second transaction waits here for the first to - // commit or roll back, then replaces the account again from its own - // file, which is what it was asked to do. // Every connection carries the request limits from `src/lib/db.ts` // (60 s per statement, 60 s idle inside a transaction). A restore is // not a request: deleting two million readings is one statement, and @@ -697,7 +692,18 @@ export async function restoreBackup( // has its own limit; SET LOCAL gives each statement the same one, and // ends with the transaction. await tx.$executeRaw`SELECT set_config('statement_timeout', ${String(transactionTimeoutMs)}, true), set_config('idle_in_transaction_session_timeout', ${String(transactionTimeoutMs)}, true)`; - await tx.$executeRaw`SELECT pg_advisory_xact_lock(hashtext(${`backup-restore:${ownerId}`}))`; + // One restore of an account at a time, whatever let a second one + // start: the job table admits one queued or running job per account, + // and this lock holds even if a worker that was thought gone is + // still writing. A second transaction waits here for the first to + // commit or roll back, then replaces the account again from its own + // file, which is what it was asked to do. + // + // The same lock orders the restore against the background passes that + // rewrite readings a day at a time (`restore-lock.ts`): they hold it + // shared for each day, so the delete below never meets a day half + // folded, which is where the two used to deadlock. + await takeRestoreLock(tx, ownerId); // Every reference this restore writes stays inside the account: the // references that leave it are counted now and again after the last // write, and any that grew rolls the whole transaction back diff --git a/src/lib/export/restore-lock.ts b/src/lib/export/restore-lock.ts new file mode 100644 index 000000000..e57c979b3 --- /dev/null +++ b/src/lib/export/restore-lock.ts @@ -0,0 +1,80 @@ +/** + * The account lock a restore holds, and the shared side of it for the passes + * that rewrite an account's readings in the background. + * + * A restore deletes every reading of the account in one statement, inside + * one long transaction. The consolidation passes (step and mean + * consolidation, the per-sample cumulative drain, the dense intraday + * retention) fold a day at a time: each day's transaction updates or mints + * the day's total and then soft-deletes or deletes the day's per-sample + * rows. Both take row locks on the same readings, each in its own order, and + * when the restore reached a day's samples before its total while a pass was + * folding that day, each held what the other needed next. Postgres broke the + * cycle by cancelling one of them (`40P01`, #1031), and on a freshly booted + * instance, where the step consolidation runs over every account with legacy + * rows, that was the restore. + * + * So the two are ordered by one advisory lock before either touches a row. + * The restore takes it exclusively as the first act of its transaction + * ({@link takeRestoreLock}); every day transaction of those passes takes it in + * shared mode as its first act ({@link holdAccountAgainstRestore}). Passes + * share the lock with each other, so nothing about how they run together + * changes. A restore that starts while a day is being folded waits for that + * one day to commit, milliseconds. A day that would start while a restore is + * running does not wait for it, which could take many minutes: it is refused + * with {@link AccountRestoreInProgressError}, and the pass leaves the account + * for its next run, which reads the restored rows rather than the ones it had + * scanned before the restore replaced them. + */ +import type { Prisma } from "@/generated/prisma/client"; + +type LockClient = Pick; + +/** The lock's key. One per account; the restore has always used this one. */ +function restoreLockKey(userId: string): string { + return `backup-restore:${userId}`; +} + +/** + * Take the account's restore lock exclusively, for the rest of the calling + * transaction. Waits for any day transaction of a background pass to commit, + * and for another restore of the same account to finish. + */ +export async function takeRestoreLock( + tx: LockClient, + userId: string, +): Promise { + // `pg_advisory_xact_lock` returns void, which the client cannot read as a + // column; selecting from it yields a plain row. + await tx.$queryRaw` + SELECT 1 AS locked + FROM pg_advisory_xact_lock(hashtext(${restoreLockKey(userId)})) + `; +} + +/** A restore of the account is running; the pass should leave it for now. */ +export class AccountRestoreInProgressError extends Error { + readonly userId: string; + constructor(userId: string) { + super("A restore of this account is running"); + this.name = "AccountRestoreInProgressError"; + this.userId = userId; + } +} + +/** + * Hold the account against a restore for the rest of the calling + * transaction: the shared side of {@link takeRestoreLock}. Call it before the + * transaction's first read or write of the account's readings. Throws + * {@link AccountRestoreInProgressError} at once when a restore holds the + * lock, rather than waiting for it. + */ +export async function holdAccountAgainstRestore( + tx: LockClient, + userId: string, +): Promise { + const [row] = await tx.$queryRaw>` + SELECT pg_try_advisory_xact_lock_shared(hashtext(${restoreLockKey(userId)})) AS held + `; + if (!row?.held) throw new AccountRestoreInProgressError(userId); +} diff --git a/src/lib/measurements/__tests__/consolidate-daily-mean.test.ts b/src/lib/measurements/__tests__/consolidate-daily-mean.test.ts index 94206a4b4..4eec27d3b 100644 --- a/src/lib/measurements/__tests__/consolidate-daily-mean.test.ts +++ b/src/lib/measurements/__tests__/consolidate-daily-mean.test.ts @@ -16,6 +16,14 @@ import { canonicalDailyTimestamp } from "../consolidation-tz"; import type { PerSampleRow } from "../drain-per-sample-cumulative"; import type { MeasurementType, PrismaClient } from "@/generated/prisma/client"; +// The restore lock is a Postgres advisory lock; the mocked client has no +// database behind it. The lock itself is pinned against a real one in +// `tests/integration/restore-consolidation-deadlock.test.ts`. +vi.mock("@/lib/export/restore-lock", async (importOriginal) => ({ + ...(await importOriginal()), + holdAccountAgainstRestore: vi.fn(async () => {}), +})); + function row( id: string, value: number, diff --git a/src/lib/measurements/__tests__/dense-intraday-pulse-summary.test.ts b/src/lib/measurements/__tests__/dense-intraday-pulse-summary.test.ts index 76f03037d..ea5f764ba 100644 --- a/src/lib/measurements/__tests__/dense-intraday-pulse-summary.test.ts +++ b/src/lib/measurements/__tests__/dense-intraday-pulse-summary.test.ts @@ -35,6 +35,14 @@ import type { PerSampleRow } from "../drain-per-sample-cumulative"; import type { MeasurementType, PrismaClient } from "@/generated/prisma/client"; import { candidateLookup, createManyVia } from "./hourly-mint-mock"; +// The restore lock is a Postgres advisory lock; the mocked client has no +// database behind it. The lock itself is pinned against a real one in +// `tests/integration/restore-consolidation-deadlock.test.ts`. +vi.mock("@/lib/export/restore-lock", async (importOriginal) => ({ + ...(await importOriginal()), + holdAccountAgainstRestore: vi.fn(async () => {}), +})); + function pulseRow(id: string, value: number, iso: string): PerSampleRow { return { id, diff --git a/src/lib/measurements/__tests__/dense-intraday-resting-mint.test.ts b/src/lib/measurements/__tests__/dense-intraday-resting-mint.test.ts index fb0579290..dd7c9218c 100644 --- a/src/lib/measurements/__tests__/dense-intraday-resting-mint.test.ts +++ b/src/lib/measurements/__tests__/dense-intraday-resting-mint.test.ts @@ -34,6 +34,14 @@ import { import type { MeasurementType, PrismaClient } from "@/generated/prisma/client"; import { candidateLookup, createManyVia } from "./hourly-mint-mock"; +// The restore lock is a Postgres advisory lock; the mocked client has no +// database behind it. The lock itself is pinned against a real one in +// `tests/integration/restore-consolidation-deadlock.test.ts`. +vi.mock("@/lib/export/restore-lock", async (importOriginal) => ({ + ...(await importOriginal()), + holdAccountAgainstRestore: vi.fn(async () => {}), +})); + /** A day well outside the retention window, so the fold definitely runs. */ const FOLD_DAY = new Date( Date.now() - (DENSE_INTRADAY_RETENTION_DAYS + 10) * 24 * 60 * 60 * 1000, diff --git a/src/lib/measurements/__tests__/dense-intraday-retention.test.ts b/src/lib/measurements/__tests__/dense-intraday-retention.test.ts index 4ea0b1250..d4e747507 100644 --- a/src/lib/measurements/__tests__/dense-intraday-retention.test.ts +++ b/src/lib/measurements/__tests__/dense-intraday-retention.test.ts @@ -27,6 +27,14 @@ import type { PerSampleRow } from "../drain-per-sample-cumulative"; import type { MeasurementType, PrismaClient } from "@/generated/prisma/client"; import { candidateLookup, createManyVia } from "./hourly-mint-mock"; +// The restore lock is a Postgres advisory lock; the mocked client has no +// database behind it. The lock itself is pinned against a real one in +// `tests/integration/restore-consolidation-deadlock.test.ts`. +vi.mock("@/lib/export/restore-lock", async (importOriginal) => ({ + ...(await importOriginal()), + holdAccountAgainstRestore: vi.fn(async () => {}), +})); + function row( id: string, value: number, diff --git a/src/lib/measurements/__tests__/dense-intraday-spo2-summary.test.ts b/src/lib/measurements/__tests__/dense-intraday-spo2-summary.test.ts index a61d7a1d8..32f5a1478 100644 --- a/src/lib/measurements/__tests__/dense-intraday-spo2-summary.test.ts +++ b/src/lib/measurements/__tests__/dense-intraday-spo2-summary.test.ts @@ -31,6 +31,14 @@ import type { PerSampleRow } from "../drain-per-sample-cumulative"; import type { MeasurementType, PrismaClient } from "@/generated/prisma/client"; import { candidateLookup, createManyVia } from "./hourly-mint-mock"; +// The restore lock is a Postgres advisory lock; the mocked client has no +// database behind it. The lock itself is pinned against a real one in +// `tests/integration/restore-consolidation-deadlock.test.ts`. +vi.mock("@/lib/export/restore-lock", async (importOriginal) => ({ + ...(await importOriginal()), + holdAccountAgainstRestore: vi.fn(async () => {}), +})); + function spo2Row(id: string, value: number, iso: string): PerSampleRow { return { id, diff --git a/src/lib/measurements/__tests__/drain-per-sample-cumulative.test.ts b/src/lib/measurements/__tests__/drain-per-sample-cumulative.test.ts index 2d8962534..ba1a18552 100644 --- a/src/lib/measurements/__tests__/drain-per-sample-cumulative.test.ts +++ b/src/lib/measurements/__tests__/drain-per-sample-cumulative.test.ts @@ -22,6 +22,14 @@ import { } from "@/lib/tz/local-day"; import type { PrismaClient } from "@/generated/prisma/client"; +// The restore lock is a Postgres advisory lock; the mocked client has no +// database behind it. The lock itself is pinned against a real one in +// `tests/integration/restore-consolidation-deadlock.test.ts`. +vi.mock("@/lib/export/restore-lock", async (importOriginal) => ({ + ...(await importOriginal()), + holdAccountAgainstRestore: vi.fn(async () => {}), +})); + describe("dayKeyForUserTz", () => { it("anchors the calendar day to the user's IANA zone", () => { // 23:45 NZST on 2026-05-16 is still 11:45 UTC on the same day, but diff --git a/src/lib/measurements/consolidate-daily-mean.ts b/src/lib/measurements/consolidate-daily-mean.ts index dd8c376e0..9b90312a4 100644 --- a/src/lib/measurements/consolidate-daily-mean.ts +++ b/src/lib/measurements/consolidate-daily-mean.ts @@ -45,6 +45,7 @@ * production standalone image strips `tsx`. Modelled on the * `step-consolidation` boot-time converging-backfill pattern. */ +import { holdAccountAgainstRestore } from "@/lib/export/restore-lock"; import type { MeasurementType, Prisma, @@ -273,6 +274,8 @@ export async function consolidateDailyMean( const unit = dayRows[0]?.unit ?? "unknown"; let removed = 0; await pc.$transaction(async (tx) => { + // First, before any reading is touched: see `restore-lock.ts`. + await holdAccountAgainstRestore(tx, userId); // The upsert below arbiters on (userId, type, source, externalId), // but the Measurement model carries a SECOND full unique — // (userId, type, measuredAt, source, sleepStage), NULLS NOT diff --git a/src/lib/measurements/consolidate-legacy-steps.ts b/src/lib/measurements/consolidate-legacy-steps.ts index 3cdcecea7..06349c1f5 100644 --- a/src/lib/measurements/consolidate-legacy-steps.ts +++ b/src/lib/measurements/consolidate-legacy-steps.ts @@ -44,6 +44,7 @@ * production standalone image strips `tsx`. Modelled on the * `rollup-full-backfill` boot-time converging-backfill pattern. */ +import { holdAccountAgainstRestore } from "@/lib/export/restore-lock"; import type { MeasurementSource, PrismaClient, @@ -318,6 +319,8 @@ export async function consolidateLegacySteps( try { let removed = 0; await pc.$transaction(async (tx) => { + // First, before any reading is touched: see `restore-lock.ts`. + await holdAccountAgainstRestore(tx, userId); if (shouldMint) { // Mint the canonical daily-total row. The unique index // (userId, type, source, externalId) makes the upsert diff --git a/src/lib/measurements/consolidation-base.ts b/src/lib/measurements/consolidation-base.ts index 8c7a447f3..592f787b8 100644 --- a/src/lib/measurements/consolidation-base.ts +++ b/src/lib/measurements/consolidation-base.ts @@ -26,6 +26,8 @@ * is no value-level cycle between the two — a cycle would otherwise trip * a module-init TDZ in the production bundle. */ +import { annotate } from "@/lib/logging/context"; +import { AccountRestoreInProgressError } from "@/lib/export/restore-lock"; import type { MeasurementType, PrismaClient, @@ -425,6 +427,8 @@ export async function runConsolidation( dryRun: boolean; daysFailed: number; stoppedEarly: boolean; + /** Accounts left for the next run because a restore of them was running. */ + usersDeferredForRestore: number; }> { const { prismaClient, options } = params; const dryRun = options.dryRun ?? false; @@ -434,6 +438,7 @@ export async function runConsolidation( let daysFailed = 0; let daysWalked = 0; let stoppedEarly = false; + let usersDeferredForRestore = 0; const users = await loadConsolidationUsers(prismaClient, options.userId); @@ -533,6 +538,20 @@ export async function runConsolidation( outcome, }); } catch (err) { + // A restore of this account started while the pass was walking + // it. The rows scanned so far are about to be replaced, so the + // account is left for the next run, which reads the restored ones + // (`restore-lock.ts`). Not a failed day, and no reason to stop the + // other accounts. + if (err instanceof AccountRestoreInProgressError) { + usersDeferredForRestore += 1; + annotate({ + meta: { + consolidation_deferred_for_restore: usersDeferredForRestore, + }, + }); + continue walk; + } // No boundary supplied → preserve the historical abort-the-run // behaviour for the drains that classify errors in `writeDay`. if (!params.onBucketError) throw err; @@ -561,5 +580,11 @@ export async function runConsolidation( params.onUserComplete?.({ userId: user.id, tz, dryRun }); } - return { usersScanned: users.length, dryRun, daysFailed, stoppedEarly }; + return { + usersScanned: users.length, + dryRun, + daysFailed, + stoppedEarly, + usersDeferredForRestore, + }; } diff --git a/src/lib/measurements/dense-intraday-retention.ts b/src/lib/measurements/dense-intraday-retention.ts index 869010135..032b173b8 100644 --- a/src/lib/measurements/dense-intraday-retention.ts +++ b/src/lib/measurements/dense-intraday-retention.ts @@ -47,6 +47,7 @@ * The one-shot history rebuild that converts ALREADY-folded daily rows to * the hourly grain lives in `dense-intraday-hourly-rebuild.ts`. */ +import { holdAccountAgainstRestore } from "@/lib/export/restore-lock"; import type { MeasurementType, Prisma, @@ -617,6 +618,8 @@ export async function runDenseIntradayRetention( let removed = 0; let retiredDaily = false; await pc.$transaction(async (tx) => { + // First, before any reading is touched: see `restore-lock.ts`. + await holdAccountAgainstRestore(tx, userId); const canonicalRowIds = await adoptOrMintHourlyRows(tx, { userId, type, diff --git a/src/lib/measurements/drain-per-sample-cumulative.ts b/src/lib/measurements/drain-per-sample-cumulative.ts index e22fb5f5a..882d7bf0f 100644 --- a/src/lib/measurements/drain-per-sample-cumulative.ts +++ b/src/lib/measurements/drain-per-sample-cumulative.ts @@ -34,6 +34,7 @@ * touching the DB so the operator can inspect what would change before * committing. */ +import { holdAccountAgainstRestore } from "@/lib/export/restore-lock"; import { Prisma } from "@/generated/prisma/client"; import type { MeasurementAggregationProvenance, @@ -291,6 +292,8 @@ export async function drainPerSampleCumulative( const foldOnce = async (): Promise => { let removed = 0; await pc.$transaction(async (tx) => { + // First, before any reading is touched: see `restore-lock.ts`. + await holdAccountAgainstRestore(tx, userId); // Single-flight the per-(user, type, day) fold. Two overlapping // drains — a rolling-deploy nightly overlap, or an operator's manual // `POST /api/admin/drain-per-sample-cumulative` racing the cron — diff --git a/tests/integration/restore-consolidation-deadlock.test.ts b/tests/integration/restore-consolidation-deadlock.test.ts new file mode 100644 index 000000000..c07e422de --- /dev/null +++ b/tests/integration/restore-consolidation-deadlock.test.ts @@ -0,0 +1,277 @@ +/** + * #1031 — a restore and a consolidation pass over the same account must not + * deadlock. + * + * Seen on a freshly booted instance: the restore's clearing step failed with + * `40P01` while the boot-time step consolidation was folding the same + * account. Each takes row locks on the account's readings in its own order. + * A consolidation day first updates (or mints) the day's `stats:` total, then + * soft-deletes the day's per-sample rows; the restore deletes every reading + * of the account in one statement, in the order the table gives them. When + * the restore reaches the day's samples before the total, each side holds + * what the other one needs next. + * + * This file builds exactly that order, deterministically. The day's samples + * are written before its (soft-deleted) total, so the restore's delete meets + * them first, and a row trigger holds the consolidation for three seconds + * right after it has locked the total. The restore is started while it waits. + * + * What must hold: both finish, neither is chosen as a deadlock victim, and + * the account ends up as the backup says. Every writer of the consolidation + * family now takes the restore's account lock in shared mode at the start of + * its day transaction, and the restore takes it exclusively before its first + * delete, so the two are ordered by that lock before either touches a row. A + * consolidation that finds a restore running leaves the account for its next + * run instead of waiting on it. + * + * Mutation check: remove the `holdAccountAgainstRestore` call from the step + * consolidation's day transaction and the first test fails with the + * consolidation rejected on `40P01`. + */ +import { + afterAll, + afterEach, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; + +import { getPrismaClient, truncateAllTables } from "./setup"; +import { acceptingBoss } from "./restore-job-driver"; + +vi.mock("@/lib/db-compat", () => ({ + ensureDbCompatibility: vi.fn().mockResolvedValue(undefined), +})); + +const OWNER = "restore-deadlock-owner"; +const DAY = "2026-01-10"; +const TOTAL_ID = "deadlock-total"; +const TOTAL_EXTERNAL_ID = `stats:HKQuantityTypeIdentifierStepCount:${DAY}`; + +async function dropTrigger() { + const prisma = getPrismaClient(); + await prisma.$executeRawUnsafe( + `DROP TRIGGER IF EXISTS restore_deadlock_hold ON measurements`, + ); + await prisma.$executeRawUnsafe( + `DROP FUNCTION IF EXISTS restore_deadlock_hold()`, + ); +} + +/** Hold the transaction that updates the day's total for three seconds. */ +async function holdTheTotalsUpdate() { + const prisma = getPrismaClient(); + await prisma.$executeRawUnsafe( + `CREATE FUNCTION restore_deadlock_hold() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN PERFORM pg_sleep(3); RETURN NEW; END $$`, + ); + await prisma.$executeRawUnsafe( + `CREATE TRIGGER restore_deadlock_hold BEFORE UPDATE ON measurements FOR EACH ROW WHEN (OLD.id = '${TOTAL_ID}') EXECUTE FUNCTION restore_deadlock_hold()`, + ); +} + +async function seedAccountWithBackup(): Promise { + const prisma = getPrismaClient(); + await prisma.user.create({ + data: { id: OWNER, username: OWNER, role: "ADMIN", timezone: "UTC" }, + }); + // The day's per-sample rows first, then its total, so a delete of the + // account meets the samples first. All samples before noon, so an index + // walk in time order meets them first too. + for (const [i, hour] of [6, 8, 10].entries()) { + await prisma.measurement.create({ + data: { + id: `deadlock-sample-${i}`, + userId: OWNER, + type: "ACTIVITY_STEPS", + value: 1000 + i, + unit: "steps", + source: "APPLE_HEALTH", + measuredAt: new Date( + `${DAY}T${String(hour).padStart(2, "0")}:00:00.000Z`, + ), + externalId: `hk-sample-${i}`, + }, + }); + } + // A total a previous pass minted and something since removed: the next + // pass mints it again through the upsert's update branch, which locks it. + await prisma.measurement.create({ + data: { + id: TOTAL_ID, + userId: OWNER, + type: "ACTIVITY_STEPS", + value: 1, + unit: "steps", + source: "MANUAL", + measuredAt: new Date(`${DAY}T12:00:00.000Z`), + externalId: TOTAL_EXTERNAL_ID, + deletedAt: new Date(), + }, + }); + const { storeBackupBlob } = await import("@/lib/export/store-backup-blob"); + const { streamFullBackupJson } = + await import("@/lib/export/full-backup-stream"); + const { id } = await storeBackupBlob( + prisma, + { userId: OWNER, type: "WEEKLY_AUTO" }, + (write) => + streamFullBackupJson(prisma, OWNER, write, { + purpose: "disaster-recovery", + }), + ); + return id; +} + +async function startRestore(backupId: string) { + const prisma = getPrismaClient(); + const jobs = await import("@/lib/jobs/backup-restore"); + const { storedBackupIdentity, STORED_BACKUP_SELECT } = + await import("@/lib/export/stored-backup"); + const backup = await prisma.dataBackup.findUniqueOrThrow({ + where: { id: backupId }, + select: STORED_BACKUP_SELECT, + }); + const admitted = await jobs.admitBackupRestore({ + userId: OWNER, + actorUserId: OWNER, + backupId, + backupDigest: jobs.backupDigest(storedBackupIdentity(backup)), + restoreInstanceSettings: false, + boss: acceptingBoss, + }); + if (!admitted.admitted) throw new Error("restore not admitted"); + return { + jobId: admitted.jobId, + done: jobs.runBackupRestoreJob(admitted.jobId), + }; +} + +/** Wait until a backend sits in the trigger's sleep. */ +async function untilHeldInTrigger() { + const prisma = getPrismaClient(); + for (let i = 0; i < 100; i++) { + const rows = await prisma.$queryRaw>` + SELECT count(*) AS n FROM pg_stat_activity WHERE wait_event = 'PgSleep' + `; + if (Number(rows[0]?.n ?? 0) > 0) return; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + throw new Error("the consolidation never reached the day's total"); +} + +beforeEach(async () => { + await dropTrigger(); + await truncateAllTables(getPrismaClient()); +}); + +afterEach(async () => { + await dropTrigger(); +}); + +afterAll(async () => { + await truncateAllTables(getPrismaClient()); +}); + +describe("restore against a running consolidation (#1031)", () => { + it("does not deadlock when the restore starts while a day is being folded", async () => { + const prisma = getPrismaClient(); + const backupId = await seedAccountWithBackup(); + await holdTheTotalsUpdate(); + + const { runStepConsolidationForUser } = + await import("@/lib/jobs/step-consolidation"); + const consolidation = runStepConsolidationForUser(OWNER).then( + (value) => ({ ok: true as const, value }), + (error: unknown) => ({ ok: false as const, error }), + ); + await untilHeldInTrigger(); + const restore = await startRestore(backupId); + + const [folded] = await Promise.all([consolidation, restore.done]); + const { readBackupRestoreJob } = await import("@/lib/jobs/backup-restore"); + const job = await readBackupRestoreJob(restore.jobId); + + expect( + folded.ok, + `the consolidation failed: ${folded.ok ? "" : String((folded.error as Error)?.message ?? folded.error)}`, + ).toBe(true); + expect(job?.status, JSON.stringify(job?.failure)).toBe("succeeded"); + // The day was folded before the restore went ahead, and the restore then + // put the account back exactly as the backup holds it. + if (folded.ok) expect(folded.value.daysConsolidated).toBe(1); + const rows = await prisma.measurement.findMany({ + where: { userId: OWNER }, + orderBy: { id: "asc" }, + select: { id: true, deletedAt: true }, + }); + expect(rows.map((row) => [row.id, row.deletedAt === null])).toEqual([ + ["deadlock-sample-0", true], + ["deadlock-sample-1", true], + ["deadlock-sample-2", true], + [TOTAL_ID, false], + ]); + }, 60_000); + + it("leaves an account alone while its restore holds the lock", async () => { + const prisma = getPrismaClient(); + await seedAccountWithBackup(); + const { holdAccountAgainstRestore, takeRestoreLock } = + await import("@/lib/export/restore-lock"); + + // A restore's transaction, standing in for the real one: the exclusive + // account lock is all a consolidation can see of it. + let release!: () => void; + const released = new Promise((resolve) => { + release = resolve; + }); + let locked!: () => void; + const isLocked = new Promise((resolve) => { + locked = resolve; + }); + const restoreTx = prisma.$transaction( + async (tx) => { + await takeRestoreLock(tx, OWNER); + locked(); + await released; + }, + { timeout: 30_000 }, + ); + await isLocked; + + try { + // Refused at once, rather than waiting on the restore. + const started = Date.now(); + await expect( + prisma.$transaction((tx) => holdAccountAgainstRestore(tx, OWNER)), + ).rejects.toMatchObject({ name: "AccountRestoreInProgressError" }); + expect(Date.now() - started).toBeLessThan(2_000); + + // The pass leaves the account's rows as they are and says so. + const { consolidateLegacySteps } = + await import("@/lib/measurements/consolidate-legacy-steps"); + const summary = await consolidateLegacySteps(prisma, { + userId: OWNER, + log: () => {}, + }); + expect(summary.totals.daysConsolidated).toBe(0); + expect( + await prisma.measurement.count({ + where: { userId: OWNER, deletedAt: null }, + }), + ).toBe(3); + } finally { + release(); + await restoreTx; + } + + // Two passes of the family can still share the account. + await prisma.$transaction(async (tx) => { + await holdAccountAgainstRestore(tx, OWNER); + await prisma.$transaction((other) => + holdAccountAgainstRestore(other, OWNER), + ); + }); + }, 60_000); +});