async_hooks: set unhandledRejection async context

This commit now executes `process.on('unhandledRejection')` in the
async execution context of the concerned `Promise`.

PR-URL: https://github.com/nodejs/node/pull/37281
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: Benjamin Gruenbaum <benjamingr@gmail.com>

5 files changed, 130 insertions, 11 deletions

diff --git a/lib/internal/async_hooks.js b/lib/internal/async_hooks.js
index f301431807..43ba749cd6 100644
--- a/lib/internal/async_hooks.js
+++ b/lib/internal/async_hooks.js
@@ -573,6 +573,8 @@ module.exports = {
   emitBefore: emitBeforeScript,
   emitAfter: emitAfterScript,
   emitDestroy: emitDestroyScript,
+  pushAsyncContext,
+  popAsyncContext,
   registerDestroyHook,
   useDomainTrampoline,
   nativeHooks: {
diff --git a/lib/internal/process/promises.js b/lib/internal/process/promises.js
index 023f7df036..c4eb9c53e2 100644
--- a/lib/internal/process/promises.js
+++ b/lib/internal/process/promises.js
@@ -24,6 +24,12 @@ const {
   triggerUncaughtException
 } = internalBinding('errors');
 
+const {
+  pushAsyncContext,
+  popAsyncContext,
+} = require('internal/async_hooks');
+const async_hooks = require('async_hooks');
+
 // *Must* match Environment::TickInfo::Fields in src/env.h.
 const kHasRejectionToWarn = 1;
 
@@ -116,11 +122,28 @@ function resolveError(type, promise, reason) {
 }
 
 function unhandledRejection(promise, reason) {
+  const asyncId = async_hooks.executionAsyncId();
+  const triggerAsyncId = async_hooks.triggerAsyncId();
+  const resource = promise;
+
+  const emit = (reason, promise, promiseInfo) => {
+    try {
+      pushAsyncContext(asyncId, triggerAsyncId, resource);
+      if (promiseInfo.domain) {
+        return promiseInfo.domain.emit('error', reason);
+      }
+      return process.emit('unhandledRejection', reason, promise);
+    } finally {
+      popAsyncContext(asyncId);
+    }
+  };
+
   maybeUnhandledPromises.set(promise, {
     reason,
     uid: ++lastPromiseId,
     warned: false,
-    domain: process.domain
+    domain: process.domain,
+    emit
   });
   // This causes the promise to be referenced at least for one tick.
   ArrayPrototypePush(pendingUnhandledRejections, promise);
@@ -194,13 +217,8 @@ function processPromiseRejections() {
       continue;
     }
     promiseInfo.warned = true;
-    const { reason, uid } = promiseInfo;
-    function emit(reason, promise, promiseInfo) {
-      if (promiseInfo.domain) {
-        return promiseInfo.domain.emit('error', reason);
-      }
-      return process.emit('unhandledRejection', reason, promise);
-    }
+    const { reason, uid, emit } = promiseInfo;
+
     switch (unhandledRejectionsMode) {
       case kStrictUnhandledRejections: {
         const err = reason instanceof Error ?
diff --git a/src/node_task_queue.cc b/src/node_task_queue.cc
index 0aa021fc5b..004eb6a652 100644
--- a/src/node_task_queue.cc
+++ b/src/node_task_queue.cc
@@ -1,3 +1,4 @@
+#include "async_wrap.h"
 #include "env-inl.h"
 #include "node.h"
 #include "node_errors.h"
@@ -16,11 +17,13 @@ using v8::Context;
 using v8::Function;
 using v8::FunctionCallbackInfo;
 using v8::Isolate;
+using v8::Just;
 using v8::kPromiseHandlerAddedAfterReject;
 using v8::kPromiseRejectAfterResolved;
 using v8::kPromiseRejectWithNoHandler;
 using v8::kPromiseResolveAfterResolved;
 using v8::Local;
+using v8::Maybe;
 using v8::Number;
 using v8::Object;
 using v8::Promise;
@@ -28,6 +31,40 @@ using v8::PromiseRejectEvent;
 using v8::PromiseRejectMessage;
 using v8::Value;
 
+static Maybe<double> GetAssignedPromiseAsyncId(Environment* env,
+                                               Local<Promise> promise,
+                                               Local<Value> id_symbol) {
+  Local<Value> maybe_async_id;
+  if (!promise->Get(env->context(), id_symbol).ToLocal(&maybe_async_id)) {
+    return v8::Just(AsyncWrap::kInvalidAsyncId);
+  }
+  return maybe_async_id->IsNumber()
+      ? maybe_async_id->NumberValue(env->context())
+      : v8::Just(AsyncWrap::kInvalidAsyncId);
+}
+
+static Maybe<double> GetAssignedPromiseWrapAsyncId(Environment* env,
+                                                   Local<Promise> promise,
+                                                   Local<Value> id_symbol) {
+  // This check is imperfect. If the internal field is set, it should
+  // be an object. If it's not, we just ignore it. Ideally v8 would
+  // have had GetInternalField returning a MaybeLocal but this works
+  // for now.
+  Local<Value> promiseWrap = promise->GetInternalField(0);
+  if (promiseWrap->IsObject()) {
+        Local<Value> maybe_async_id;
+    if (!promiseWrap.As<Object>()->Get(env->context(), id_symbol)
+        .ToLocal(&maybe_async_id)) {
+      return v8::Just(AsyncWrap::kInvalidAsyncId);
+    }
+    return maybe_async_id->IsNumber()
+        ? maybe_async_id->NumberValue(env->context())
+        : v8::Just(AsyncWrap::kInvalidAsyncId);
+  } else {
+      return v8::Just(AsyncWrap::kInvalidAsyncId);
+  }
+}
+
 void PromiseRejectCallback(PromiseRejectMessage message) {
   static std::atomic<uint64_t> unhandledRejections{0};
   static std::atomic<uint64_t> rejectionsHandledAfter{0};
@@ -76,12 +113,46 @@ void PromiseRejectCallback(PromiseRejectMessage message) {
 
   Local<Value> args[] = { type, promise, value };
 
-  // V8 does not expect this callback to have a scheduled exceptions once it
-  // returns, so we print them out in a best effort to do something about it
-  // without failing silently and without crashing the process.
+  double async_id = AsyncWrap::kInvalidAsyncId;
+  double trigger_async_id = AsyncWrap::kInvalidAsyncId;
   TryCatchScope try_catch(env);
+
+  if (!GetAssignedPromiseAsyncId(env, promise, env->async_id_symbol())
+          .To(&async_id)) return;
+  if (!GetAssignedPromiseAsyncId(env, promise, env->trigger_async_id_symbol())
+          .To(&trigger_async_id)) return;
+
+  if (async_id == AsyncWrap::kInvalidAsyncId &&
+      trigger_async_id == AsyncWrap::kInvalidAsyncId) {
+    // That means that promise might be a PromiseWrap, so we'll
+    // check there as well.
+    if (!GetAssignedPromiseWrapAsyncId(env, promise, env->async_id_symbol())
+              .To(&async_id)) return;
+    if (!GetAssignedPromiseWrapAsyncId(
+          env, promise, env->trigger_async_id_symbol())
+              .To(&trigger_async_id)) return;
+  }
+
+  if (async_id != AsyncWrap::kInvalidAsyncId &&
+      trigger_async_id != AsyncWrap::kInvalidAsyncId) {
+    env->async_hooks()->push_async_context(
+        async_id, trigger_async_id, promise);
+  }
+
   USE(callback->Call(
       env->context(), Undefined(isolate), arraysize(args), args));
+
+  if (async_id != AsyncWrap::kInvalidAsyncId &&
+      trigger_async_id != AsyncWrap::kInvalidAsyncId &&
+      env->execution_async_id() == async_id) {
+    // This condition might not be true if async_hooks was enabled during
+    // the promise callback execution.
+    env->async_hooks()->pop_async_context(async_id);
+  }
+
+  // V8 does not expect this callback to have a scheduled exceptions once it
+  // returns, so we print them out in a best effort to do something about it
+  // without failing silently and without crashing the process.
   if (try_catch.HasCaught() && !try_catch.HasTerminated()) {
     fprintf(stderr, "Exception in PromiseRejectCallback:\n");
     PrintCaughtException(isolate, env->context(), try_catch);
diff --git a/test/async-hooks/test-unhandled-rejection-context.js b/test/async-hooks/test-unhandled-rejection-context.js
new file mode 100644
index 0000000000..b70d97f99f
--- /dev/null
+++ b/test/async-hooks/test-unhandled-rejection-context.js
@@ -0,0 +1,27 @@
+'use strict';
+
+const common = require('../common');
+
+const assert = require('assert');
+const initHooks = require('./init-hooks');
+const async_hooks = require('async_hooks');
+
+if (!common.isMainThread)
+  common.skip('Worker bootstrapping works differently -> different async IDs');
+
+const promiseAsyncIds = [];
+const hooks = initHooks({
+  oninit(asyncId, type) {
+    if (type === 'PROMISE') {
+      promiseAsyncIds.push(asyncId);
+    }
+  }
+});
+
+hooks.enable();
+Promise.reject();
+
+process.on('unhandledRejection', common.mustCall(() => {
+  assert.strictEqual(promiseAsyncIds.length, 1);
+  assert.strictEqual(async_hooks.executionAsyncId(), promiseAsyncIds[0]);
+}));
diff --git a/test/parallel/test-bootstrap-modules.js b/test/parallel/test-bootstrap-modules.js
index 105bfb1086..4d47de41c2 100644
--- a/test/parallel/test-bootstrap-modules.js
+++ b/test/parallel/test-bootstrap-modules.js
@@ -102,6 +102,7 @@ const expectedModules = new Set([
   'NativeModule internal/worker/io',
   'NativeModule internal/worker/js_transferable',
   'NativeModule internal/blob',
+  'NativeModule async_hooks',
   'NativeModule path',
   'NativeModule stream',
   'NativeModule timers',