async_hooks: use new v8::Context PromiseHook API

PR-URL: #36394 Reviewed-By: Bryan English <bryan@bryanenglish.com> Reviewed-By: Gus Caplan <me@gus.host> Reviewed-By: Vladimir de Turckheim <vlad2t@hotmail.com> Reviewed-By: Gerhard Stöbich <deb2001-github@yahoo.de> Reviewed-By: James M Snell <jasnell@gmail.com>
nodejs · Qard · Apr 21, 2021 · Apr 22, 2021 · Apr 27, 2021 · Apr 29, 2021
commit fb36f98ff5a926c143ac7f4799a77922f325484e
diff --git a/lib/internal/async_hooks.js b/lib/internal/async_hooks.js
@@ -5,7 +5,6 @@ const {
   ErrorCaptureStackTrace,
   ObjectPrototypeHasOwnProperty,
   ObjectDefineProperty,
-  Promise,
   Symbol,
 } = primordials;
 
@@ -53,7 +52,7 @@ const {
   clearAsyncIdStack,
 } = async_wrap;
 // For performance reasons, only track Promises when a hook is enabled.
-const { enablePromiseHook, disablePromiseHook } = async_wrap;
+const { enablePromiseHook, disablePromiseHook, setPromiseHooks } = async_wrap;
 // Properties in active_hooks are used to keep track of the set of hooks being
 // executed in case another hook is enabled/disabled. The new set of hooks is
 // then restored once the active set of hooks is finished executing.
@@ -303,71 +302,68 @@ function restoreActiveHooks() {
   active_hooks.tmp_fields = null;
 }
 
-function trackPromise(promise, parent, silent) {
-  const asyncId = getOrSetAsyncId(promise);
+function trackPromise(promise, parent) {
+  if (promise[async_id_symbol]) {
+    return;
+  }
 
+  promise[async_id_symbol] = newAsyncId();
   promise[trigger_async_id_symbol] = parent ? getOrSetAsyncId(parent) :
     getDefaultTriggerAsyncId();
+}
 
-  if (!silent && initHooksExist()) {
-    const triggerId = promise[trigger_async_id_symbol];
-    emitInitScript(asyncId, 'PROMISE', triggerId, promise);
-  }
+function promiseInitHook(promise, parent) {
+  trackPromise(promise, parent);
+  const asyncId = promise[async_id_symbol];
+  const triggerAsyncId = promise[trigger_async_id_symbol];
+  emitInitScript(asyncId, 'PROMISE', triggerAsyncId, promise);
 }
 
-function fastPromiseHook(type, promise, parent) {
-  if (type === kInit || !promise[async_id_symbol]) {
-    const silent = type !== kInit;
-    if (parent instanceof Promise) {
-      trackPromise(promise, parent, silent);
-    } else {
-      trackPromise(promise, null, silent);
-    }
+function promiseBeforeHook(promise) {
+  trackPromise(promise);
+  const asyncId = promise[async_id_symbol];
+  const triggerId = promise[trigger_async_id_symbol];
+  emitBeforeScript(asyncId, triggerId, promise);
+}
 
-    if (!silent) return;
+function promiseAfterHook(promise) {
+  trackPromise(promise);
+  const asyncId = promise[async_id_symbol];
+  if (hasHooks(kAfter)) {
+    emitAfterNative(asyncId);
   }
+  if (asyncId === executionAsyncId()) {
+    // This condition might not be true if async_hooks was enabled during
+    // the promise callback execution.
+    // Popping it off the stack can be skipped in that case, because it is
+    // known that it would correspond to exactly one call with
+    // PromiseHookType::kBefore that was not witnessed by the PromiseHook.
+    popAsyncContext(asyncId);
+  }
+}
 
+function promiseResolveHook(promise) {
+  trackPromise(promise);
   const asyncId = promise[async_id_symbol];
-  switch (type) {
-    case kBefore:
-      const triggerId = promise[trigger_async_id_symbol];
-      emitBeforeScript(asyncId, triggerId, promise);
-      break;
-    case kAfter:
-      if (hasHooks(kAfter)) {
-        emitAfterNative(asyncId);
-      }
-      if (asyncId === executionAsyncId()) {
-        // This condition might not be true if async_hooks was enabled during
-        // the promise callback execution.
-        // Popping it off the stack can be skipped in that case, because it is
-        // known that it would correspond to exactly one call with
-        // PromiseHookType::kBefore that was not witnessed by the PromiseHook.
-        popAsyncContext(asyncId);
-      }
-      break;
-    case kPromiseResolve:
-      emitPromiseResolveNative(asyncId);
-      break;
-  }
+  emitPromiseResolveNative(asyncId);
 }
 
 let wantPromiseHook = false;
 function enableHooks() {
   async_hook_fields[kCheck] += 1;
 }
 
-let promiseHookMode = -1;
 function updatePromiseHookMode() {
   wantPromiseHook = true;
   if (destroyHooksExist()) {
-    if (promiseHookMode !== 1) {
-      promiseHookMode = 1;
-      enablePromiseHook();
-    }
-  } else if (promiseHookMode !== 0) {
-    promiseHookMode = 0;
-    enablePromiseHook(fastPromiseHook);
+    enablePromiseHook();
+  } else {
+    setPromiseHooks(
+      initHooksExist() ? promiseInitHook : undefined,
+      promiseBeforeHook,
+      promiseAfterHook,
+      promiseResolveHooksExist() ? promiseResolveHook : undefined,
+    );
   }
 }
 
@@ -383,8 +379,8 @@ function disableHooks() {
 
 function disablePromiseHookIfNecessary() {
   if (!wantPromiseHook) {
-    promiseHookMode = -1;
     disablePromiseHook();
+    setPromiseHooks(undefined, undefined, undefined, undefined);
   }
 }
 
@@ -458,6 +454,10 @@ function destroyHooksExist() {
   return hasHooks(kDestroy);
 }
 
+function promiseResolveHooksExist() {
+  return hasHooks(kPromiseResolve);
+}
+
 
 function emitInitScript(asyncId, type, triggerAsyncId, resource) {
   // Short circuit all checks for the common case. Which is that no hooks have

diff --git a/src/async_wrap.cc b/src/async_wrap.cc
@@ -491,6 +491,15 @@ static void EnablePromiseHook(const FunctionCallbackInfo<Value>& args) {
   }
 }
 
+static void SetPromiseHooks(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  Local<Context> ctx = env->context();
+  ctx->SetPromiseHooks(
+    args[0]->IsFunction() ? args[0].As<Function>() : Local<Function>(),
+    args[1]->IsFunction() ? args[1].As<Function>() : Local<Function>(),
+    args[2]->IsFunction() ? args[2].As<Function>() : Local<Function>(),
+    args[3]->IsFunction() ? args[3].As<Function>() : Local<Function>());
+}
 
 static void DisablePromiseHook(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);
@@ -670,6 +679,7 @@ void AsyncWrap::Initialize(Local<Object> target,
   env->SetMethod(target, "clearAsyncIdStack", ClearAsyncIdStack);
   env->SetMethod(target, "queueDestroyAsyncId", QueueDestroyAsyncId);
   env->SetMethod(target, "enablePromiseHook", EnablePromiseHook);
+  env->SetMethod(target, "setPromiseHooks", SetPromiseHooks);
   env->SetMethod(target, "disablePromiseHook", DisablePromiseHook);
   env->SetMethod(target, "registerDestroyHook", RegisterDestroyHook);