async_hooks: fix async/await context loss in AsyncLocalStorage

Author: Qard

lib/async_hooks.js

@@ -2,6 +2,7 @@
 
 const {
   NumberIsSafeInteger,
+  PromiseResolve,
   ReflectApply,
   Symbol,
 } = primordials;
@@ -211,13 +212,44 @@ class AsyncResource {
 }
 
 const storageList = [];
+const seenLayer = [];
+let depth = 0;
+
+function patchPromiseBarrier(currentResource) {
+  PromiseResolve({
+    then(resolve) {
+      const resource = executionAsyncResource();
+      propagateToStorageLists(resource, currentResource);
+      resolve();
+    }
+  });
+}
+
+function propagateToStorageLists(resource, currentResource) {
+  for (let i = 0; i < storageList.length; ++i) {
+    storageList[i]._propagate(resource, currentResource);
+  }
+}
+
 const storageHook = createHook({
   init(asyncId, type, triggerAsyncId, resource) {
     const currentResource = executionAsyncResource();
     // Value of currentResource is always a non null object
-    for (let i = 0; i < storageList.length; ++i) {
-      storageList[i]._propagate(resource, currentResource);
+    propagateToStorageLists(resource, currentResource);
+
+    if (type === 'PROMISE' && !seenLayer[depth]) {
+      seenLayer[depth] = true;
+      patchPromiseBarrier(currentResource);
     }
+  },
+
+  before(asyncId) {
+    depth++;
+    seenLayer[depth] = false;
+  },
+
+  after(asyncId) {
+    depth--;
   }
 });
 

test/parallel/test-async-local-storage-async-await.js

@@ -0,0 +1,51 @@
+'use strict';
+const common = require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const store = new AsyncLocalStorage();
+let checked = 0;
+
+function thenable(expected, count) {
+  return {
+    then: common.mustCall((cb) => {
+      assert.strictEqual(expected, store.getStore());
+      checked++;
+      cb();
+    }, count)
+  };
+}
+
+function main(n) {
+  const firstData = Symbol('first-data');
+  const secondData = Symbol('second-data');
+
+  const first = thenable(firstData, 1);
+  const second = thenable(secondData, 1);
+  const third = thenable(firstData, 2);
+
+  return store.run(firstData, common.mustCall(async () => {
+    assert.strictEqual(firstData, store.getStore());
+    await first;
+
+    await store.run(secondData, common.mustCall(async () => {
+      assert.strictEqual(secondData, store.getStore());
+      await second;
+      assert.strictEqual(secondData, store.getStore());
+    }));
+
+    await Promise.all([ third, third ]);
+    assert.strictEqual(firstData, store.getStore());
+  }));
+}
+
+const outerData = Symbol('outer-data');
+
+Promise.all([
+  store.run(outerData, () => Promise.resolve(thenable(outerData))),
+  Promise.resolve(3).then(common.mustCall(main)),
+  main(1),
+  main(2)
+]).then(common.mustCall(() => {
+  assert.strictEqual(checked, 13);
+}));