[async_wrap] enable/disable PromiseHook depending on kTotals

lib/async_hooks.js

@@ -38,8 +38,8 @@ var tmp_async_hook_fields = null;
 // Each constant tracks how many callbacks there are for any given step of
 // async execution. These are tracked so if the user didn't include callbacks
 // for a given step, that step can bail out early.
-const { kInit, kBefore, kAfter, kDestroy, kCurrentAsyncId, kCurrentTriggerId,
-    kAsyncUidCntr, kInitTriggerId } = async_wrap.constants;
+const { kInit, kBefore, kAfter, kDestroy, kTotals, kCurrentAsyncId,
+  kCurrentTriggerId, kAsyncUidCntr, kInitTriggerId } = async_wrap.constants;
 
 const { async_id_symbol, trigger_id_symbol } = async_wrap;
 
@@ -49,7 +49,14 @@ const before_symbol = Symbol('before');
 const after_symbol = Symbol('after');
 const destroy_symbol = Symbol('destroy');
 
-let setupHooksCalled = false;
+// Setup the callbacks that node::AsyncWrap will call when there are hooks to
+// process. They use the same functions as the JS embedder API. These callbacks
+// are setup immediately to prevent async_wrap.setupHooks() from being hijacked
+// and the cost of doing so is negligible.
+async_wrap.setupHooks({ init,
+                        before: emitBeforeN,
+                        after: emitAfterN,
+                        destroy: emitDestroyN });
 
 // Used to fatally abort the process if a callback throws.
 function fatalError(e) {
@@ -98,24 +105,21 @@ class AsyncHook {
     if (hooks_array.includes(this))
       return this;
 
-    if (!setupHooksCalled) {
-      setupHooksCalled = true;
-      // Setup the callbacks that node::AsyncWrap will call when there are
-      // hooks to process. They use the same functions as the JS embedder API.
-      async_wrap.setupHooks({ init,
-                              before: emitBeforeN,
-                              after: emitAfterN,
-                              destroy: emitDestroyN });
-    }
+    const prev_kTotals = hook_fields[kTotals];
+    hook_fields[kTotals] = 0;
 
     // createHook() has already enforced that the callbacks are all functions,
     // so here simply increment the count of whether each callbacks exists or
     // not.
-    hook_fields[kInit] += +!!this[init_symbol];
-    hook_fields[kBefore] += +!!this[before_symbol];
-    hook_fields[kAfter] += +!!this[after_symbol];
-    hook_fields[kDestroy] += +!!this[destroy_symbol];
+    hook_fields[kTotals] += hook_fields[kInit] += +!!this[init_symbol];
+    hook_fields[kTotals] += hook_fields[kBefore] += +!!this[before_symbol];
+    hook_fields[kTotals] += hook_fields[kAfter] += +!!this[after_symbol];
+    hook_fields[kTotals] += hook_fields[kDestroy] += +!!this[destroy_symbol];
     hooks_array.push(this);
+
+    if (prev_kTotals === 0 && hook_fields[kTotals] > 0)
+      async_wrap.enablePromiseHook();
+
     return this;
   }
 
@@ -126,11 +130,18 @@ class AsyncHook {
     if (index === -1)
       return this;
 
-    hook_fields[kInit] -= +!!this[init_symbol];
-    hook_fields[kBefore] -= +!!this[before_symbol];
-    hook_fields[kAfter] -= +!!this[after_symbol];
-    hook_fields[kDestroy] -= +!!this[destroy_symbol];
+    const prev_kTotals = hook_fields[kTotals];
+    hook_fields[kTotals] = 0;
+
+    hook_fields[kTotals] += hook_fields[kInit] -= +!!this[init_symbol];
+    hook_fields[kTotals] += hook_fields[kBefore] -= +!!this[before_symbol];
+    hook_fields[kTotals] += hook_fields[kAfter] -= +!!this[after_symbol];
+    hook_fields[kTotals] += hook_fields[kDestroy] -= +!!this[destroy_symbol];
     hooks_array.splice(index, 1);
+
+    if (prev_kTotals > 0 && hook_fields[kTotals] === 0)
+      async_wrap.disablePromiseHook();
+
     return this;
   }
 }

src/async-wrap.cc

