Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

async_hooks: fix async/await context loss in AsyncLocalStorage #33189

Closed
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 65 additions & 6 deletions lib/async_hooks.js
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

const {
NumberIsSafeInteger,
PromiseResolve,
ReflectApply,
Symbol,
} = primordials;
Expand Down Expand Up @@ -211,19 +212,73 @@ 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) {
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);
}
});

const storageHookWithTracking = createHook({
init(asyncId, type, triggerAsyncId, resource) {
const currentResource = executionAsyncResource();
// Value of currentResource is always a non null object
propagateToStorageLists(resource, currentResource);

if (type === 'PROMISE' && !seenLayer[depth]) {
seenLayer[depth] = true;
patchPromiseBarrier(currentResource);
}
},

before(asyncId) {
depth++;
seenLayer[depth] = false;
},

after(asyncId) {
depth--;
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I’m -1 to this fix. This is going to slow everything down because of the use use of before and after hooks.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm thinking a config option like trackAsyncAwait which would switch between two hook sets so it only does the before/after if a user of AsyncLocalStorage has explicitly requested it. What do you think?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure. I think we should investigate if there is another possible fix first, this API is extremely nice and losing so much perf would not be good for the ecosystem.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, I'm trying to figure out a lower-level fix. I just made this as a possible higher-level solution for now, until we can come up with something better. Agreed it's not great though, needing the extra before/after hooks.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Considering that this problem is not new of AL, I don't think this should be rushed in.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, that's fine. Just putting something up that I can iterate on. If the possible lower-level fix is what it needs to be to land, so be it. :)

});

class AsyncLocalStorage {
constructor() {
constructor({ trackAsyncAwait = false } = {}) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not blocking comment: I'm not a big fan of user facing options which could be even named disableProblems...
Anyhow, some doc update is needed. And maybe we should use a more clear name as async await is tracked in general.

this.kResourceStore = Symbol('kResourceStore');
this.trackAsyncAwait = trackAsyncAwait;
this.enabled = false;
}

Expand All @@ -232,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();
}
}

Expand All @@ -250,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;
Expand Down
51 changes: 51 additions & 0 deletions test/parallel/test-async-local-storage-async-await.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
'use strict';
const common = require('../common');
const assert = require('assert');
const { AsyncLocalStorage } = require('async_hooks');

const store = new AsyncLocalStorage({ trackAsyncAwait: true });
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);
}));