diagnostics_channel: return original thenable

lib/diagnostics_channel.js

@@ -10,6 +10,8 @@
   ObjectDefineProperty,
   ObjectGetPrototypeOf,
   ObjectSetPrototypeOf,
+  Promise,
+  PromisePrototypeThen,
   ReflectApply,
   SafeFinalizationRegistry,
   SafeMap,
@@ -26,6 +28,9 @@
 } = require('internal/validators');
 
 const { triggerUncaughtException } = internalBinding('errors');
+const { isPromise } = require('internal/util/types');
+
+const PromisePrototype = Promise.prototype;
 
 const dc_binding = internalBinding('diagnostics_channel');
 const { subscribers: subscriberCounts } = dc_binding;
@@ -369,16 +374,20 @@
 
     const { start, end, asyncStart, asyncEnd, error } = this;
 
-    function reject(err) {
+    function onReject(err) {
       context.error = err;
       error.publish(context);
       asyncStart.publish(context);
       // TODO: Is there a way to have asyncEnd _after_ the continuation?
       asyncEnd.publish(context);
+    }
+
+    function onRejectWithRethrow(err) {
+      onReject(err);
       throw err;
     }
 
-    function resolve(result) {
+    function onResolve(result) {
       context.result = result;
       asyncStart.publish(context);
       // TODO: Is there a way to have asyncEnd _after_ the continuation?
@@ -396,7 +405,17 @@
           context.result = result;
           return result;
         }
-        return result.then(resolve, reject);
+        // isPromise() matches sub-classes, but we need to match only direct
+        // instances of the native Promise type to safely use PromisePrototypeThen.
+        if (isPromise(result) && ObjectGetPrototypeOf(result) === PromisePrototype) {
+          return PromisePrototypeThen(result, onResolve, onRejectWithRethrow);
+        }
+        // For non-native thenables, subscribe to the result but return the
+        // original thenable so the consumer can continue handling it directly.
+        // Non-native thenables don't have unhandledRejection tracking, so
+        // swallowing the rejection here doesn't change existing behaviour.
+        result.then(onResolve, onReject);
+        return result;
       } catch (err) {
         context.error = err;
         error.publish(context);

test/parallel/test-diagnostics-channel-tracing-channel-promise-spoofed-constructor.js

@@ -0,0 +1,52 @@
+'use strict';
+
+const common = require('../common');
+const dc = require('diagnostics_channel');
+const assert = require('assert');
+
+class SpoofedPromise extends Promise {
+  customMethod() {
+    return 'works';
+  }
+}
+
+const channel = dc.tracingChannel('test');
+
+const expectedResult = { foo: 'bar' };
+const input = { foo: 'bar' };
+const thisArg = { baz: 'buz' };
+
+function check(found) {
+  assert.strictEqual(found, input);
+}
+
+function checkAsync(found) {
+  check(found);
+  assert.strictEqual(found.error, undefined);
+  assert.deepStrictEqual(found.result, expectedResult);
+}
+
+const handlers = {
+  start: common.mustCall(check),
+  end: common.mustCall(check),
+  asyncStart: common.mustCall(checkAsync),
+  asyncEnd: common.mustCall(checkAsync),
+  error: common.mustNotCall()
+};
+
+channel.subscribe(handlers);
+
+let innerPromise;
+
+const result = channel.tracePromise(common.mustCall(function() {
+  innerPromise = SpoofedPromise.resolve(expectedResult);
+  // Spoof the constructor to try to trick the brand check
+  innerPromise.constructor = Promise;
+  return innerPromise;
+}), input, thisArg);
+
+// Despite the spoofed constructor, the original subclass instance should be
+// returned directly so that custom methods remain accessible.
+assert(result instanceof SpoofedPromise);
+assert.strictEqual(result, innerPromise);
+assert.strictEqual(result.customMethod(), 'works');

test/parallel/test-diagnostics-channel-tracing-channel-promise-thenable.js

@@ -12,6 +12,9 @@ class ResolvedThenable {
   then(resolve) {
     return new ResolvedThenable(resolve(this.#result));
   }
+  customMethod() {
+    return this.#result;
+  }
 }
 
 const channel = dc.tracingChannel('test');
@@ -49,7 +52,10 @@ const result = channel.tracePromise(common.mustCall(function(value) {
 }), input, thisArg, expectedResult);
 
 assert(result instanceof ResolvedThenable);
-assert.notStrictEqual(result, innerThenable);
+// With branching then, the original thenable is returned directly so that
+// extra methods defined on it remain accessible to the caller.
+assert.strictEqual(result, innerThenable);
+assert.deepStrictEqual(result.customMethod(), expectedResult);
 result.then(common.mustCall((value) => {
   assert.deepStrictEqual(value, expectedResult);
 }));