async_hooks: avoid extra frame copies in AsyncLocalStorage.run()

With the AsyncContextFrame implementation, every
AsyncLocalStorage.run() built two copies of the current frame: one to
enter the new value and another one in the finally block to put the
previous value back. Both copies went through the Map iterator
protocol, which calls next() and allocates an entry array per store.

Copy frames with forEach() instead, and when fn() neither replaced the
current frame (enterWith()) nor removed a store from it (disable()),
put back the frame that was current before run() instead of building
a new one. Otherwise keep building a new frame as before, so that
those changes stay visible after run() returns. getStore() now does a
single lookup unless the stored value is undefined.

A side effect is that after a top-level run() returns the current
frame is undefined again, as it was before run(), instead of a frame
that holds the previous value of the store. getStore() returns the
same values either way.

Measured locally:

  als.run(i, noop)                          100 ns ->  16 ns
  als.run(i, noop), two other stores set    145 ns ->  35 ns
  await als.run(i, async fn with 2 awaits)  250 ns ->  89 ns
  async_hooks/async-local-storage-run.js    7.0M  ->  20M ops/s

Assisted-by: claude:opus-5.5
Signed-off-by: Gürgün Dayıoğlu <hey@gurgun.day>

Author: gurgunday

lib/internal/async_context_frame.js

@@ -9,9 +9,18 @@ const {
   setContinuationPreservedEmbedderData,
 } = internalBinding('async_context_frame');
 
