AsyncLocalStorage: don't retain the outer store from exit() / nested run()

src/js/node/async_hooks.ts

@@ -20,9 +20,11 @@
 //
 // AsyncContextData is the innermost Frame of a persistent linked list managed in
 // here: each Frame binds one AsyncLocalStorage to a value and points at the frame
-// it was pushed onto, so run() allocates one three-field object and never copies,
-// getStore() walks the (short) chain, and a captured context is a single
-// reference that shares its tail with every other capture.
+// it was pushed onto, so run() allocates one small object on entry and, when
+// nested, one on exit (re-linking only the frames of other storages above an
+// existing binding of the same storage), getStore() walks a chain no longer
+// than the number of storages, and a captured context is a single reference
+// that shares its tail with every other capture.
 //
 const setAsyncHooksEnabled = $newCppFunction("NodeAsyncHooks.cpp", "jsSetAsyncHooksEnabled", 1);
 const { validateFunction, validateString, validateObject } = require("internal/validators");
@@ -131,24 +133,18 @@ function push(head: Frame | undefined, storage: AsyncLocalStorage, value: unknow
   return new Frame(storage, value, head, head === undefined ? undefined : unmask(head.masked, storage));
 }
 
-// `frame` with the innermost binding of `storage` removed. Frames above it are
-// copied (they may be shared with other captures); the tail below it is shared.
-// The view from the result is the view from `frame` minus that binding, masks
-// included.
+// `frame` with the binding of `storage` removed. Frames above it are copied
+// (they may be shared with other captures); the tail below it is shared. The
+// view from the result is the view from `frame` minus that binding, masks
+// included. run() (inlined on entry) and enterWith() both drop the old binding
+// this way, so a chain never holds two bindings of one storage and a capture
+// never keeps a shadowed value alive.
 function without(frame: Frame | undefined, storage: AsyncLocalStorage): Frame | undefined {
   var found = find(frame, storage);
   if (found === undefined) return frame;
   return copyUntil(frame!, found, found.prev);
 }
 
-// `frame` with every binding of `storage` removed (nested run() of one storage
-// stacks shadowed bindings).
-function withoutAll(frame: Frame | undefined, storage: AsyncLocalStorage): Frame | undefined {
-  var found = find(frame, storage);
-  if (found === undefined) return frame;
-  return copyUntil(frame!, found, withoutAll(found.prev, storage));
-}
-
 // Copies [from, stop) onto tail so that the view from the result is the view
 // from `from` minus what was cut out: the new head carries `from`'s mask (a
 // deeper frame's own mask describes lookups that start *there* and is not
