From ec6b04867b73fe839c76fbbace204d7bfd37ec9e Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 30 Apr 2020 22:48:43 -0700 Subject: [PATCH 1/2] async_hooks: fix async/await context loss in AsyncLocalStorage --- lib/async_hooks.js | 36 ++++++++++++- .../test-async-local-storage-async-await.js | 51 +++++++++++++++++++ 2 files changed, 85 insertions(+), 2 deletions(-) create mode 100644 test/parallel/test-async-local-storage-async-await.js diff --git a/lib/async_hooks.js b/lib/async_hooks.js index 0943534790550c..f88225527d01b5 100644 --- a/lib/async_hooks.js +++ b/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--; } }); diff --git a/test/parallel/test-async-local-storage-async-await.js b/test/parallel/test-async-local-storage-async-await.js new file mode 100644 index 00000000000000..fdeb8af9a42648 --- /dev/null +++ b/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); +})); From cb867d7d6a258c13923642e904c3c1e39e2c9f0d Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Mon, 4 May 2020 17:52:36 -0700 Subject: [PATCH 2/2] async_hooks: add trackAsyncAwait option to AsyncLocalStorage --- lib/async_hooks.js | 35 ++++++++++++++++--- .../test-async-local-storage-async-await.js | 2 +- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/lib/async_hooks.js b/lib/async_hooks.js index f88225527d01b5..cb6fc183004dee 100644 --- a/lib/async_hooks.js +++ b/lib/async_hooks.js @@ -213,8 +213,22 @@ class AsyncResource { const storageList = []; const seenLayer = []; +let trackerCount = 0; let depth = 0; +function refreshStorageHooks() { + if (storageList.length === 0) { + storageHookWithTracking.disable(); + storageHook.disable(); + } else if (trackerCount > 0) { + storageHookWithTracking.enable(); + storageHook.disable(); + } else { + storageHookWithTracking.disable(); + storageHook.enable(); + } +} + function patchPromiseBarrier(currentResource) { PromiseResolve({ then(resolve) { @@ -232,6 +246,14 @@ function propagateToStorageLists(resource, currentResource) { } const storageHook = createHook({ + init(asyncId, type, triggerAsyncId, resource) { + const currentResource = executionAsyncResource(); + // Value of currentResource is always a non null object + propagateToStorageLists(resource, currentResource); + } +}); + +const storageHookWithTracking = createHook({ init(asyncId, type, triggerAsyncId, resource) { const currentResource = executionAsyncResource(); // Value of currentResource is always a non null object @@ -254,8 +276,9 @@ const storageHook = createHook({ }); class AsyncLocalStorage { - constructor() { + constructor({ trackAsyncAwait = false } = {}) { this.kResourceStore = Symbol('kResourceStore'); + this.trackAsyncAwait = trackAsyncAwait; this.enabled = false; } @@ -264,9 +287,10 @@ class AsyncLocalStorage { this.enabled = false; // If this.enabled, the instance must be in storageList storageList.splice(storageList.indexOf(this), 1); - if (storageList.length === 0) { - storageHook.disable(); + if (this.trackAsyncAwait) { + trackerCount--; } + refreshStorageHooks(); } } @@ -282,7 +306,10 @@ class AsyncLocalStorage { if (!this.enabled) { this.enabled = true; storageList.push(this); - storageHook.enable(); + if (this.trackAsyncAwait) { + trackerCount++; + } + refreshStorageHooks(); } const resource = executionAsyncResource(); resource[this.kResourceStore] = store; diff --git a/test/parallel/test-async-local-storage-async-await.js b/test/parallel/test-async-local-storage-async-await.js index fdeb8af9a42648..f0ec016fa0792c 100644 --- a/test/parallel/test-async-local-storage-async-await.js +++ b/test/parallel/test-async-local-storage-async-await.js @@ -3,7 +3,7 @@ const common = require('../common'); const assert = require('assert'); const { AsyncLocalStorage } = require('async_hooks'); -const store = new AsyncLocalStorage(); +const store = new AsyncLocalStorage({ trackAsyncAwait: true }); let checked = 0; function thenable(expected, count) {