diff --git a/src/review/visual/preview-poll-budget.ts b/src/review/visual/preview-poll-budget.ts index 559b963110..9c8b834a37 100644 --- a/src/review/visual/preview-poll-budget.ts +++ b/src/review/visual/preview-poll-budget.ts @@ -31,26 +31,43 @@ const BUDGET_MARKER_MAX_AGE_MS = 24 * 60 * 60 * 1000; export const MAX_PREVIEW_POLL_ATTEMPTS = 5; type BudgetMarker = { count: number; firstAttemptAt: number }; +// A read of the marker plus the R2 httpEtag it was stored under (null when no object exists yet). The etag is +// what recordPreviewPollAttempt's conditional write compares-and-swaps against so two triggers racing for the +// same head SHA can't both read count=N and both write count=N+1, silently losing one increment (#7780). +type BudgetRead = { marker: BudgetMarker | null; etag: string | null }; +// How many times recordPreviewPollAttempt re-reads + retries its conditional write when another trigger wins +// the compare-and-swap first. Small: the race window is a single R2 round-trip and realistically at most a +// handful of triggers ever contend for one head SHA at once, so a couple of retries converges; exhausting them +// just degrades to the pre-#7780 best-effort "this attempt didn't count" outcome, the same safe direction the +// module already accepts for a genuine write failure. +const BUDGET_CAS_MAX_ATTEMPTS = 3; async function budgetR2Key(headSha: string): Promise { const fingerprint = await sha256Hex(`${headSha}:preview-poll-budget`); return `${BUDGET_R2_NAMESPACE}${fingerprint.slice(0, 40)}.json`; } -/** Shared read path for both public functions below. Returns null (fail-open toward "no attempts yet") on - * any read error, a malformed marker, or one older than BUDGET_MARKER_MAX_AGE_MS -- a stale marker is - * treated as absent, not as "budget still exhausted from a previous, unrelated review cycle". */ -async function readBudgetMarker(env: Env, headSha: string): Promise { - if (!env.REVIEW_AUDIT) return null; +/** Validate a raw stored payload into a BudgetMarker, or null when it's malformed or older than + * BUDGET_MARKER_MAX_AGE_MS -- a stale marker is treated as absent, not as "budget still exhausted from a + * previous, unrelated review cycle". */ +function parseBudgetMarker(text: string): BudgetMarker | null { + const marker = JSON.parse(text) as Partial; + if (typeof marker.count !== "number" || typeof marker.firstAttemptAt !== "number") return null; + if (Date.now() - marker.firstAttemptAt >= BUDGET_MARKER_MAX_AGE_MS) return null; + return { count: marker.count, firstAttemptAt: marker.firstAttemptAt }; +} + +/** Shared read path for both public functions below. Returns a fail-open read (marker null, etag null) on any + * read error or a malformed/stale marker. Also surfaces the object's httpEtag so the increment path can do a + * compare-and-swap write against exactly the version it read (#7780). */ +async function readBudgetMarker(env: Env, headSha: string): Promise { + if (!env.REVIEW_AUDIT) return { marker: null, etag: null }; try { const object = await env.REVIEW_AUDIT.get(await budgetR2Key(headSha)); - if (!object) return null; - const marker = JSON.parse(await new Response(object.body).text()) as Partial; - if (typeof marker.count !== "number" || typeof marker.firstAttemptAt !== "number") return null; - if (Date.now() - marker.firstAttemptAt >= BUDGET_MARKER_MAX_AGE_MS) return null; - return { count: marker.count, firstAttemptAt: marker.firstAttemptAt }; + if (!object) return { marker: null, etag: null }; + return { marker: parseBudgetMarker(await new Response(object.body).text()), etag: object.httpEtag }; } catch { - return null; + return { marker: null, etag: null }; } } @@ -58,7 +75,7 @@ async function readBudgetMarker(env: Env, headSha: string): Promise { - return (await readBudgetMarker(env, headSha))?.count ?? 0; + return (await readBudgetMarker(env, headSha)).marker?.count ?? 0; } /** Record one more preview-poll attempt for `headSha`, preserving the marker's original `firstAttemptAt` @@ -70,9 +87,20 @@ export async function previewPollAttemptCount(env: Env, headSha: string): Promis export async function recordPreviewPollAttempt(env: Env, headSha: string): Promise { if (!env.REVIEW_AUDIT) return; try { - const existing = await readBudgetMarker(env, headSha); - const marker: BudgetMarker = { count: (existing?.count ?? 0) + 1, firstAttemptAt: existing?.firstAttemptAt ?? Date.now() }; - await env.REVIEW_AUDIT.put(await budgetR2Key(headSha), JSON.stringify(marker), { httpMetadata: { contentType: "application/json" } }); + const key = await budgetR2Key(headSha); + for (let attempt = 0; attempt < BUDGET_CAS_MAX_ATTEMPTS; attempt += 1) { + const existing = await readBudgetMarker(env, headSha); + const marker: BudgetMarker = { count: (existing.marker?.count ?? 0) + 1, firstAttemptAt: existing.marker?.firstAttemptAt ?? Date.now() }; + // Compare-and-swap against exactly the version we just read: only overwrite the existing object if its + // etag is unchanged (etagMatches), or -- when we read no object -- only create one if none exists yet + // (etagDoesNotMatch: "*"). If another trigger wrote in between, R2 returns null instead of writing, and + // we loop to re-read its newer count and retry, so no increment is lost (#7780). + const onlyIf: R2Conditional = existing.etag !== null ? { etagMatches: existing.etag } : { etagDoesNotMatch: "*" }; + const written = await env.REVIEW_AUDIT.put(key, JSON.stringify(marker), { httpMetadata: { contentType: "application/json" }, onlyIf }); + if (written) return; + } + // Exhausted retries under sustained contention -- degrade to "this attempt didn't count", the same safe + // failure direction the module already accepts for a genuine write failure (see doc comment above). } catch { // best effort -- see doc comment above } diff --git a/test/unit/preview-poll-budget.test.ts b/test/unit/preview-poll-budget.test.ts index 646abaa878..36082b9894 100644 --- a/test/unit/preview-poll-budget.test.ts +++ b/test/unit/preview-poll-budget.test.ts @@ -4,21 +4,38 @@ import { createTestEnv } from "../helpers/d1"; const HEAD_SHA = "budget-head-sha-1234567890"; -function memoryBudgetStore(options: { failGet?: boolean; failPut?: boolean; forcedValue?: string } = {}): R2Bucket { - const store = new Map(); +// An etag-aware in-memory R2 stand-in: each key holds its value plus a monotonic etag, `get` surfaces the +// current httpEtag, and `put` honors the compare-and-swap conditions recordPreviewPollAttempt relies on +// (#7780) -- `etagMatches` writes only when the etag is unchanged, `etagDoesNotMatch: "*"` writes only when +// the object is still absent -- returning null (no write) on a precondition miss exactly as real R2 does. +// `onPut` is an optional hook fired at the START of every put, used to interleave two concurrent writers. +function memoryBudgetStore( + options: { failGet?: boolean; failPut?: boolean; forcedValue?: string; onPut?: (key: string) => Promise | void } = {}, +): R2Bucket { + const store = new Map(); + let etagSeq = 0; return { async get(key: string) { if (options.failGet) throw new Error("simulated budget-marker read failure"); // forcedValue bypasses the real per-key store entirely -- used to simulate a corrupted/malformed stored // marker without needing to know the module's own private R2-key derivation. - if (options.forcedValue !== undefined) return { body: new Response(options.forcedValue).body } as unknown as R2ObjectBody; - const value = store.get(key); - return value === undefined ? null : ({ body: new Response(value).body } as unknown as R2ObjectBody); + if (options.forcedValue !== undefined) return { body: new Response(options.forcedValue).body, httpEtag: "forced-etag" } as unknown as R2ObjectBody; + const entry = store.get(key); + return entry === undefined ? null : ({ body: new Response(entry.value).body, httpEtag: entry.etag } as unknown as R2ObjectBody); }, - async put(key: string, value: unknown) { + async put(key: string, value: unknown, putOptions?: R2PutOptions) { + if (options.onPut) await options.onPut(key); if (options.failPut) throw new Error("simulated budget-marker write failure"); - store.set(key, await new Response(value as BodyInit).text()); - return { key } as unknown as R2Object; + const onlyIf = putOptions?.onlyIf as R2Conditional | undefined; + const current = store.get(key); + // Enforce the two compare-and-swap preconditions the production code sends; a miss returns null (real R2 + // signals "not written" by returning null rather than throwing). + if (onlyIf?.etagMatches !== undefined && current?.etag !== onlyIf.etagMatches) return null; + if (onlyIf?.etagDoesNotMatch === "*" && current !== undefined) return null; + etagSeq += 1; + const etag = `etag-${etagSeq}`; + store.set(key, { value: await new Response(value as BodyInit).text(), etag }); + return { key, etag } as unknown as R2Object; }, } as unknown as R2Bucket; } @@ -118,4 +135,52 @@ describe("previewPollAttemptCount / recordPreviewPollAttempt (#6323 -- durable p const env = createTestEnv({ REVIEW_AUDIT: memoryBudgetStore({ failGet: true }) }); await expect(recordPreviewPollAttempt(env, HEAD_SHA)).resolves.toBeUndefined(); }); + + it("does NOT lose an increment when two triggers race for the same head SHA -- both are counted (#7780)", async () => { + // Interleave two concurrent recordPreviewPollAttempt calls so BOTH read the marker before EITHER writes -- + // the classic read-modify-write TOCTOU. The first put to reach the store is stalled at a barrier until the + // second writer has read (and is about to write); whichever writes second must have re-read the newer + // count via the compare-and-swap retry, so the final count reflects BOTH increments, not one. + let releaseFirstPut: () => void = () => {}; + const firstPutStalled = new Promise((resolve) => { + releaseFirstPut = resolve; + }); + let putCalls = 0; + const onPut = async () => { + putCalls += 1; + // Only the very first put blocks; every later put (the retry, and the second writer) runs immediately. + if (putCalls === 1) await firstPutStalled; + }; + const env = createTestEnv({ REVIEW_AUDIT: memoryBudgetStore({ onPut }) }); + + const first = recordPreviewPollAttempt(env, HEAD_SHA); // reads count=0, stalls at its put barrier + // Let the first writer reach (and block at) its put before the second even starts, guaranteeing both read + // count=0 against the same (absent) etag. + await new Promise((r) => setTimeout(r, 0)); + const second = recordPreviewPollAttempt(env, HEAD_SHA); // reads count=0 too, writes count=1 (wins the CAS) + await second; + releaseFirstPut(); // the first writer's stalled put now runs; its etagDoesNotMatch:"*" precondition misses + await first; // ...so it retries, re-reads count=1, and writes count=2 + + await expect(previewPollAttemptCount(env, HEAD_SHA)).resolves.toBe(2); + }); + + it("gives up after the bounded CAS retries under sustained contention, without throwing (#7780)", async () => { + // A pathological store whose conditional put NEVER succeeds (every etagDoesNotMatch precondition is treated + // as a miss): recordPreviewPollAttempt must exhaust its bounded retries and degrade to "this attempt didn't + // count" -- the same safe failure direction as a genuine write failure -- rather than throw or loop forever. + let putAttempts = 0; + const store = { + async get() { + return null; // always "no marker yet" -> the write path always uses etagDoesNotMatch:"*" + }, + async put() { + putAttempts += 1; + return null; // precondition perpetually "misses" -> forces the retry loop to run to exhaustion + }, + } as unknown as R2Bucket; + const env = createTestEnv({ REVIEW_AUDIT: store }); + await expect(recordPreviewPollAttempt(env, HEAD_SHA)).resolves.toBeUndefined(); + expect(putAttempts).toBe(3); // BUDGET_CAS_MAX_ATTEMPTS + }); });