async_hooks: use new v8::Context PromiseHook API

Author: Qard

lib/internal/async_hooks.js

@@ -8,7 +8,6 @@ const {
   FunctionPrototypeBind,
   ObjectPrototypeHasOwnProperty,
   ObjectDefineProperty,
-  Promise,
   ReflectApply,
   Symbol,
 } = primordials;
@@ -57,7 +56,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.
@@ -298,71 +297,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,
+    );
   }
 }
 
@@ -378,8 +374,8 @@ function disableHooks() {
 
 function disablePromiseHookIfNecessary() {
   if (!wantPromiseHook) {
-    promiseHookMode = -1;
     disablePromiseHook();
+    setPromiseHooks(undefined, undefined, undefined, undefined);
   }
 }
 
@@ -453,6 +449,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

src/async_wrap.cc

@@ -452,6 +452,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();
+  Local<Value> init_hook = args[0];
+  Local<Value> before_hook = args[1];
+  Local<Value> after_hook = args[2];
+  Local<Value> resolve_hook = args[3];
+  ctx->SetPromiseHooks(init_hook, before_hook, after_hook, resolve_hook);
+}
 
 static void DisablePromiseHook(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);
@@ -631,6 +640,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);
 
@@ -725,6 +735,7 @@ void AsyncWrap::RegisterExternalReferences(
   registry->Register(ClearAsyncIdStack);
   registry->Register(QueueDestroyAsyncId);
   registry->Register(EnablePromiseHook);
+  registry->Register(SetPromiseHooks);
   registry->Register(DisablePromiseHook);
   registry->Register(RegisterDestroyHook);
   registry->Register(AsyncWrap::GetAsyncId);