src: add environment cleanup hooks

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.

We may want to replace `AtExit` with this API at some point.

Many thanks for Stephen Belanger for reviewing the original version of
this commit in the Ayo.js project.

Refs: ayojs/ayo#82
PR-URL: nodejs#19377
Reviewed-By: Ben Noordhuis <[email protected]>
Reviewed-By: James M Snell <[email protected]>

Author: addaleax

doc/api/n-api.md

@@ -901,6 +901,58 @@ If still valid, this API returns the `napi_value` representing the
 JavaScript Object associated with the `napi_ref`. Otherwise, result
 will be NULL.
 
+### Cleanup on exit of the current Node.js instance
+
+While a Node.js process typically releases all its resources when exiting,
+embedders of Node.js, or future Worker support, may require addons to register
+clean-up hooks that will be run once the current Node.js instance exits.
+
+N-API provides functions for registering and un-registering such callbacks.
+When those callbacks are run, all resources that are being held by the addon
+should be freed up.
+
+#### napi_add_env_cleanup_hook
+<!-- YAML
+added: REPLACEME
+-->
+```C
+NODE_EXTERN napi_status napi_add_env_cleanup_hook(napi_env env,
+                                                  void (*fun)(void* arg),
+                                                  void* arg);
+```
+
+Registers `fun` as a function to be run with the `arg` parameter once the
+current Node.js environment exits.
+
+A function can safely be specified multiple times with different
+`arg` values. In that case, it will be called multiple times as well.
+Providing the same `fun` and `arg` values multiple times is not allowed
+and will lead the process to abort.
+
+The hooks will be called in reverse order, i.e. the most recently added one
+will be called first.
+
+Removing this hook can be done by using `napi_remove_env_cleanup_hook`.
+Typically, that happens when the resource for which this hook was added
+is being torn down anyway.
+
+#### napi_remove_env_cleanup_hook
+<!-- YAML
+added: REPLACEME
+-->
+```C
+NAPI_EXTERN napi_status napi_remove_env_cleanup_hook(napi_env env,
+                                                     void (*fun)(void* arg),
+                                                     void* arg);
+```
+
+Unregisters `fun` as a function to be run with the `arg` parameter once the
+current Node.js environment exits. Both the argument and the function value
+need to be exact matches.
+
+The function must have originally been registered
+with `napi_add_env_cleanup_hook`, otherwise the process will abort.
+
 ## Module registration
 N-API modules are registered in a manner similar to other modules
 except that instead of using the `NODE_MODULE` macro the following

src/env-inl.h

@@ -565,6 +565,29 @@ inline v8::Local<v8::Object> Environment::NewInternalFieldObject() {
   return m_obj.ToLocalChecked();
 }
 
+void Environment::AddCleanupHook(void (*fn)(void*), void* arg) {
+  auto insertion_info = cleanup_hooks_.emplace(CleanupHookCallback {
+    fn, arg, cleanup_hook_counter_++
+  });
+  // Make sure there was no existing element with these values.
+  CHECK_EQ(insertion_info.second, true);
+}
+
+void Environment::RemoveCleanupHook(void (*fn)(void*), void* arg) {
+  CleanupHookCallback search { fn, arg, 0 };
+  cleanup_hooks_.erase(search);
+}
+
+size_t Environment::CleanupHookCallback::Hash::operator()(
+    const CleanupHookCallback& cb) const {
+  return std::hash<void*>()(cb.arg_);
+}
+
+bool Environment::CleanupHookCallback::Equal::operator()(
+    const CleanupHookCallback& a, const CleanupHookCallback& b) const {
+  return a.fn_ == b.fn_ && a.arg_ == b.arg_;
+}
+
 #define VP(PropertyName, StringValue) V(v8::Private, PropertyName, StringValue)
 #define VS(PropertyName, StringValue) V(v8::String, PropertyName, StringValue)
 #define V(TypeName, PropertyName, StringValue)                                \

src/env.cc

@@ -8,6 +8,7 @@
 #endif
 
 #include <stdio.h>
+#include <algorithm>
 
 namespace node {
 
@@ -61,4 +62,34 @@ void Environment::PrintSyncTrace() const {
   fflush(stderr);
 }
 
+void Environment::RunCleanup() {
+  while (!cleanup_hooks_.empty()) {
+    // Copy into a vector, since we can't sort an unordered_set in-place.
+    std::vector<CleanupHookCallback> callbacks(
+        cleanup_hooks_.begin(), cleanup_hooks_.end());
+    // We can't erase the copied elements from `cleanup_hooks_` yet, because we
+    // need to be able to check whether they were un-scheduled by another hook.
+
+    std::sort(callbacks.begin(), callbacks.end(),
+              [](const CleanupHookCallback& a, const CleanupHookCallback& b) {
+      // Sort in descending order so that the most recently inserted callbacks
+      // are run first.
+      return a.insertion_order_counter_ > b.insertion_order_counter_;
+    });
+
+    for (const CleanupHookCallback& cb : callbacks) {
+      if (cleanup_hooks_.count(cb) == 0) {
+        // This hook was removed from the `cleanup_hooks_` set during another
+        // hook that was run earlier. Nothing to do here.
+        continue;
+      }
+
+      cb.fn_(cb.arg_);
+      cleanup_hooks_.erase(cb);
+      // TODO(addaleax): Not calling CleanupHandles() here because it hangs in a
+      // busy loop.
+    }
+  }
+}
+
 }  // namespace node

