diff --git a/lib/async_hooks.js b/lib/async_hooks.js index 65ecaf8a620fdd..47deffe84e233f 100644 --- a/lib/async_hooks.js +++ b/lib/async_hooks.js @@ -256,23 +256,21 @@ class AsyncLocalStorage { resource[this.kResourceStore] = store; } - _exit() { - const resource = executionAsyncResource(); - if (resource) { - resource[this.kResourceStore] = undefined; - } - } - runSyncAndReturn(store, callback, ...args) { + const resource = executionAsyncResource(); + const outerStore = resource[this.kResourceStore]; this._enter(store); try { return callback(...args); } finally { - this._exit(); + resource[this.kResourceStore] = outerStore; } } exitSyncAndReturn(callback, ...args) { + if (!this.enabled) { + return callback(...args); + } this.enabled = false; try { return callback(...args); @@ -289,12 +287,17 @@ class AsyncLocalStorage { } run(store, callback, ...args) { + const resource = executionAsyncResource(); + const outerStore = resource[this.kResourceStore]; this._enter(store); process.nextTick(callback, ...args); - this._exit(); + resource[this.kResourceStore] = outerStore; } exit(callback, ...args) { + if (!this.enabled) { + return process.nextTick(callback, ...args); + } this.enabled = false; process.nextTick(callback, ...args); this.enabled = true; diff --git a/test/async-hooks/test-async-local-storage-enable-disable.js b/test/async-hooks/test-async-local-storage-enable-disable.js index 22a3f5f6c8c43f..93132079827eeb 100644 --- a/test/async-hooks/test-async-local-storage-enable-disable.js +++ b/test/async-hooks/test-async-local-storage-enable-disable.js @@ -12,8 +12,16 @@ asyncLocalStorage.runSyncAndReturn(new Map(), () => { process.nextTick(() => { assert.strictEqual(asyncLocalStorage.getStore(), undefined); }); + asyncLocalStorage.disable(); assert.strictEqual(asyncLocalStorage.getStore(), undefined); + + // Calls to exit() should not mess with enabled status + asyncLocalStorage.exit(() => { + assert.strictEqual(asyncLocalStorage.getStore(), undefined); + }); + assert.strictEqual(asyncLocalStorage.getStore(), undefined); + process.nextTick(() => { assert.strictEqual(asyncLocalStorage.getStore(), undefined); asyncLocalStorage.runSyncAndReturn(new Map(), () => { diff --git a/test/async-hooks/test-async-local-storage-nested.js b/test/async-hooks/test-async-local-storage-nested.js index 1409a8ebc82a04..143d5d45de9e25 100644 --- a/test/async-hooks/test-async-local-storage-nested.js +++ b/test/async-hooks/test-async-local-storage-nested.js @@ -4,19 +4,35 @@ const assert = require('assert'); const { AsyncLocalStorage } = require('async_hooks'); const asyncLocalStorage = new AsyncLocalStorage(); +const outer = {}; +const inner = {}; -setTimeout(() => { - asyncLocalStorage.run(new Map(), () => { - const asyncLocalStorage2 = new AsyncLocalStorage(); - asyncLocalStorage2.run(new Map(), () => { - const store = asyncLocalStorage.getStore(); - const store2 = asyncLocalStorage2.getStore(); - store.set('hello', 'world'); - store2.set('hello', 'foo'); - setTimeout(() => { - assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'world'); - assert.strictEqual(asyncLocalStorage2.getStore().get('hello'), 'foo'); - }, 200); - }); +function testInner() { + assert.strictEqual(asyncLocalStorage.getStore(), outer); + + asyncLocalStorage.run(inner, () => { + assert.strictEqual(asyncLocalStorage.getStore(), inner); + }); + assert.strictEqual(asyncLocalStorage.getStore(), outer); + + asyncLocalStorage.exit(() => { + assert.strictEqual(asyncLocalStorage.getStore(), undefined); + }); + assert.strictEqual(asyncLocalStorage.getStore(), outer); + + asyncLocalStorage.runSyncAndReturn(inner, () => { + assert.strictEqual(asyncLocalStorage.getStore(), inner); }); -}, 100); + assert.strictEqual(asyncLocalStorage.getStore(), outer); + + asyncLocalStorage.exitSyncAndReturn(() => { + assert.strictEqual(asyncLocalStorage.getStore(), undefined); + }); + assert.strictEqual(asyncLocalStorage.getStore(), outer); +} + +asyncLocalStorage.run(outer, testInner); +assert.strictEqual(asyncLocalStorage.getStore(), undefined); + +asyncLocalStorage.runSyncAndReturn(outer, testInner); +assert.strictEqual(asyncLocalStorage.getStore(), undefined);