nodejs
diff --git a/‎src/api/hooks.cc‎
Lines changed: 55 additions & 4 deletions b/‎src/api/hooks.cc‎
Lines changed: 55 additions & 4 deletions
diff --git a/‎test/addons/worker-addon/binding.cc‎
Lines changed: 17 additions & 2 deletions b/‎test/addons/worker-addon/binding.cc‎
Lines changed: 17 additions & 2 deletions
@@ -115,20 +115,71 @@ struct ACHHandle final {
 // this.
 void DeleteACHHandle::operator ()(ACHHandle* handle) const { delete handle; }
 
+// TODO(addaleax): Having this extra set of data structures is far from
+// ideal, but unfortunately the public synchronous cleanup hook API was
+// slightly mis-designed; in particular, RemoveEnvironmentCleanupHook() needs
+// to keep working when the Isolate either has no active context (such as
+// during GC) or that context is associated with another Node.js Environment.
+// We should align this with the asynchronous API, which handles this properly
+// through an explicit reference to the cleanup hook instead of requiring
+// lookups in internal maps.
+struct CleanupHookThunk final {
+  Isolate* isolate;
+  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;
+  }
+};
+struct CleanupHookThunkHash {
+  size_t operator()(const CleanupHookThunk& thunk) const {
+    return std::hash<void*>()(thunk.arg);
+  }
+};
+using CleanupHookRegistry =
+    std::unordered_set<CleanupHookThunk, CleanupHookThunkHash>;
+static ExclusiveAccess<CleanupHookRegistry> cleanup_hook_registry;
+
+static void CleanupHookThunkRun(void* arg) {
+  const CleanupHookThunk* thunk = static_cast<CleanupHookThunk*>(arg);
+  thunk->fun(thunk->arg);
+  RemoveEnvironmentCleanupHook(thunk->isolate, thunk->fun, thunk->arg);
+}
+
 void AddEnvironmentCleanupHook(Isolate* isolate,
                                CleanupHook fun,
                                void* arg) {
   Environment* env = Environment::GetCurrent(isolate);
   CHECK_NOT_NULL(env);
-  env->AddCleanupHook(fun, arg);
+  void* wrapped_arg;
+  {
+    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);
+  }
+  env->AddCleanupHook(CleanupHookThunkRun, wrapped_arg);
 }
 
 void RemoveEnvironmentCleanupHook(Isolate* isolate,
                                   CleanupHook fun,
                                   void* arg) {
-  Environment* env = Environment::GetCurrent(isolate);
-  CHECK_NOT_NULL(env);
-  env->RemoveCleanupHook(fun, arg);
+  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);
+  }
+  thunk.env->RemoveCleanupHook(CleanupHookThunkRun, wrapped_arg);
 }
 
 static void FinishAsyncCleanupHook(void* arg) {
 
@@ -54,8 +54,23 @@ void Initialize(Local<Object> exports,
   Isolate* isolate = Isolate::GetCurrent();
   node::AddEnvironmentCleanupHook(
       isolate, Cleanup, const_cast<void*>(static_cast<const void*>("cleanup")));
-  node::AddEnvironmentCleanupHook(isolate, Dummy, nullptr);
-  node::RemoveEnvironmentCleanupHook(isolate, Dummy, nullptr);
+
+  // Test that adding and removing a cleanup hook works as expected
+  {
+    node::AddEnvironmentCleanupHook(isolate, Dummy, nullptr);
+    node::RemoveEnvironmentCleanupHook(isolate, Dummy, nullptr);
+  }
+
+  // Test that adding and removing a cleanup hook also works if there
+  // is no active context during removal
+  {
+    node::AddEnvironmentCleanupHook(isolate, Dummy, nullptr);
+    {
+      context->Exit();
+      node::RemoveEnvironmentCleanupHook(isolate, Dummy, nullptr);
+      context->Enter();
+    }
+  }
 
   if (getenv("addExtraItemToEventLoop") != nullptr) {
     // Add an item to the event loop that we do not clean up in order to make