fix(cache): order a set() of a key behind its in-flight write on Node (LAB-12299)

packages/cachekit/README.md

@@ -434,6 +434,14 @@ pending write, which then never settles. After the cap the op runs anyway, so
 there the ordering holds for 30 s: a write still in flight at the cap can
 land after the op.
 
+With or without write-behind, on Node a `set()` or `delete()` of a key whose
+`set()` or `delete()` is still in flight waits for it to settle before
+sending, so concurrent `set()` calls of a key land in call order (on an
+encrypted cache, the order they finish encrypting) and L1 ends on the value L2
+holds. On Workers, where workerd can cancel the op in flight, concurrent
+`set()` calls of a key with no pending write stay unordered: they can land in
+either order, and L1 can serve the older value for up to its TTL.
+
 **The write stays on the caller's path:**
 
 - while `stampede.distributedLock` holds the lease: a write deferred past the
@@ -444,8 +452,9 @@ land after the op.
   `cache.withExecutionContext(ctx)`, which registers the write with
   `ctx.waitUntil`. workerd cancels unregistered work at response return;
 - while the key already has a write or delete pending or in flight in this
-  process: the write runs on the caller's path (behind a pending one,
-  alongside one in flight) and fills L1 after its L2 write, as with the
+  process: the write runs on the caller's path (behind a pending one, and on
+  Node behind one in flight too; on Workers alongside one in flight) and
+  fills L1 after its L2 write, as with the
   option off, so a write-behind never fills L1 ahead of an earlier write. On Workers
   a background write queued there would also spend workerd's 30 s
   `waitUntil` budget waiting and be cancelled as it sends; for that reason an

packages/cachekit/src/cache-core.ts

@@ -116,6 +116,13 @@ type DeferredWrite =
  */
 type L1Mode = 'fill' | 'skip' | 'refresh';
 
+/** A key's idle-key ops in flight (see keysInFlight). */
+interface InFlight {
+  count: number;
+  /** The promise the latest of them handed its caller (see inFlightTail). */
+  last: Promise<unknown>;
+}
+
 /**
  * Sentinel for "the lock path did not resolve the miss — compute without
  * it". Distinct from null: the wrapped function may legitimately resolve
@@ -453,13 +460,14 @@ export class CacheImpl implements SecureCache {
   /**
    * Per key, the L2 writes and deletes running on an idle key (no tail), so
    * not in keyTails. A write-behind of such a key is refused and runs on the
-   * caller's path, concurrently, as with the option off: its early L1 fill
-   * could otherwise land before the running op's later one, leaving L1
-   * older than L2. Kept out of keyTails so that an op workerd cancels never
+   * caller's path: its early L1 fill could otherwise land before the running
+   * op's later one, leaving L1 older than L2. On Node a later op of the key
+   * queues behind the running one (see inKeyOrder); on Workers it runs
+   * concurrently. Kept out of keyTails so that an op workerd cancels never
    * leaves a tail that never settles; a count it leaks only keeps that key's
    * writes on the caller's path.
    */
