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,8 +10,8 @@ const {
   ObjectDefineProperty,
   ObjectGetPrototypeOf,
   ObjectSetPrototypeOf,
+  PromisePrototype,
   PromisePrototypeThen,
-  PromiseReject,
   ReflectApply,
   SafeFinalizationRegistry,
   SafeMap,
@@ -29,12 +29,12 @@ 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;
 
 const { WeakReference } = require('internal/util');
-const { isPromise } = require('internal/util/types');
 
 // Can't delete when weakref count reaches 0 as it could increment again.
 // Only GC can be used as a valid time to clean up the channels map.
@@ -546,17 +546,21 @@ class TracingChannel {
     const { error } = this;
     const continuationWindow = this.#continuationWindow;
 
-    function reject(err) {
+    function onReject(err) {
       context.error = err;
       error.publish(context);
       // Use continuation window for asyncStart/asyncEnd
       // eslint-disable-next-line no-unused-vars
       using scope = continuationWindow.withScope(context);
       // TODO: Is there a way to have asyncEnd _after_ the continuation?
-      return PromiseReject(err);
     }
 
-    function resolve(result) {
+    function onRejectWithRethrow(err) {
+      onReject(err);
+      throw err;
+    }
+
+    function onResolve(result) {
       context.result = result;
       // Use continuation window for asyncStart/asyncEnd
       // eslint-disable-next-line no-unused-vars
@@ -576,12 +580,17 @@ class TracingChannel {
         context.result = result;
         return result;
       }
-      // For native Promises use PromisePrototypeThen to avoid user overrides.
-      if (isPromise(result)) {
-        return PromisePrototypeThen(result, 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 custom thenables, call .then() directly to preserve the thenable type.
-      return result.then(resolve, reject);
+      // 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);
 }));