diff --git a/src/lib/documents/__tests__/pdf-isolation.test.ts b/src/lib/documents/__tests__/pdf-isolation.test.ts index 4bf197d77..84fed27f2 100644 --- a/src/lib/documents/__tests__/pdf-isolation.test.ts +++ b/src/lib/documents/__tests__/pdf-isolation.test.ts @@ -12,6 +12,7 @@ vi.mock("@/lib/logging/context", () => ({ annotate })); import { gif } from "../../__tests__/gif-fixture"; import { MAX_QUEUED_PER_OWNER, + READER_WAIT_MS, runOnce, runPdfChild, withPdfWork, @@ -301,3 +302,76 @@ describe("the reader's queue", () => { } }); }); + +describe("the caller's clock while it waits for the reader", () => { + const request = { op: "text", data: Buffer.from("x"), maxPages: 1 } as const; + + it("moves by the wait plus the read's allowance when the slot is granted, then gives back what the read did not use", async () => { + let open: () => void = () => {}; + const gate = new Promise((resolve) => (open = resolve)); + const holder = vi.fn(async (): Promise> => { + await gate; + return { ok: true, value: "first" }; + }); + const quick = vi.fn(async (): Promise> => ({ + ok: true, + value: "second", + })); + const moves: number[] = []; + const first = withPdfWork({ owner: "a", lane: "interactive" }, () => + runPdfChild(request, { timeoutMs: 1_000, run: holder }), + ); + const second = withPdfWork( + { + owner: "b", + lane: "interactive", + extendDeadline: async (ms) => { + moves.push(ms); + }, + }, + () => runPdfChild(request, { timeoutMs: 1_000, run: quick }), + ); + await new Promise((resolve) => setTimeout(resolve, 150)); + open(); + await expect(first).resolves.toMatchObject({ ok: true }); + await expect(second).resolves.toMatchObject({ ok: true, value: "second" }); + + // Granted after about 150 ms of waiting, with two attempts' allowance. + expect(moves).toHaveLength(2); + expect(moves[0]!).toBeGreaterThanOrEqual(2_000 + 140); + expect(moves[0]!).toBeLessThan(2_000 + 1_000); + // The quick read gives back nearly all of its allowance. + expect(moves[1]!).toBeLessThan(0); + expect(moves[1]!).toBeGreaterThan(-2_000 - 1); + // Net: the wait and the read's own time, nothing more. + expect(moves[0]! + moves[1]!).toBeGreaterThanOrEqual(140); + expect(moves[0]! + moves[1]!).toBeLessThan(1_000); + }); + + it("reads on when moving the clock fails", async () => { + const run = vi.fn(async (): Promise> => ({ + ok: true, + value: "text", + })); + await expect( + withPdfWork( + { + owner: "c", + lane: "interactive", + extendDeadline: async () => { + throw new Error("database gone"); + }, + }, + () => runPdfChild(request, { timeoutMs: 1_000, run }), + ), + ).resolves.toEqual({ ok: true, value: "text" }); + }); + + it("waits no longer for the reader than a run's margin can absorb before the slot moves its deadline", async () => { + // The deadline moves when the slot is granted; until then the run's + // armed margin is all that stands between a waiting read and the + // reaper. + const { AI_RUN_MARGIN_MS } = await import("../ai-runs/types"); + expect(READER_WAIT_MS.interactive).toBeLessThanOrEqual(AI_RUN_MARGIN_MS); + }); +}); diff --git a/src/lib/documents/pdf-isolation.ts b/src/lib/documents/pdf-isolation.ts index c0a2ea8dc..74167c6b7 100644 --- a/src/lib/documents/pdf-isolation.ts +++ b/src/lib/documents/pdf-isolation.ts @@ -143,6 +143,16 @@ export interface PdfWork { /** The account the work is done for; the queue takes accounts in turn. */ owner: string; lane: ReaderLane; + /** + * Moves the caller's own clock by `ms`. A background AI run has a deadline + * measured in model calls, and time spent waiting for the reader or + * reading is not model time: when the slot is granted the clock moves by + * the wait plus the most the read may take, and when the read ends it moves + * back by the part of that allowance the read did not use. So the run's + * model-time budget starts once the reader is done with it, however long + * the queue was. Failures are the caller's to swallow; the read goes on. + */ + extendDeadline?: (ms: number) => Promise; } const workContext = new AsyncLocalStorage(); @@ -265,10 +275,15 @@ export async function runPdfChild( } const work = currentPdfWork(); + const queuedAt = Date.now(); const ticket = await queue.acquire(work.owner, work.lane); if (!ticket.ok) { return { ok: false, reason: "busy", retryAfterMs: ticket.retryAfterMs }; } + const grantedAt = Date.now(); + // The most this read may take: every attempt to its time limit. + const allowance = limits.timeoutMs * MAX_ATTEMPTS; + await moveCallerClock(work, grantedAt - queuedAt + allowance); try { for (let attempt = 1; ; attempt++) { const result = limits.run @@ -279,6 +294,19 @@ export async function runPdfChild( } } finally { ticket.release(); + await moveCallerClock( + work, + -Math.max(0, allowance - (Date.now() - grantedAt)), + ); + } +} + +async function moveCallerClock(work: PdfWork, ms: number): Promise { + if (!work.extendDeadline || ms === 0) return; + try { + await work.extendDeadline(ms); + } catch { + // The caller's clock is the caller's; the read is not failed for it. } } diff --git a/src/lib/jobs/__tests__/document-ai-run.test.ts b/src/lib/jobs/__tests__/document-ai-run.test.ts index 7de034a2d..bd933e31e 100644 --- a/src/lib/jobs/__tests__/document-ai-run.test.ts +++ b/src/lib/jobs/__tests__/document-ai-run.test.ts @@ -293,7 +293,7 @@ describe("runDocumentAiRun", () => { return { ok: true, data: {} }; }); await runDocumentAiRun("run1", NOW + 2_700_000, now); - expect(seen).toEqual({ owner: "u1", lane: "interactive" }); + expect(seen).toMatchObject({ owner: "u1", lane: "interactive" }); }); it("answers a busy document reader with a retryable code and hands the reservation back", async () => { diff --git a/src/lib/jobs/document-ai-run.ts b/src/lib/jobs/document-ai-run.ts index 69949a2b0..0c5bd62ab 100644 --- a/src/lib/jobs/document-ai-run.ts +++ b/src/lib/jobs/document-ai-run.ts @@ -187,6 +187,34 @@ async function armDeadline( }); } +/** + * Move a running run's deadline by `ms`, never past the job's own expiry. + * The document reader calls it (`PdfWork.extendDeadline`): the deadline is + * model time, and a wait in the reader's queue or the read itself is not, so + * a person's read is not timed out because other work held the reader. + */ +export async function moveRunDeadline( + runId: string, + ms: number, + deadlineAt: number | undefined, +): Promise { + const delta = Math.round(ms); + if (deadlineAt === undefined) { + await prisma.$executeRaw` + UPDATE document_ai_runs + SET expires_at = expires_at + (${delta}::double precision * interval '1 millisecond') + WHERE id = ${runId} AND status = 'RUNNING'`; + return; + } + await prisma.$executeRaw` + UPDATE document_ai_runs + SET expires_at = LEAST( + expires_at + (${delta}::double precision * interval '1 millisecond'), + ${new Date(deadlineAt)} + ) + WHERE id = ${runId} AND status = 'RUNNING'`; +} + async function runDocumentIndex( run: ClaimedAiRun, deadlineAt: number | undefined, @@ -619,7 +647,11 @@ export async function runDocumentAiRun( // A person asked for this read and is waiting on it: the document // reader serves it ahead of background work, in turn with other people. outcome = await withPdfWork( - { owner: run.userId, lane: "interactive" }, + { + owner: run.userId, + lane: "interactive", + extendDeadline: (ms) => moveRunDeadline(run.id, ms, deadlineAt), + }, () => RUN_DISPATCH[run.kind](run, deadlineAt, now), ); } catch (err) { diff --git a/tests/integration/document-ai-runs.test.ts b/tests/integration/document-ai-runs.test.ts index 68a3a922c..177493de2 100644 --- a/tests/integration/document-ai-runs.test.ts +++ b/tests/integration/document-ai-runs.test.ts @@ -540,6 +540,68 @@ describe("the reaper", () => { }); }); +/** + * A run's deadline is model time. The document reader moves it by the time + * a read waited for the reader and took, so a person's lab scan is not + * timed out because other work held the reader; never past the job's own + * expiry, and only while the run is running. + */ +describe("the deadline while the document reader works", () => { + async function runningRun(expiresAt: Date) { + const user = await makeUser(); + await grantConsent(user.id); + await signIn(user.id); + const queued = await postLabScanText("Glucose 95 mg/dL"); + const { runId } = queued.body.data as { runId: string }; + await getPrismaClient().documentAiRun.update({ + where: { id: runId }, + data: { status: "RUNNING", expiresAt }, + }); + return runId; + } + const expiry = async (runId: string) => + ( + await getPrismaClient().documentAiRun.findUniqueOrThrow({ + where: { id: runId }, + }) + ).expiresAt.getTime(); + + it("moves a running run's deadline by the reader's time, so the reaper leaves it alone", async () => { + const { moveRunDeadline } = await import("@/lib/jobs/document-ai-run"); + const { reapAiRuns } = await import("@/lib/documents/ai-runs/store"); + const armed = new Date(Date.now() + 5_000); + const runId = await runningRun(armed); + + // The slot is granted after a 60 s wait, with 70 s allowed for the read. + await moveRunDeadline(runId, 130_000, undefined); + expect(await expiry(runId)).toBe(armed.getTime() + 130_000); + // Without the move the reaper would have taken it a minute later. + expect( + (await reapAiRuns(new Date(armed.getTime() + 60_000))).timedOut, + ).toBe(0); + + // The read took 10 s of its 70: the unused 60 s go back. + await moveRunDeadline(runId, -60_000, undefined); + expect(await expiry(runId)).toBe(armed.getTime() + 70_000); + }); + + it("never moves it past the job's own expiry, and leaves a run that is not running alone", async () => { + const { moveRunDeadline } = await import("@/lib/jobs/document-ai-run"); + const armed = new Date(Date.now() + 5_000); + const runId = await runningRun(armed); + const jobEnds = armed.getTime() + 20_000; + await moveRunDeadline(runId, 130_000, jobEnds); + expect(await expiry(runId)).toBe(jobEnds); + + await getPrismaClient().documentAiRun.update({ + where: { id: runId }, + data: { status: "FAILED" }, + }); + await moveRunDeadline(runId, 130_000, undefined); + expect(await expiry(runId)).toBe(jobEnds); + }); +}); + describe("a summary in the background", () => { it("marks a stored summary PENDING while queued and READY when it lands", async () => { const user = await makeUser();