async_wrap: use kTotals to enable PromiseHook

Keep a total of enabled hook callbacks in kTotals. This value is used to track whether node::PromiseHook (src/async-wrap.cc) should be enabled or disabled. Don't enable node::PromiseHook, using enablePromiseHook(), until a hook has been added. Then, using disablePromiseHook(), disable node::PromiseHook when all hooks have been disabled. Need to use a native test in order to check the internal field of the Promise and check for a PromiseWrap. PR-URL: #13509
nodejs · trevnorris · Jun 6, 2017 · Jun 6, 2017 · Jun 6, 2017 · Jun 10, 2017
commit 5fa59ccd7e42a4bfe8d8e3bc37d0e67bd8575abd
diff --git a/lib/async_hooks.js b/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;
 
@@ -50,7 +50,9 @@ const after_symbol = Symbol('after');
 const destroy_symbol = Symbol('destroy');
 
 // Setup the callbacks that node::AsyncWrap will call when there are hooks to
-// process. They use the same functions as the JS embedder API.
+// 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,
@@ -103,14 +105,21 @@ class AsyncHook {
     if (hooks_array.includes(this))
       return this;
 
+    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;
   }
 
@@ -121,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;
   }
 }

diff --git a/src/async-wrap.cc b/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
 
   {
@@ -542,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);

diff --git a/src/env.h b/src/env.h
@@ -357,6 +357,7 @@ class Environment {
       kBefore,
       kAfter,
       kDestroy,
+      kTotals,
       kFieldsCount,
     };
 

diff --git a/test/addons/async-hooks-promise/binding.cc b/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
diff --git a/test/addons/async-hooks-promise/binding.gyp b/test/addons/async-hooks-promise/binding.gyp
@@ -0,0 +1,9 @@
+{
+  'targets': [
+    {
+      'target_name': 'binding',
+      'defines': [ 'V8_DEPRECATION_WARNINGS=1' ],
+      'sources': [ 'binding.cc' ]
+    }
+  ]
+}
diff --git a/test/addons/async-hooks-promise/test.js b/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');
diff --git a/test/parallel/test-async-hooks-promise-enable-disable.js b/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);