From 3befe80c4fa864fbef445dc8eee7b81f22b184de Mon Sep 17 00:00:00 2001
From: Andrey Pechkurov <apechkurov@gmail.com>
Date: Wed, 4 Mar 2020 14:57:56 +0300
Subject: [PATCH] async_hooks: fix ctx loss after nested ALS calls

PR-URL: https://github.com/nodejs/node/pull/32085
Reviewed-By: Stephen Belanger <admin@stephenbelanger.com>
Reviewed-By: Vladimir de Turckheim <vlad2t@hotmail.com>
Reviewed-By: Michael Dawson <michael_dawson@ca.ibm.com>
---
 lib/async_hooks.js                            | 21 +++++----
 ...test-async-local-storage-enable-disable.js |  8 ++++
 .../test-async-local-storage-nested.js        | 44 +++++++++++++------
 3 files changed, 50 insertions(+), 23 deletions(-)

diff --git a/lib/async_hooks.js b/lib/async_hooks.js
index 3797baf183250a..bd3cd57d022f4b 100644
--- a/lib/async_hooks.js
+++ b/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;
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);