+function copyEntry(value, key) {
+  this.set(key, value);
+}
+
 class AsyncContextFrame extends SafeMap {
   constructor(store, data) {
-    super(AsyncContextFrame.current());
+    super();
+    // Avoid the iterator protocol.
+    const prior = AsyncContextFrame.current();
+    if (prior !== undefined) {
+      prior.forEach(copyEntry, this);
+    }
     this.set(store, data);
   }
 

lib/internal/async_local_storage/async_context_frame.js

@@ -61,11 +61,20 @@ class AsyncLocalStorage {
     if (ObjectIs(prior, data)) {
       return ReflectApply(fn, null, args);
     }
-    this.enterWith(data);
+    const priorFrame = AsyncContextFrame.current();
+    const frame = new AsyncContextFrame(this, data);
+    const size = frame.size;
+    AsyncContextFrame.set(frame);
     try {
       return ReflectApply(fn, null, args);
     } finally {
-      this.enterWith(prior);
+      // Frames are only ever deleted from, so the same frame with the same
+      // size is untouched. Otherwise keep what fn() changed.
+      if (AsyncContextFrame.current() === frame && frame.size === size) {
+        AsyncContextFrame.set(priorFrame);
+      } else {
+        this.enterWith(prior);
+      }
     }
   }
 
@@ -75,10 +84,14 @@ class AsyncLocalStorage {
 
   getStore() {
     const frame = AsyncContextFrame.current();
-    if (!frame?.has(this)) {
+    if (frame === undefined) {
+      return this.#defaultValue;
+    }
+    const value = frame.get(this);
+    if (value === undefined && !frame.has(this)) {
       return this.#defaultValue;
     }
-    return frame?.get(this);
+    return value;
   }
 
   withScope(store) {

test/parallel/test-async-local-storage-run-restore.js

@@ -0,0 +1,187 @@
+'use strict';
+const common = require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+// Verify which store values are visible after asyncLocalStorage.run() returns.
+
+{
+  const als1 = new AsyncLocalStorage();
+  const als2 = new AsyncLocalStorage();
+
+  als1.run('outer', common.mustCall(() => {
+    als1.run('inner', common.mustCall(() => {
+      als2.enterWith('entered');
+      als1.enterWith('replaced');
+      assert.strictEqual(als1.getStore(), 'replaced');
+      assert.strictEqual(als2.getStore(), 'entered');
+    }));
+    assert.strictEqual(als1.getStore(), 'outer');
+    assert.strictEqual(als2.getStore(), 'entered');
+  }));
+  assert.strictEqual(als1.getStore(), undefined);
+  assert.strictEqual(als2.getStore(), 'entered');
+  als2.disable();
+}
+
+{
+  // enterWith() inside a run() with an unchanged value.
+  const als1 = new AsyncLocalStorage();
+  const als2 = new AsyncLocalStorage();
+
+  als1.run(NaN, common.mustCall(() => {
+    als1.run(NaN, common.mustCall(() => {
+      als2.enterWith('entered');
+    }));
+    assert.strictEqual(als2.getStore(), 'entered');
+  }));
+  assert.strictEqual(als1.getStore(), undefined);
+  assert.strictEqual(als2.getStore(), 'entered');
+  als2.disable();
+}
+
+{
+  const als1 = new AsyncLocalStorage();
+  const als2 = new AsyncLocalStorage();
+  const als3 = new AsyncLocalStorage();
+
+  als2.enterWith('store2');
+  als1.run('store1', common.mustCall(() => {
+    als3.disable();
+    als2.disable();
+    assert.strictEqual(als1.getStore(), 'store1');
+    assert.strictEqual(als2.getStore(), undefined);
+  }));
+  assert.strictEqual(als1.getStore(), undefined);
+  assert.strictEqual(als2.getStore(), undefined);
+  assert.strictEqual(als3.getStore(), undefined);
+
+  als1.run('outer', common.mustCall(() => {
+    als1.run('inner', common.mustCall(() => {
+      als1.disable();
+      assert.strictEqual(als1.getStore(), undefined);
+    }));
+    assert.strictEqual(als1.getStore(), 'outer');
+  }));
+  assert.strictEqual(als1.getStore(), undefined);
+
+  als2.enterWith('store2');
+  als1.run('store1', common.mustCall(() => {
+    als3.disable();
+  }));
+  assert.strictEqual(als1.getStore(), undefined);
+  assert.strictEqual(als2.getStore(), 'store2');
+  als2.disable();
+}
+
+{
+  // A store holding undefined does not fall back to defaultValue.
+  const als = new AsyncLocalStorage({ defaultValue: 'default' });
+  const other = new AsyncLocalStorage();
+
+  assert.strictEqual(als.getStore(), 'default');
+  als.run(undefined, common.mustCall(() => {
+    assert.strictEqual(als.getStore(), undefined);
+  }));
+  assert.strictEqual(als.getStore(), 'default');
+
+  als.run('value', common.mustCall(() => {
+    als.exit(common.mustCall(() => {
+      assert.strictEqual(als.getStore(), undefined);
+    }));
+    als.run(undefined, common.mustCall(() => {
+      assert.strictEqual(als.getStore(), undefined);
+    }));
+    assert.strictEqual(als.getStore(), 'value');
+  }));
+  assert.strictEqual(als.getStore(), 'default');
+
+  other.run('other', common.mustCall(() => {
+    assert.strictEqual(als.getStore(), 'default');
+    als.run(undefined, common.mustCall(() => {
+      assert.strictEqual(als.getStore(), undefined);
+      assert.strictEqual(other.getStore(), 'other');
+    }));
+    assert.strictEqual(als.getStore(), 'default');
+  }));
+  assert.strictEqual(als.getStore(), 'default');
+  assert.strictEqual(other.getStore(), undefined);
+
+  als.run('value', common.mustCall(() => {
+    als.disable();
+    assert.strictEqual(als.getStore(), 'default');
+  }));
+  assert.strictEqual(als.getStore(), 'default');
+}
+
+{
+  const als1 = new AsyncLocalStorage();
+  const als2 = new AsyncLocalStorage();
+
+  als1.run('outer', common.mustCall(() => {
+    als2.run('store2', common.mustCall(() => {
+      assert.throws(() => {
+        als1.run('inner', common.mustCall(() => {
+          assert.strictEqual(als1.getStore(), 'inner');
+          throw new Error('boom');
+        }));
+      }, /boom/);
+      assert.strictEqual(als1.getStore(), 'outer');
+      assert.strictEqual(als2.getStore(), 'store2');
+    }));
+  }));
+
+  assert.throws(() => {
+    als1.run('value', common.mustCall(() => {
+      throw new Error('boom');
+    }));
+  }, /boom/);
+  assert.strictEqual(als1.getStore(), undefined);
+  assert.strictEqual(als2.getStore(), undefined);
+}
+
+{
+  const als1 = new AsyncLocalStorage();
+  const als2 = new AsyncLocalStorage();
+
+  als1.run('a1', common.mustCall(() => {
+    als2.run('b1', common.mustCall(() => {
+      als1.run('a2', common.mustCall(() => {
+        assert.strictEqual(als1.getStore(), 'a2');
+        assert.strictEqual(als2.getStore(), 'b1');
+        setImmediate(common.mustCall(() => {
+          assert.strictEqual(als1.getStore(), 'a2');
+          assert.strictEqual(als2.getStore(), 'b1');
+        }));
+      }));
+      assert.strictEqual(als1.getStore(), 'a1');
+      assert.strictEqual(als2.getStore(), 'b1');
+      als2.run('b2', common.mustCall(() => {
+        assert.strictEqual(als1.getStore(), 'a1');
+        assert.strictEqual(als2.getStore(), 'b2');
+      }));
+      assert.strictEqual(als2.getStore(), 'b1');
+    }));
+    assert.strictEqual(als1.getStore(), 'a1');
+    assert.strictEqual(als2.getStore(), undefined);
+  }));
+  assert.strictEqual(als1.getStore(), undefined);
+  assert.strictEqual(als2.getStore(), undefined);
+
+  const promise = als1.run('async', common.mustCall(async () => {
+    await 1;
+    assert.strictEqual(als1.getStore(), 'async');
+    await als2.run('nested', common.mustCall(async () => {
+      await 1;
+      assert.strictEqual(als1.getStore(), 'async');
+      assert.strictEqual(als2.getStore(), 'nested');
+    }));
+    assert.strictEqual(als2.getStore(), undefined);
+    return als1.getStore();
+  }));
+  assert.strictEqual(als1.getStore(), undefined);
+  promise.then(common.mustCall((value) => {
+    assert.strictEqual(value, 'async');
+    assert.strictEqual(als1.getStore(), undefined);
+  }));
+}