stream: preserve AsyncLocalStorage on finished only when needed

PR-URL: #59873
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Daniel Lemire <daniel@lemire.me>
Reviewed-By: Yagiz Nizipli <yagiz@nizipli.com>

Author: avcribl

benchmark/streams/finished.js

@@ -0,0 +1,33 @@
+'use strict';
+
+constcommon=require('../common');
+const{ Readable, Writable }=require('stream');
+const{ finished }=require('stream/promises');
+
+constbench=common.createBenchmark(main,{
+n: [1e7],
+streamType: ['readable','writable'],
+});
+
+asyncfunctionmain({ n, streamType }){
+bench.start();
+
+for(leti=0;i<n;i++){
+letstream;
+
+switch(streamType){
+case'readable':
+stream=newReadable({read(){this.push(null);}});
+stream.resume();
+break;
+case'writable':
+stream=newWritable({write(chunk,enc,cb){cb();}});
+stream.end();
+break;
+}
+
+awaitfinished(stream);
+}
+
+bench.end(n);
+}

lib/internal/streams/end-of-stream.js

@@ -44,6 +44,9 @@ const {
  kIsClosedPromise,
 }=require('internal/streams/utils');
 
+const{ getHookArrays }=require('internal/async_hooks');
+constAsyncContextFrame=require('internal/async_context_frame');
+
 // Lazy load
 letAsyncResource;
 letaddAbortListener;
@@ -74,9 +77,14 @@ function eos(stream, options, callback) {
 validateFunction(callback,'callback');
 validateAbortSignal(options.signal,'options.signal');
 
-// Avoid AsyncResource.bind() because it calls ObjectDefineProperties which
-// is a bottleneck here.
-callback=once(bindAsyncResource(callback,'STREAM_END_OF_STREAM'));
+if(AsyncContextFrame.current()||
+getHookArrays()[0].length>0){
+// Avoid AsyncResource.bind() because it calls ObjectDefineProperties which
+// is a bottleneck here.
+callback=once(bindAsyncResource(callback,'STREAM_END_OF_STREAM'));
+}else{
+callback=once(callback);
+}
 
 if(isReadableStream(stream)||isWritableStream(stream)){
 returneosWeb(stream,options,callback);

test/parallel/test-stream-finished-async-local-storage.js

@@ -0,0 +1,24 @@
+// Flags: --expose-internals
+'use strict';
+
+constcommon=require('../common');
+const{ Readable, finished }=require('stream');
+const{ AsyncLocalStorage }=require('async_hooks');
+const{ strictEqual }=require('assert');
+constAsyncContextFrame=require('internal/async_context_frame');
+constinternalAsyncHooks=require('internal/async_hooks');
+
+// This test verifies that ALS context is preserved when using stream.finished()
+
+constals=newAsyncLocalStorage();
+constreadable=newReadable();
+
+als.run('test-context-1',()=>{
+finished(readable,common.mustCall(()=>{
+strictEqual(AsyncContextFrame.enabled||internalAsyncHooks.getHookArrays()[0].length>0,
+true,'One of AsyncContextFrame or async hooks criteria should be met');
+strictEqual(als.getStore(),'test-context-1','ALS context should be preserved');
+}));
+});
+
+readable.destroy();

test/parallel/test-stream-finished-bindAsyncResource-path.js

@@ -0,0 +1,35 @@
+// Flags: --expose-internals
+'use strict';
+
+constcommon=require('../common');
+const{ Readable, finished }=require('stream');
+const{ createHook, executionAsyncId }=require('async_hooks');
+const{ strictEqual }=require('assert');
+constinternalAsyncHooks=require('internal/async_hooks');
+
+// This test verifies that when there are active async hooks, stream.finished() uses
+// the bindAsyncResource path
+
+createHook({
+init(asyncId,type,triggerAsyncId){
+if(type==='STREAM_END_OF_STREAM'){
+constparentContext=contextMap.get(triggerAsyncId);
+contextMap.set(asyncId,parentContext);
+}
+}
+}).enable();
+
+constcontextMap=newMap();
+constasyncId=executionAsyncId();
+contextMap.set(asyncId,'abc-123');
+constreadable=newReadable();
+
+finished(readable,common.mustCall(()=>{
+constcurrentAsyncId=executionAsyncId();
+constctx=contextMap.get(currentAsyncId);
+strictEqual(internalAsyncHooks.getHookArrays()[0].length>0,
+true,'Should have active user async hook');
+strictEqual(ctx,'abc-123','Context should be preserved');
+}));
+
+readable.destroy();

test/parallel/test-stream-finished-default-path.js

@@ -0,0 +1,21 @@
+// Flags: --expose-internals --no-async-context-frame
+'use strict';
+
+constcommon=require('../common');
+const{ Readable, finished }=require('stream');
+const{ strictEqual }=require('assert');
+constAsyncContextFrame=require('internal/async_context_frame');
+constinternalAsyncHooks=require('internal/async_hooks');
+
+// This test verifies that when there are no active async hooks, stream.finished() uses the default callback path
+
+constreadable=newReadable();
+
+finished(readable,common.mustCall(()=>{
+strictEqual(internalAsyncHooks.getHookArrays()[0].length===0,
+true,'Should not have active user async hook');
+strictEqual(AsyncContextFrame.current()||internalAsyncHooks.getHookArrays()[0].length>0,
+false,'Default callback path should be used');
+}));
+
+readable.destroy();