diff --git a/lib/internal/async_context_frame.js b/lib/internal/async_context_frame.js index c12e74124a25..2807d297c7e7 100644 --- a/lib/internal/async_context_frame.js +++ b/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); } diff --git a/lib/internal/async_local_storage/async_context_frame.js b/lib/internal/async_local_storage/async_context_frame.js index 0fd20b79535d..fdb34f3d4c48 100644 --- a/lib/internal/async_local_storage/async_context_frame.js +++ b/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) { diff --git a/test/parallel/test-async-local-storage-run-restore.js b/test/parallel/test-async-local-storage-run-restore.js new file mode 100644 index 000000000000..5f10fd95efce --- /dev/null +++ b/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); + })); +}