From a6304d97aacb1f9b6a62ca6ee1b323fd71902889 Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna@addaleax.net>
Date: Fri, 2 Jun 2017 22:41:32 +0200
Subject: [PATCH] async_hooks: tear down hooks if none are in use

Basically, reset everything into its original state once all
async hooks have been disabled as closely as possible.

This has been requested by Andreas, and I agree that it makes sense.
---
 lib/async_hooks.js | 12 ++++++++++++
 src/async-wrap.cc  | 22 +++++++++++++++++++---
 src/env.cc         | 18 ++++++++++++++++++
 src/env.h          |  1 +
 4 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/lib/async_hooks.js b/lib/async_hooks.js
index d5d5407b1de19b..db3ebec83d85d3 100644
--- a/lib/async_hooks.js
+++ b/lib/async_hooks.js
@@ -131,6 +131,18 @@ class AsyncHook {
     hook_fields[kAfter] -= +!!this[after_symbol];
     hook_fields[kDestroy] -= +!!this[destroy_symbol];
     hooks_array.splice(index, 1);
+
+    if (hooks_array.length === 0) {
+      if (hook_fields[kInit] + hook_fields[kBefore] +
+          hook_fields[kAfter] + hook_fields[kDestroy] !== 0) {
+        const { inspect } = require('util');
+        throw new Error(`Invalid async_hooks state: ${inspect(hook_fields)}`);
+      }
+
+      async_wrap.tearDownHooks();
+      setupHooksCalled = false;
+    }
+
     return this;
   }
 }
diff --git a/src/async-wrap.cc b/src/async-wrap.cc
index 0ea6c64a8be582..68a653d46329c0 100644
--- a/src/async-wrap.cc
+++ b/src/async-wrap.cc
@@ -291,8 +291,7 @@ class PromiseWrap : public AsyncWrap {
 
 static void PromiseHook(PromiseHookType type, Local<Promise> promise,
                         Local<Value> parent, void* arg) {
-  Local<Context> context = promise->CreationContext();
-  Environment* env = Environment::GetCurrent(context);
+  Environment* env = static_cast<Environment*>(arg);
   PromiseWrap* wrap = Unwrap<PromiseWrap>(promise);
   if (type == PromiseHookType::kInit || wrap == nullptr) {
     bool silent = type != PromiseHookType::kInit;
@@ -334,8 +333,24 @@ 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
+
+  env->AddPromiseHook(PromiseHook, static_cast<void*>(env));
+}
+
+
+static void TearDownHooks(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+
+  CHECK(!env->async_hooks_init_function().IsEmpty());
+
+  env->set_async_hooks_init_function(Local<Function>());
+  env->set_async_hooks_before_function(Local<Function>());
+  env->set_async_hooks_after_function(Local<Function>());
+  env->set_async_hooks_destroy_function(Local<Function>());
+
+  bool removed = env->RemovePromiseHook(PromiseHook, static_cast<void*>(env));
+  CHECK(removed);
 }
 
 
@@ -391,6 +406,7 @@ void AsyncWrap::Initialize(Local<Object> target,
   HandleScope scope(isolate);
 
   env->SetMethod(target, "setupHooks", SetupHooks);
+  env->SetMethod(target, "tearDownHooks", TearDownHooks);
   env->SetMethod(target, "pushAsyncIds", PushAsyncIds);
   env->SetMethod(target, "popAsyncIds", PopAsyncIds);
   env->SetMethod(target, "clearIdStack", ClearIdStack);
diff --git a/src/env.cc b/src/env.cc
index 034625b375446c..8c90f67512a1f9 100644
--- a/src/env.cc
+++ b/src/env.cc
@@ -11,6 +11,7 @@
 #endif
 
 #include <stdio.h>
+#include <algorithm>
 
 namespace node {
 
@@ -195,6 +196,23 @@ void Environment::AddPromiseHook(promise_hook_func fn, void* arg) {
   }
 }
 
+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 c8c8232cc07fd4..cd7222b9d05365 100644
--- a/src/env.h
+++ b/src/env.h
@@ -653,6 +653,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>),