@@ -261,34 +257,46 @@ class AsyncLocalStorage {
   run(store_value, callback, ...args) {
     $debug("run " + (this as any).__id__);
     var prior = get();
-    var before = lookup(prior, this);
-    var beforeValue = before !== undefined ? before.value : this.#defaultValue;
+    var bound = find(prior, this);
+    var beforeValue = bound !== undefined && !isMasked(prior, this) ? bound.value : this.#defaultValue;
     // Node short-circuits when the value is unchanged: no enterWith, no
     // finally-restore. Observable when the callback calls enterWith() —
     // the new value survives past run() (verified against Node v22/v26).
     if (sameValue(beforeValue, store_value)) {
       return callback.$apply(undefined, args);
     }
     var mutations = frameMutations;
-    // Shadows any outer binding of this storage: lookups stop at the innermost.
-    var frame = push(prior, this, store_value);
+    // Replace rather than shadow an outer binding of this storage, so a callback
+    // captured inside exit() or a nested run() does not retain the outer value.
+    // Only frames of other storages pushed since that binding are copied; when
+    // it is the innermost frame (or absent) this is just `prior` / `prior.prev`.
+    var frame = push(bound === undefined ? prior : copyUntil(prior!, bound, bound.prev), this, store_value);
     set(frame);
     try {
       // $apply, not a spread: spreading goes through Array.prototype[Symbol.iterator],
       // which userland can delete (node uses ReflectApply here for the same reason).
       return callback.$apply(undefined, args);
     } finally {
-      if (get() === frame && mutations === frameMutations) {
-        set(prior);
+      var head = get();
+      if (
+        mutations === frameMutations &&
+        (head === frame ||
+          (head !== undefined && head.prev === frame.prev && head.storage === this && head.masked === frame.masked))
+      ) {
+        // Node exits through enterWith(prior), a fresh frame object: a later
+        // disable() reaches continuations captured after run() returned but not
+        // ones captured before it, so `prior` itself must not become current
+        // again. An enclosing run() recognises the copy of its frame above.
+        set(prior === undefined ? undefined : new Frame(prior.storage, prior.value, prior.prev, prior.masked));
       } else {
         // enterWith()/disable() ran inside the callback. Node's finally is
         // enterWith(prior store): keep whatever else the callback installed and
         // rebind this storage to what getStore() returned on entry. Frames may
         // have been copied since (enterWith() of a storage bound further down
-        // copies everything above it), so go by value, not identity: drop every
+        // copies everything above it), so go by value, not identity: drop the
         // binding of this storage and put the prior one back on top. Enclosing
         // run()s of the same storage restore their own value likewise.
-        set(push(withoutAll(get(), this), this, beforeValue));
+        set(push(without(head, this), this, beforeValue));
       }
       $assert(sameValue(this.getStore(), beforeValue), "run: previous value was not restored");
     }

test/js/node/async_hooks/AsyncLocalStorage.test.ts

@@ -296,6 +296,33 @@ test("re-entering a storage inside run() does not grow the context", () => {
   });
 });
 
+// Node's run() exits through enterWith(prior), a fresh frame: a disable() after
+// run() returned reaches continuations captured since, not ones from before it.
+test("disable() after a run() does not reach continuations captured before the run()", () => {
+  const a = new AsyncLocalStorage();
+  const b = new AsyncLocalStorage();
+  const seen: unknown[] = [];
+  a.run("outer", () => {
+    b.run(2, () => {
+      const before = AsyncLocalStorage.snapshot();
+      a.run(1, () => {
+        b.run(5, () => {});
+      });
+      const after = AsyncLocalStorage.snapshot();
+      b.disable();
+      seen.push(
+        before(() => b.getStore()),
+        after(() => b.getStore()),
+        b.getStore(),
+        a.getStore(),
+      );
+    });
+    seen.push(a.getStore(), b.getStore());
+  });
+  seen.push(a.getStore(), b.getStore());
+  expect(seen).toEqual([2, undefined, undefined, "outer", "outer", undefined, undefined, undefined]);
+});
+
 // Node's run() enters a fresh frame, so a disable() inside it only reaches
 // continuations captured since; ones captured before keep their binding.
 test("disable() inside another storage's run() does not reach earlier continuations", async () => {
@@ -1556,3 +1583,45 @@ test("an active store adds no per-await / per-then helper allocations", () => {
   expect(keep.length).toBe(N * 3);
   expect(delta).toBeLessThan(50);
 });
+
+// A callback captured inside exit() or a nested run() cannot see the outer
+// store, so it must not keep it alive either.
+test("exit() and nested run() release the shadowed outer store", async () => {
+  const als = new AsyncLocalStorage();
+  const other = new AsyncLocalStorage();
+  const N = 20;
+  const timers: ReturnType<typeof setTimeout>[] = [];
+  const pending: Promise<unknown>[] = [];
+  const never = new Promise(() => {});
+  function arm() {
+    timers.push(setTimeout(() => {}, 1_000_000).unref());
+    pending.push(never.then(() => {}));
+  }
+  const modes = [
+    (fn: () => void) => als.exit(fn),
+    (fn: () => void) => als.run(undefined, fn),
+    (fn: () => void) => als.run({ inner: true }, fn),
+    (fn: () => void) => other.run(1, () => als.run({ inner: true }, () => other.run(2, fn))),
+  ];
+  const refs: WeakRef<object>[][] = modes.map(() => []);
+  for (const [m, mode] of modes.entries()) {
+    for (let i = 0; i < N; i++) {
+      const outer = { payload: new Uint8Array(64 * 1024) };
+      refs[m].push(new WeakRef(outer));
+      als.run(outer, () => {
+        mode(arm);
+        expect(als.getStore()).toBe(outer);
+      });
+    }
+  }
+  expect(als.getStore()).toBeUndefined();
+  const alive = () => refs.map(mode => mode.filter(r => r.deref() !== undefined).length);
+  for (let i = 0; i < 20 && Math.max(...alive()) > N / 2; i++) {
+    Bun.gc(true);
+    await new Promise(r => setTimeout(r, 10));
+  }
+  expect(timers.length + pending.length).toBe(modes.length * N * 2);
+  // without the fix every store of an affected mode is retained
+  for (const n of alive()) expect(n).toBeLessThanOrEqual(N / 2);
+  for (const t of timers) clearTimeout(t);
+});