diagnostics_channel: return original thenable

This makes tracePromise return the original thenable to allow custom
thenable types to retain their methods rather than producing the
chained result type.

Author: Qard

lib/diagnostics_channel.js

@@ -10,6 +10,8 @@ const {
   ObjectDefineProperty,
   ObjectGetPrototypeOf,
   ObjectSetPrototypeOf,
+  Promise,
+  PromisePrototypeThen,
   ReflectApply,
   SafeFinalizationRegistry,
   SafeMap,
@@ -26,6 +28,7 @@ const {
 } = require('internal/validators');
 
 const { triggerUncaughtException } = internalBinding('errors');
+const { isPromise } = require('internal/util/types');
 
 const dc_binding = internalBinding('diagnostics_channel');
 const { subscribers: subscriberCounts } = dc_binding;
@@ -369,16 +372,20 @@ class TracingChannel {
 
     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 +403,17 @@ class TracingChannel {
           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) && result.constructor === Promise) {
+          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-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);
 }));