@@ -335,12 +335,17 @@ static void PromiseHook(PromiseHookType type, Local<Promise> promise,
                         Local<Value> parent, void* arg) {
   Local<Context> context = promise->CreationContext();
   Environment* env = Environment::GetCurrent(context);
+
+  // PromiseHook() should never be called if no hooks have been enabled.
+  CHECK_GT(env->async_hooks()->fields()[AsyncHooks::kTotals], 0);
+
   Local<Value> resource_object_value = promise->GetInternalField(0);
   PromiseWrap* wrap = nullptr;
   if (resource_object_value->IsObject()) {
     Local<Object> resource_object = resource_object_value.As<Object>();
     wrap = Unwrap<PromiseWrap>(resource_object);
   }
+
   if (type == PromiseHookType::kInit || wrap == nullptr) {
     bool silent = type != PromiseHookType::kInit;
     PromiseWrap* parent_wrap = nullptr;
@@ -368,6 +373,7 @@ static void PromiseHook(PromiseHookType type, Local<Promise> promise,
   } else if (type == PromiseHookType::kResolve) {
     // TODO(matthewloring): need to expose this through the async hooks api.
   }
+
   CHECK_NE(wrap, nullptr);
   if (type == PromiseHookType::kBefore) {
     PreCallbackExecution(wrap, false);
@@ -401,7 +407,6 @@ static void SetupHooks(const FunctionCallbackInfo<Value>& args) {
   SET_HOOK_FN(before);
   SET_HOOK_FN(after);
   SET_HOOK_FN(destroy);
-  env->AddPromiseHook(PromiseHook, nullptr);
 #undef SET_HOOK_FN
 
   {
@@ -422,6 +427,18 @@ static void SetupHooks(const FunctionCallbackInfo<Value>& args) {
 }
 
 
+static void EnablePromiseHook(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  env->AddPromiseHook(PromiseHook, nullptr);
+}
+
+
+static void DisablePromiseHook(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  env->RemovePromiseHook(PromiseHook, nullptr);
+}
+
+
 void AsyncWrap::GetAsyncId(const FunctionCallbackInfo<Value>& args) {
   AsyncWrap* wrap;
   args.GetReturnValue().Set(-1);
@@ -478,6 +495,8 @@ void AsyncWrap::Initialize(Local<Object> target,
   env->SetMethod(target, "popAsyncIds", PopAsyncIds);
   env->SetMethod(target, "clearIdStack", ClearIdStack);
   env->SetMethod(target, "addIdToDestroyList", QueueDestroyId);
+  env->SetMethod(target, "enablePromiseHook", EnablePromiseHook);
+  env->SetMethod(target, "disablePromiseHook", DisablePromiseHook);
 
   v8::PropertyAttribute ReadOnlyDontDelete =
       static_cast<v8::PropertyAttribute>(v8::ReadOnly | v8::DontDelete);
@@ -528,6 +547,7 @@ void AsyncWrap::Initialize(Local<Object> target,
   SET_HOOKS_CONSTANT(kBefore);
   SET_HOOKS_CONSTANT(kAfter);
   SET_HOOKS_CONSTANT(kDestroy);
+  SET_HOOKS_CONSTANT(kTotals);
   SET_HOOKS_CONSTANT(kCurrentAsyncId);
   SET_HOOKS_CONSTANT(kCurrentTriggerId);
   SET_HOOKS_CONSTANT(kAsyncUidCntr);

src/env.cc

@@ -11,6 +11,7 @@
 #endif
 
 #include <stdio.h>
+#include <algorithm>
 
 namespace node {
 
@@ -178,12 +179,36 @@ void Environment::AtExit(void (*cb)(void* arg), void* arg) {
 }
 
 void Environment::AddPromiseHook(promise_hook_func fn, void* arg) {
+  auto it = std::find_if(
+      promise_hooks_.begin(), promise_hooks_.end(),
+      [&](const PromiseHookCallback& hook) {
+        return hook.cb_ == fn && hook.arg_ == arg;
+      });
+  CHECK_EQ(it, promise_hooks_.end());
   promise_hooks_.push_back(PromiseHookCallback{fn, arg});
+
   if (promise_hooks_.size() == 1) {
     isolate_->SetPromiseHook(EnvPromiseHook);
   }
 }
 
+bool Environment::RemovePromiseHook(promise_hook_func fn, void* arg) {
+  auto it = std::find_if(
+      promise_hooks_.begin(), promise_hooks_.end(),
+      [&](const PromiseHookCallback& hook) {
+        return hook.cb_ == fn && hook.arg_ == arg;
+      });
+
+  if (it == promise_hooks_.end()) return false;
+
+  promise_hooks_.erase(it);
+  if (promise_hooks_.empty()) {
+    isolate_->SetPromiseHook(nullptr);
+  }
+
+  return true;
+}
+
 void Environment::EnvPromiseHook(v8::PromiseHookType type,
                                  v8::Local<v8::Promise> promise,
                                  v8::Local<v8::Value> parent) {

src/env.h

@@ -357,6 +357,7 @@ class Environment {
       kBefore,
       kAfter,
       kDestroy,
+      kTotals,
       kFieldsCount,
     };
 
@@ -667,6 +668,7 @@ class Environment {
   static const int kContextEmbedderDataIndex = NODE_CONTEXT_EMBEDDER_DATA_INDEX;
 
   void AddPromiseHook(promise_hook_func fn, void* arg);
+  bool RemovePromiseHook(promise_hook_func fn, void* arg);
 
  private:
   inline void ThrowError(v8::Local<v8::Value> (*fun)(v8::Local<v8::String>),

test/addons/async-hooks-promise/binding.cc

@@ -0,0 +1,43 @@
+#include <node.h>
+#include <v8.h>
+
+namespace {
+
+using v8::FunctionCallbackInfo;
+using v8::Isolate;
+using v8::Local;
+using v8::NewStringType;
+using v8::Object;
+using v8::Promise;
+using v8::String;
+using v8::Value;
+
+static void ThrowError(Isolate* isolate, const char* err_msg) {
+  Local<String> str = String::NewFromOneByte(
+      isolate,
+      reinterpret_cast<const uint8_t*>(err_msg),
+      NewStringType::kNormal).ToLocalChecked();
+  isolate->ThrowException(str);
+}
+
+static void GetPromiseField(const FunctionCallbackInfo<Value>& args) {
+  auto isolate = args.GetIsolate();
+
+  if (!args[0]->IsPromise())
+    return ThrowError(isolate, "arg is not an Promise");
+
+  auto p = args[0].As<Promise>();
+  if (p->InternalFieldCount() < 1)
+    return ThrowError(isolate, "Promise has no internal field");
+
+  auto l = p->GetInternalField(0);
+  args.GetReturnValue().Set(l);
+}
+
+inline void Initialize(v8::Local<v8::Object> binding) {
+  NODE_SET_METHOD(binding, "getPromiseField", GetPromiseField);
+}
+
+NODE_MODULE(binding, Initialize)
+
+}  // anonymous namespace

test/addons/async-hooks-promise/binding.gyp

@@ -0,0 +1,9 @@
+{
+  'targets': [
+    {
+      'target_name': 'binding',
+      'defines': [ 'V8_DEPRECATION_WARNINGS=1' ],
+      'sources': [ 'binding.cc' ]
+    }
+  ]
+}

test/addons/async-hooks-promise/test.js

@@ -0,0 +1,43 @@
+'use strict';
+
+const common = require('../../common');
+const assert = require('assert');
+const async_hooks = require('async_hooks');
+const binding = require(`./build/${common.buildType}/binding`);
+
+// Baseline to make sure the internal field isn't being set.
+assert.strictEqual(
+  binding.getPromiseField(Promise.resolve(1)),
+  0,
+  'Promise internal field used despite missing enabled AsyncHook');
+
+const hook0 = async_hooks.createHook({}).enable();
+
+// Check that no PromiseWrap is created when there are no hook callbacks.
+assert.strictEqual(
+  binding.getPromiseField(Promise.resolve(1)),
+  0,
+  'Promise internal field used despite missing enabled AsyncHook');
+
+hook0.disable();
+
+let pwrap = null;
+const hook1 = async_hooks.createHook({
+  init(id, type, tid, resource) {
+    pwrap = resource;
+  }
+}).enable();
+
+// Check that the internal field returns the same PromiseWrap passed to init().
+assert.strictEqual(
+  binding.getPromiseField(Promise.resolve(1)),
+  pwrap,
+  'Unexpected PromiseWrap');
+
+hook1.disable();
+
+// Check that internal fields are no longer being set.
+assert.strictEqual(
+  binding.getPromiseField(Promise.resolve(1)),
+  0,
+  'Promise internal field used despite missing enabled AsyncHook');

test/parallel/test-async-hooks-promise-enable-disable.js

@@ -0,0 +1,47 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const async_hooks = require('async_hooks');
+const EXPECTED_INITS = 2;
+let p_resource = null;
+let p_er = null;
+let p_inits = 0;
+
+// Not useful to place common.mustCall() around 'exit' event b/c it won't be
+// able to check it anway.
+process.on('exit', (code) => {
+  if (code !== 0)
+    return;
+  if (p_er !== null)
+    throw p_er;
+  // Expecint exactly 2 PROMISE types to reach init.
+  assert.strictEqual(p_inits, EXPECTED_INITS);
+});
+
+const mustCallInit = common.mustCall(function init(id, type, tid, resource) {
+  if (type !== 'PROMISE')
+    return;
+  p_inits++;
+  p_resource = resource.promise;
+}, EXPECTED_INITS);
+
+const hook = async_hooks.createHook({
+  init: mustCallInit
+// Enable then disable to test whether disable() actually works.
+}).enable().disable().disable();
+
+new Promise(common.mustCall((res) => {
+  res(42);
+})).then(common.mustCall((val) => {
+  hook.enable().enable();
+  const p = new Promise((res) => res(val));
+  assert.strictEqual(p, p_resource);
+  hook.disable();
+  return p;
+})).then(common.mustCall((val2) => {
+  hook.enable();
+  const p = new Promise((res) => res(val2));
+  hook.disable();
+  return p;
+})).catch((er) => p_er = er);