From 96279e83e73375fcbf4833e1810ff811f36416cc Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna@addaleax.net>
Date: Tue, 6 Jun 2017 13:22:48 -0600
Subject: [PATCH] async_wrap: expose enable/disablePromiseHook API

Allow node::PromiseHook (src/async-wrap.cc) to be enabled/disabled from
the JavaScript API.

PR-URL: https://github.com/nodejs/node/pull/13509
Reviewed-By: Andreas Madsen <amwebdk@gmail.com>
Reviewed-By: Anna Henningsen <anna@addaleax.net>
---
 src/async-wrap.cc | 14 ++++++++++++++
 src/env.cc        | 25 +++++++++++++++++++++++++
 src/env.h         |  1 +
 3 files changed, 40 insertions(+)

diff --git a/src/async-wrap.cc b/src/async-wrap.cc
index 2dcfdc799720ed..c4777ac827a685 100644
--- a/src/async-wrap.cc
+++ b/src/async-wrap.cc
@@ -422,6 +422,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 +490,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);
diff --git a/src/env.cc b/src/env.cc
index 0bf4b573aafd69..c97868a5882b41 100644
--- a/src/env.cc
+++ b/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) {
diff --git a/src/env.h b/src/env.h
index 2eacf68ef5e93b..2f62663be6ab33 100644
--- a/src/env.h
+++ b/src/env.h
@@ -654,6 +654,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>),