src: add environment cleanup hooks

Original PR: ayojs/ayo#82

> This adds pairs of methods to the `Environment` class and to public APIs
> which can add and remove cleanup handlers.

> Unlike `AtExit`, this API targets addon developers rather than
> embedders, giving them (and Node’s internals) the ability to register
> per-`Environment` cleanup work.

> PR-URL: ayojs/ayo#82
> Reviewed-By: Stephen Belanger <[email protected]>

src/env-inl.h

@@ -599,6 +599,26 @@ inline void Environment::SetTemplateMethod(v8::Local<v8::FunctionTemplate> that,
   t->SetClassName(name_string);  // NODE_SET_METHOD() compatibility.
 }
 
+void Environment::AddCleanupHook(void (*fn)(void*), void* arg) {
+  cleanup_hooks_[arg].push_back(
+      CleanupHookCallback { fn, arg, cleanup_hook_counter_++ });
+}
+
+void Environment::RemoveCleanupHook(void (*fn)(void*), void* arg) {
+  auto map_it = cleanup_hooks_.find(arg);
+  if (map_it == cleanup_hooks_.end())
+    return;
+
+  for (auto it = map_it->second.begin(); it != map_it->second.end(); ++it) {
+    if (it->fun_ == fn && it->arg_ == arg) {
+      map_it->second.erase(it);
+      if (map_it->second.empty())
+        cleanup_hooks_.erase(arg);
+      return;
+    }
+  }
+}
+
 #define VP(PropertyName, StringValue) V(v8::Private, PropertyName)
 #define VS(PropertyName, StringValue) V(v8::String, PropertyName)
 #define V(TypeName, PropertyName)                                             \

src/env.cc

@@ -166,6 +166,26 @@ void Environment::PrintSyncTrace() const {
   fflush(stderr);
 }
 
+void Environment::RunCleanup() {
+  while (!cleanup_hooks_.empty()) {
+    std::vector<CleanupHookCallback> callbacks;
+    // Concatenate all vectors in cleanup_hooks_
+    for (const auto& pair : cleanup_hooks_)
+      callbacks.insert(callbacks.end(), pair.second.begin(), pair.second.end());
+    cleanup_hooks_.clear();
+    std::sort(callbacks.begin(), callbacks.end(),
+              [](const CleanupHookCallback& a, const CleanupHookCallback& b) {
+      // Sort in descending order so that the last-inserted callbacks get run
+      // first.
+      return a.insertion_order_counter_ > b.insertion_order_counter_;
+    });
+
+    for (const CleanupHookCallback& cb : callbacks) {
+      cb.fun_(cb.arg_);
+    }
+  }
+}
+
 void Environment::RunAtExitCallbacks() {
   for (AtExitCallback at_exit : at_exit_functions_) {
     at_exit.cb_(at_exit.arg_);

src/env.h

@@ -680,6 +680,10 @@ class Environment {
   bool RemovePromiseHook(promise_hook_func fn, void* arg);
   bool EmitNapiWarning();
 
+  inline void AddCleanupHook(void (*fn)(void*), void* arg);
+  inline void RemoveCleanupHook(void (*fn)(void*), void* arg);
+  void RunCleanup();
+
  private:
   inline void ThrowError(v8::Local<v8::Value> (*fun)(v8::Local<v8::String>),
                          const char* errmsg);
@@ -738,6 +742,15 @@ class Environment {
   };
   std::vector<PromiseHookCallback> promise_hooks_;
 
+  struct CleanupHookCallback {
+    void (*fun_)(void*);
+    void* arg_;
+    int64_t insertion_order_counter_;
+  };
+
+  std::unordered_map<void*, std::vector<CleanupHookCallback>> cleanup_hooks_;
+  int64_t cleanup_hook_counter_ = 0;
+
   static void EnvPromiseHook(v8::PromiseHookType type,
                              v8::Local<v8::Promise> promise,
                              v8::Local<v8::Value> parent);

src/node.cc

@@ -1339,6 +1339,22 @@ void AddPromiseHook(v8::Isolate* isolate, promise_hook_func fn, void* arg) {
   env->AddPromiseHook(fn, arg);
 }
 
+void AddEnvironmentCleanupHook(v8::Isolate* isolate,
+                               void (*fun)(void* arg),
+                               void* arg) {
+  Environment* env = Environment::GetCurrent(isolate);
+  env->AddCleanupHook(fun, arg);
+}
+
+
+void RemoveEnvironmentCleanupHook(v8::Isolate* isolate,
+                                  void (*fun)(void* arg),
+                                  void* arg) {
+  Environment* env = Environment::GetCurrent(isolate);
+  env->RemoveCleanupHook(fun, arg);
+}
+
+
 CallbackScope::CallbackScope(Isolate* isolate,
                              Local<Object> object,
                              async_context asyncContext)
@@ -4747,6 +4763,8 @@ inline int Start(Isolate* isolate, IsolateData* isolate_data,
   env.set_trace_sync_io(false);
 
   const int exit_code = EmitExit(&env);
+
+  env.RunCleanup();
   RunAtExit(&env);
   uv_key_delete(&thread_local_env);
 

src/node.h

@@ -546,6 +546,19 @@ NODE_EXTERN void AddPromiseHook(v8::Isolate* isolate,
                                 promise_hook_func fn,
                                 void* arg);
 
+/* This is a lot like node::AtExit, except that the hooks added via this
+ * function are run before the AtExit ones and will always be registered
+ * for the current Environment instance.
+ * These functions are safe to use in an addon supporting multiple
+ * threads/isolates. */
+NODE_EXTERN void AddEnvironmentCleanupHook(v8::Isolate* isolate,
+                                           void (*fun)(void* arg),
+                                           void* arg);
+
+NODE_EXTERN void RemoveEnvironmentCleanupHook(v8::Isolate* isolate,
+                                              void (*fun)(void* arg),
+                                              void* arg);
+
 /* Returns the id of the current execution context. If the return value is
  * zero then no execution has been set. This will happen if the user handles
  * I/O from native code. */

src/node_api.cc

@@ -868,6 +868,28 @@ void napi_module_register(napi_module* mod) {
   node::node_module_register(nm);
 }
 
+napi_status napi_add_env_cleanup_hook(napi_env env,
+                                      void (*fun)(void* arg),
+                                      void* arg) {
+  CHECK_ENV(env);
+  CHECK_ARG(env, fun);
+
+  node::AddEnvironmentCleanupHook(env->isolate, fun, arg);
+
+  return napi_ok;
+}
+
+napi_status napi_remove_env_cleanup_hook(napi_env env,
+                                         void (*fun)(void* arg),
+                                         void* arg) {
+  CHECK_ENV(env);
+  CHECK_ARG(env, fun);
+
+  node::RemoveEnvironmentCleanupHook(env->isolate, fun, arg);
+
+  return napi_ok;
+}
+
 // Warning: Keep in-sync with napi_status enum
 const char* error_messages[] = {nullptr,
                                 "Invalid argument",

src/node_api.h

@@ -106,6 +106,13 @@ EXTERN_C_START
 
 NAPI_EXTERN void napi_module_register(napi_module* mod);
 
+NAPI_EXTERN napi_status napi_add_env_cleanup_hook(napi_env env,
+                                                  void (*fun)(void* arg),
+                                                  void* arg);
+NAPI_EXTERN napi_status napi_remove_env_cleanup_hook(napi_env env,
+                                                     void (*fun)(void* arg),
+                                                     void* arg);
+
 NAPI_EXTERN napi_status
 napi_get_last_error_info(napi_env env,
                          const napi_extended_error_info** result);