src/env.h

@@ -17,6 +17,7 @@
 
 #include <stdint.h>
 #include <vector>
+#include <unordered_set>
 
 // Caveat emptor: we're going slightly crazy with macros here but the end
 // hopefully justifies the means. We have a lot of per-context properties
@@ -537,6 +538,10 @@ class Environment {
 
   static const int kContextEmbedderDataIndex = NODE_CONTEXT_EMBEDDER_DATA_INDEX;
 
+  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);
@@ -633,6 +638,32 @@ class Environment {
     DISALLOW_COPY_AND_ASSIGN(IsolateData);
   };
 
+  struct CleanupHookCallback {
+    void (*fn_)(void*);
+    void* arg_;
+
+    // We keep track of the insertion order for these objects, so that we can
+    // call the callbacks in reverse order when we are cleaning up.
+    uint64_t insertion_order_counter_;
+
+    // Only hashes `arg_`, since that is usually enough to identify the hook.
+    struct Hash {
+      inline size_t operator()(const CleanupHookCallback& cb) const;
+    };
+
+    // Compares by `fn_` and `arg_` being equal.
+    struct Equal {
+      inline bool operator()(const CleanupHookCallback& a,
+                             const CleanupHookCallback& b) const;
+    };
+  };
+
+  // Use an unordered_set, so that we have efficient insertion and removal.
+  std::unordered_set<CleanupHookCallback,
+                     CleanupHookCallback::Hash,
+                     CleanupHookCallback::Equal> cleanup_hooks_;
+  uint64_t cleanup_hook_counter_ = 0;
+
   DISALLOW_COPY_AND_ASSIGN(Environment);
 };
 

src/node.cc

@@ -1241,6 +1241,20 @@ void SetupPromises(const FunctionCallbackInfo<Value>& args) {
       FIXED_ONE_BYTE_STRING(isolate, "_setupPromises")).FromJust();
 }
 
+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);
+}
 
 Local<Value> MakeCallback(Environment* env,
                           Local<Value> recv,
@@ -3616,6 +3630,7 @@ void LoadEnvironment(Environment* env) {
 
 void FreeEnvironment(Environment* env) {
   CHECK_NE(env, nullptr);
+  env->RunCleanup();
   env->Dispose();
 }
 
@@ -4913,6 +4928,8 @@ static void StartNodeInstance(void* arg) {
     int exit_code = EmitExit(env);
     if (instance_data->is_main())
       instance_data->set_exit_code(exit_code);
+
+    env->RunCleanup();
     RunAtExit(env);
 
     WaitForInspectorDisconnect(env);

src/node.h

@@ -495,6 +495,19 @@ extern "C" NODE_EXTERN void node_module_register(void* mod);
  */
 NODE_EXTERN void AtExit(void (*cb)(void* arg), void* arg = 0);
 
+/* 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);
+
 }  // namespace node
 
 #endif  // SRC_NODE_H_

src/node_api.cc

@@ -922,6 +922,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
 static
 const char* error_messages[] = {nullptr,

src/node_api.h

@@ -99,6 +99,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);

test/addons-napi/test_cleanup_hook/binding.cc

@@ -0,0 +1,24 @@
+#include "node_api.h"
+#include "uv.h"
+#include "../common.h"
+
+namespace {
+
+void cleanup(void* arg) {
+  printf("cleanup(%d)\n", *static_cast<int*>(arg));
+}
+
+int secret = 42;
+int wrong_secret = 17;
+
+napi_value Init(napi_env env, napi_value exports) {
+  napi_add_env_cleanup_hook(env, cleanup, &wrong_secret);
+  napi_add_env_cleanup_hook(env, cleanup, &secret);
+  napi_remove_env_cleanup_hook(env, cleanup, &wrong_secret);
+
+  return nullptr;
+}
+
+}  // anonymous namespace
+
+NAPI_MODULE(NODE_GYP_MODULE_NAME, Init)

test/addons-napi/test_cleanup_hook/binding.gyp

@@ -0,0 +1,9 @@
+{
+  'targets': [
+    {
+      'target_name': 'binding',
+      'defines': [ 'V8_DEPRECATION_WARNINGS=1' ],
+      'sources': [ 'binding.cc' ]
+    }
+  ]
+}