src: let Environments on one isolate share a cleanup hook

src/api/hooks.cc

@@ -3,6 +3,9 @@
 #include "node_process-inl.h"
 #include "async_wrap.h"
 
+#include <algorithm>
+#include <unordered_map>
+
 namespace node {
 
 using v8::Context;
@@ -128,63 +131,78 @@ struct CleanupHookThunk final {
   Environment* env;
   CleanupHook fun;
   void* arg;
-
-  bool operator==(const CleanupHookThunk& other) const {
-    // `env` is intentionally not part of this comparison
-    return isolate == other.isolate && fun == other.fun && arg == other.arg;
-  }
+  bool running = false;
 };
-struct CleanupHookThunkHash {
-  size_t operator()(const CleanupHookThunk& thunk) const {
-    return std::hash<void*>()(thunk.arg);
-  }
-};
-using CleanupHookRegistry =
-    std::unordered_set<CleanupHookThunk, CleanupHookThunkHash>;
+// Keyed on `arg`. The same hook may be registered once per Environment, and
+// several Environments can share an Isolate.
+using CleanupHookRegistry = std::unordered_multimap<void*, CleanupHookThunk>;
 static ExclusiveAccess<CleanupHookRegistry> cleanup_hook_registry;
 
 static void CleanupHookThunkRun(void* arg) {
-  const CleanupHookThunk* thunk = static_cast<CleanupHookThunk*>(arg);
-  // `thunk->fun` may itself remove and free this CleanupHookThunk (e.g. via
-  // ~ObjectWrap(), which calls RemoveEnvironmentCleanupHook()), so cache the
-  // fields we still need before invoking it rather than reading them from
-  // `thunk` afterwards.
-  Isolate* isolate = thunk->isolate;
-  CleanupHook fun = thunk->fun;
-  void* fun_arg = thunk->arg;
-  fun(fun_arg);
-  RemoveEnvironmentCleanupHook(isolate, fun, fun_arg);
+  CleanupHookThunk* thunk = static_cast<CleanupHookThunk*>(arg);
+  {
+    ExclusiveAccess<CleanupHookRegistry>::Scoped registry(
+        &cleanup_hook_registry);
+    thunk->running = true;
+  }
+  thunk->fun(thunk->arg);
+  ExclusiveAccess<CleanupHookRegistry>::Scoped registry(&cleanup_hook_registry);
+  auto [begin, end] = registry->equal_range(thunk->arg);
+  auto self = std::find_if(
+      begin, end, [&](const auto& entry) { return &entry.second == thunk; });
+  CHECK(self != end);
+  registry->erase(self);
 }
 
 void AddEnvironmentCleanupHook(Isolate* isolate,
                                CleanupHook fun,
                                void* arg) {
   Environment* env = Environment::GetCurrent(isolate);
   CHECK_NOT_NULL(env);
-  void* wrapped_arg;
+  CleanupHookThunk* thunk;
   {
     ExclusiveAccess<CleanupHookRegistry>::Scoped registry(
         &cleanup_hook_registry);
-    auto result = registry->insert({isolate, env, fun, arg});
-    CHECK(result.second);
-    wrapped_arg = const_cast<CleanupHookThunk*>(&*result.first);
+    auto [begin, end] = registry->equal_range(arg);
+    // Adding the same hook twice to one Environment is documented to abort;
+    // a running hook may register itself again.
+    CHECK(std::none_of(begin, end, [&](const auto& entry) {
+      return entry.second.env == env && entry.second.fun == fun &&
+             !entry.second.running;
+    }));
+    thunk = &registry->emplace(arg, CleanupHookThunk{isolate, env, fun, arg})
+                 ->second;
   }
-  env->AddCleanupHook(CleanupHookThunkRun, wrapped_arg);
+  env->AddCleanupHook(CleanupHookThunkRun, thunk);
 }
 
 void RemoveEnvironmentCleanupHook(Isolate* isolate,
                                   CleanupHook fun,
                                   void* arg) {
+  // Prefer the current Environment's registration and otherwise take any
+  // match: there may be no current context (GC, addon threads) or it may
+  // belong to another Environment on the same isolate.
+  Environment* current =
+      isolate != nullptr && isolate == Isolate::TryGetCurrent()
+          ? Environment::GetCurrent(isolate)
+          : nullptr;
   CleanupHookThunk thunk;
   void* wrapped_arg;
   {
     ExclusiveAccess<CleanupHookRegistry>::Scoped registry(
         &cleanup_hook_registry);
-    auto result = registry->find({isolate, nullptr, fun, arg});
-    if (result == registry->end()) return;
-    wrapped_arg = const_cast<CleanupHookThunk*>(&*result);
-    thunk = *result;
-    registry->erase(result);
+    auto [begin, end] = registry->equal_range(arg);
+    auto found = end;
+    for (auto it = begin; it != end; ++it) {
+      if (it->second.isolate != isolate || it->second.fun != fun) continue;
+      if (found == end || it->second.env == current) found = it;
+      if (it->second.env == current) break;
+    }
+    // A running hook is removing itself; CleanupHookThunkRun() cleans up.
+    if (found == end || found->second.running) return;
+    wrapped_arg = &found->second;
+    thunk = found->second;
+    registry->erase(found);
   }
   thunk.env->RemoveCleanupHook(CleanupHookThunkRun, wrapped_arg);
 }

test/cctest/test_environment.cc

@@ -355,6 +355,66 @@ TEST_F(EnvironmentTest, MultipleEnvironmentsPerIsolate) {
   EXPECT_TRUE(called_cb_2);
 }
 
+static int cleanup_hook_runs = 0;
+static void CountingCleanupHook(void* arg) {
+  cleanup_hook_runs++;
+}
+
+TEST_F(EnvironmentTest, SameCleanupHookInTwoEnvironmentsOnOneIsolate) {
+  const v8::HandleScope handle_scope(isolate_);
+  const Argv argv;
+  cleanup_hook_runs = 0;
+  {
+    Env env1{handle_scope, argv};
+    {
+      Env env2{handle_scope, argv, node::EnvironmentFlags::kNoCreateInspector};
+      {
+        v8::Context::Scope context_scope(env1.context());
+        node::AddEnvironmentCleanupHook(isolate_, CountingCleanupHook, nullptr);
+      }
+      node::AddEnvironmentCleanupHook(isolate_, CountingCleanupHook, nullptr);
+    }
+    EXPECT_EQ(cleanup_hook_runs, 1);
+  }
+  EXPECT_EQ(cleanup_hook_runs, 2);
+}
+
+TEST_F(EnvironmentTest, RemoveCleanupHookOfOtherEnvironmentOnSameIsolate) {
+  const v8::HandleScope handle_scope(isolate_);
+  const Argv argv;
+  cleanup_hook_runs = 0;
+  int arg;
+  {
+    Env env1{handle_scope, argv};
+    node::AddEnvironmentCleanupHook(isolate_, CountingCleanupHook, &arg);
+    Env env2{handle_scope, argv, node::EnvironmentFlags::kNoCreateInspector};
+    // env2's context is current; the hook belongs to env1.
+    node::RemoveEnvironmentCleanupHook(isolate_, CountingCleanupHook, &arg);
+  }
+  EXPECT_EQ(cleanup_hook_runs, 0);
+}
+
+struct SelfRemovingHook {
+  v8::Isolate* isolate;
+  bool ran = false;
+  static void Run(void* arg) {
+    SelfRemovingHook* self = static_cast<SelfRemovingHook*>(arg);
+    self->ran = true;
+    node::RemoveEnvironmentCleanupHook(self->isolate, Run, arg);
+  }
+};
+
+TEST_F(EnvironmentTest, CleanupHookRemovesItselfWhileRunning) {
+  const v8::HandleScope handle_scope(isolate_);
+  const Argv argv;
+  SelfRemovingHook hook{isolate_};
+  {
+    Env env{handle_scope, argv};
+    node::AddEnvironmentCleanupHook(isolate_, SelfRemovingHook::Run, &hook);
+  }
+  EXPECT_TRUE(hook.ran);
+}
+
 TEST_F(EnvironmentTest, NoEnvironmentSanity) {
   const v8::HandleScope handle_scope(isolate_);
   v8::Local<v8::Context> context = v8::Context::New(isolate_);