-  private readonly keysInFlight = new Map<string, number>();
+  private readonly keysInFlight = new Map<string, InFlight>();
 
   constructor(options: CacheOptions, runtime: CacheRuntime) {
     // Every config check that can throw ConfigurationError runs before
@@ -1508,7 +1516,10 @@ export class CacheImpl implements SecureCache {
     // returned L1Write. A write degradation absorbed got no response: its copy
     // keeps the TTL, counted from the L1 write (API-36).
     let until: number | undefined;
-    if (this.keyTails.has(key)) {
+    // On Node a key with an op in flight queues too: two concurrent PUTs can
+    // be answered out of order, and the older value's late fill would leave
+    // L1 older than L2.
+    if (this.keyTails.has(key) || (!this.swrRequiresWaitUntil && this.keysInFlight.has(key))) {
       await this.inKeyOrder(key, async () => {
         until = await this.run('set', key, undefined, writeL2);
         if (l1Mode === 'fill') this.fillL1(key, l1Write, ttl, namespace, until);
@@ -1517,9 +1528,10 @@ export class CacheImpl implements SecureCache {
       // An idle key: what inKeyOrder does for it, inline. Every plain set()
       // takes this path, and inKeyOrder's closure and promise chain were most
       // of what per-key ordering added to it, so the two lines above repeat.
-      this.enterInFlight(key);
+      const written = this.run('set', key, undefined, writeL2);
+      this.enterInFlight(key, written);
       try {
-        until = await this.run('set', key, undefined, writeL2);
+        until = await written;
         if (l1Mode === 'fill') this.fillL1(key, l1Write, ttl, namespace, until);
       } finally {
         this.leaveInFlight(key);
@@ -1680,8 +1692,10 @@ export class CacheImpl implements SecureCache {
    * Workers when the key has an op queued. It is never put on a caller's
    * path: the revalidation runs in the background already.
    *
-   * On Node it queues behind the key's pending ops (it fills no L1, so
-   * queuing cannot leave L1 older than L2), counts against the bound, and
+   * On Node it queues behind the key's pending ops and the op in flight on
+   * it (it fills no L1, so queuing cannot leave L1 older than L2; skipping
+   * the op in flight would let a later op queued behind it overtake that
+   * op), counts against the bound, and
    * becomes the key's tail: a later set(), delete() or invalidate('params')
    * orders after it, and close() drains it.
    *
@@ -1714,46 +1728,60 @@ export class CacheImpl implements SecureCache {
     // On Workers no tail exists here (checked above), so inKeyOrder counts it in flight.
     if (this.swrRequiresWaitUntil) return this.inKeyOrder(key, send);
     return this.startPendingWrite(key, bytes, () => {
-      const tail = this.keyTails.get(key);
+      const tail = this.keyTails.get(key) ?? this.inFlightTail(key);
       return tail ? tail.then(send) : send();
     });
   }
 
   /**
    * Run an L2 write or delete of `key` after everything already queued for
-   * the key, and make it the key's new tail. An op on an idle key runs at
-   * once, as before write-behind existed, counted in keysInFlight while it
-   * runs.
+   * the key, and make it the key's new tail. On Node an op in flight on the
+   * key counts as its tail. An op on an idle key runs at once, as before
+   * write-behind existed, counted in keysInFlight while it runs.
    */
   private inKeyOrder<T>(key: string, op: () => Promise<T>): Promise<T> {
-    const tail = this.keyTails.get(key);
+    const tail = this.keyTails.get(key) ?? this.inFlightTail(key);
     if (tail === undefined) return this.countInFlight(key, op);
     const result = this.afterTail(tail).then(op);
     this.setKeyTail(key, result);
     return result;
   }
 
+  /**
+   * On Node, the settling of the latest op in flight on `key`, or undefined.
+   * Its caller awaited that promise before any waiter could, so the caller's
+   * L1 fill or eviction runs before a waiter's op starts. On Workers always
+   * undefined: workerd can cancel an idle-key op, and a wait for it would
+   * stall the key's next op.
+   */
+  private inFlightTail(key: string): Promise<void> | undefined {
+    if (this.swrRequiresWaitUntil) return undefined;
+    return this.keysInFlight.get(key)?.last.then(
+      () => {},
+      () => {}
+    );
+  }
+
   private countInFlight<T>(key: string, op: () => Promise<T>): Promise<T> {
-    this.enterInFlight(key);
-    const release = (): void => this.leaveInFlight(key);
-    let running: Promise<T>;
-    try {
-      running = op();
-    } catch (error) {
-      release();
-      throw error;
-    }
-    return running.finally(release);
+    const result = op().finally(() => this.leaveInFlight(key));
+    this.enterInFlight(key, result);
+    return result;
   }
 
-  private enterInFlight(key: string): void {
-    this.keysInFlight.set(key, (this.keysInFlight.get(key) ?? 0) + 1);
+  /** Count `op` in flight on `key`; `op` is the promise its caller awaits. */
+  private enterInFlight(key: string, op: Promise<unknown>): void {
+    const inFlight = this.keysInFlight.get(key);
+    if (inFlight === undefined) {
+      this.keysInFlight.set(key, { count: 1, last: op });
+    } else {
+      inFlight.count++;
+      inFlight.last = op;
+    }
   }
 
   private leaveInFlight(key: string): void {
-    const left = (this.keysInFlight.get(key) ?? 1) - 1;
-    if (left > 0) this.keysInFlight.set(key, left);
-    else this.keysInFlight.delete(key);
+    const inFlight = this.keysInFlight.get(key);
+    if (inFlight === undefined || --inFlight.count === 0) this.keysInFlight.delete(key);
   }
 
   /** A key's tail, capped at KEY_WAIT_CAP_MS where it may never settle. */

packages/cachekit/src/cache.stale-grace.test.ts

@@ -849,6 +849,65 @@ describe('stale-grace revalidation write through the per-key write queue', () =>
     expect(backend.log).toEqual(['write:done', 'write:start', 'write:done', 'release']);
   });
 
+  it('on Node, waits for a set() in flight on the key, so a later set() cannot overtake that set()', async () => {
+    /** Every write waits for its own release, in any order. */
+    class EachHeldBackend extends StaleGraceBackend {
+      held: Array<() => void> = [];
+      hold = false;
+
+      private gate(): Promise<void> {
+        if (!this.hold) return Promise.resolve();
+        return new Promise<void>((resolve) => this.held.push(resolve));
+      }
+
+      async set(key: string, value: Uint8Array, ttl?: number): Promise<void> {
+        await this.gate();
+        await super.set(key, value, ttl);
+      }
+
+      async setWithStaleTtl(key: string, value: Uint8Array, ttl: number, staleTtl: number) {
+        await this.gate();
+        await super.setWithStaleTtl(key, value, ttl, staleTtl);
+      }
+    }
+    const backend = new EachHeldBackend();
+    await make(backend).wrap(async () => 'v1', WRAP)();
+    backend.expire();
+    const [key] = backend.store.keys();
+    const cache = make(backend);
+    let computed!: () => void;
+    const computing = new Promise<void>((resolve) => {
+      computed = resolve;
+    });
+    const wrapped = cache.wrap(async () => {
+      await computing;
+      return 'v2';
+    }, WRAP);
+    expect(await wrapped()).toBe('v1'); // served stale; the revalidation is computing
+
+    backend.hold = true;
+    const first = cache.set(key, 'first', { ttl: 60 }); // in flight on an idle key
+    await sleep(10);
+    computed();
+    await sleep(10);
+    expect(backend.held).toHaveLength(1); // the revalidation's write waits for the set
+    const second = cache.set(key, 'second', { ttl: 60 });
+    await sleep(10);
+    expect(backend.held).toHaveLength(1); // and the second set waits for both
+
+    backend.held[0](); // the first set lands
+    await first;
+    await vi.waitFor(() => expect(backend.held).toHaveLength(2));
+    backend.held[1](); // the revalidation lands
+    await vi.waitFor(() => expect(backend.held).toHaveLength(3));
+    backend.held[2](); // the second set lands last
+    await second;
+
+    expect(await cache.get(key)).toBe('second'); // L1
+    await cache.invalidate('global');
+    expect(await cache.get(key)).toBe('second'); // L2
+  });
+
   it('on Workers, a revalidation write is never a key tail, so a later op does not wait on it', async () => {
     const backend = new HeldWritesBackend();
     await make(backend).wrap(async () => 'v1', WRAP)();

packages/cachekit/src/cache.test.ts

@@ -1344,17 +1344,15 @@ describe('Cache Integration', () => {
       await swrCache.close();
     });
 
-    it('an explicit set() during an in-flight refresh wins in L1 (refresh L1 write discarded)', async () => {
-      // Gate the refresh's L2 persist (2nd backend set: 1st = cold miss,
-      // 2nd = refresh, 3rd = the explicit set) so the explicit write can
-      // interleave between the refresh's L2 write and its L1 completion.
+    it('an explicit set() during a refresh’s compute wins in L1 (refresh L1 write discarded)', async () => {
+      // Gate the refresh's compute (call 2) so the explicit write lands in L1
+      // between the refresh's start and its completeRefresh.
       let release!: () => void;
       const gate = new Promise<void>((resolve) => {
         release = resolve;
       });
-      const gatedBackend = new GatedBackend(2, gate);
       const swrCache = createCache({
-        backend: gatedBackend,
+        backend: new InMemoryBackend(),
         defaultTtl: 60,
         l1: { swrEnabled: true, swrThresholdRatio: 2 },
       });
@@ -1363,17 +1361,17 @@ describe('Cache Integration', () => {
       const fn = swrCache.wrap(
         async () => {
           calls++;
+          if (calls === 2) await gate;
           return { gen: calls };
         },
         { namespace: 'swr:race', ttl: 60 }
       );
       const cacheKey = generateKey('swr:race', []);
 
-      expect(await fn()).toEqual({ gen: 1 }); // cold miss (set #1)
-      expect(await fn()).toEqual({ gen: 1 }); // stale hit → refresh blocks in set #2
-
+      expect(await fn()).toEqual({ gen: 1 }); // cold miss
+      expect(await fn()).toEqual({ gen: 1 }); // stale hit → refresh parks in its compute
       await vi.waitFor(() => {
-        expect(gatedBackend.setCount).toBe(2); // refresh is parked in its L2 write
+        expect(calls).toBe(2);
       });
 
       // Explicit write lands while the refresh is in flight — bumps the L1
@@ -1391,6 +1389,50 @@ describe('Cache Integration', () => {
       await swrCache.close();
     });
 
+    it('on Node, an explicit set() during a refresh’s L2 write lands after it, in L1 and L2', async () => {
+      // Gate the refresh's L2 write (2nd backend set: 1st = cold miss,
+      // 2nd = refresh, 3rd = the explicit set).
+      let release!: () => void;
+      const gate = new Promise<void>((resolve) => {
+        release = resolve;
+      });
+      const gatedBackend = new GatedBackend(2, gate);
+      const swrCache = createCache({
+        backend: gatedBackend,
+        defaultTtl: 60,
+        l1: { swrEnabled: true, swrThresholdRatio: 2 },
+      });
+
+      let calls = 0;
+      const fn = swrCache.wrap(
+        async () => {
+          calls++;
+          return { gen: calls };
+        },
+        { namespace: 'swr:race-put', ttl: 60 }
+      );
+      const cacheKey = generateKey('swr:race-put', []);
+
+      expect(await fn()).toEqual({ gen: 1 }); // cold miss (set #1)
+      expect(await fn()).toEqual({ gen: 1 }); // stale hit → refresh blocks in set #2
+      await vi.waitFor(() => {
+        expect(gatedBackend.setCount).toBe(2); // refresh is parked in its L2 write
+      });
+
+      const explicit = swrCache.set(cacheKey, { gen: 999 }, { ttl: 60, namespace: 'swr:race-put' });
+      await new Promise((resolve) => setTimeout(resolve, 20));
+      expect(gatedBackend.setCount).toBe(2); // queued behind the refresh's PUT
+
+      release();
+      await explicit;
+      await new Promise((resolve) => setTimeout(resolve, 20));
+      expect(await swrCache.get(cacheKey)).toEqual({ gen: 999 }); // L1
+      const l2View = createCache({ backend: gatedBackend, defaultTtl: 60, l1: { enabled: false } });
+      expect(await l2View.get(cacheKey)).toEqual({ gen: 999 }); // L2
+      await l2View.close();
+      await swrCache.close();
+    });
+
     it('a refresh whose value fails to encode leaves L1 alone on a plaintext cache', async () => {
       // A direct write and a cold miss store nothing for a value the encoder
       // rejects; the refresh must not be the one path that puts it in L1.