async_hooks: fix ctx loss after nested ALS calls

PR-URL: #32085 Reviewed-By: Stephen Belanger <[email protected]> Reviewed-By: Vladimir de Turckheim <[email protected]> Reviewed-By: Michael Dawson <[email protected]>

Author: puzpuzpuz

lib/async_hooks.js

@@ -255,23 +255,21 @@ class AsyncLocalStorage{
 resource[this.kResourceStore]=store;
 }
 
-_exit(){
-constresource=executionAsyncResource();
-if(resource){
-resource[this.kResourceStore]=undefined;
-}
-}
-
 runSyncAndReturn(store,callback, ...args){
+constresource=executionAsyncResource();
+constouterStore=resource[this.kResourceStore];
 this._enter(store);
 try{
 returncallback(...args);
 }finally{
-this._exit();
+resource[this.kResourceStore]=outerStore;
 }
 }
 
 exitSyncAndReturn(callback, ...args){
+if(!this.enabled){
+returncallback(...args);
+}
 this.enabled=false;
 try{
 returncallback(...args);
@@ -288,12 +286,17 @@ class AsyncLocalStorage{
 }
 
 run(store,callback, ...args){
+constresource=executionAsyncResource();
+constouterStore=resource[this.kResourceStore];
 this._enter(store);
 process.nextTick(callback, ...args);
-this._exit();
+resource[this.kResourceStore]=outerStore;
 }
 
 exit(callback, ...args){
+if(!this.enabled){
+returnprocess.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(newMap(),()=>{

test/async-hooks/test-async-local-storage-nested.js

@@ -4,19 +4,35 @@ const assert = require('assert');
 const{ AsyncLocalStorage }=require('async_hooks');
 
 constasyncLocalStorage=newAsyncLocalStorage();
+constouter={};
+constinner={};
 
-setTimeout(()=>{
-asyncLocalStorage.run(newMap(),()=>{
-constasyncLocalStorage2=newAsyncLocalStorage();
-asyncLocalStorage2.run(newMap(),()=>{
-conststore=asyncLocalStorage.getStore();
-conststore2=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);
-});
+functiontestInner(){
+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);