fix(review): make preview-poll-budget increment safe under concurrent triggers

src/review/visual/preview-poll-budget.ts

@@ -31,34 +31,51 @@ 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<string> {
   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<BudgetMarker | null> {
-  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<BudgetMarker>;
+  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<BudgetRead> {
+  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<BudgetMarker>;
-    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 };
   }
 }
 
 /** How many preview-poll attempts have already been recorded for `headSha` -- 0 when no marker exists,
  *  storage is unavailable, or the existing marker has expired. Consulted by buildCapture BEFORE treating a
  *  "still building" preview state as worth another attempt. */
 export async function previewPollAttemptCount(env: Env, headSha: string): Promise<number> {
-  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<void> {
   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
   }

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<string, string>();
+// 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> | void } = {},
+): R2Bucket {
+  const store = new Map<string, { value: string; etag: string }>();
+  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<void>((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
+  });
 });