async_hooks: fix context loss after nested calls to AsyncLocalStorage

lib/async_hooks.js

@@ -255,23 +255,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);
@@ -288,12 +286,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;

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(), () => {

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);