src: make AtExit callback's per Environment

This commit attempts to address one of the TODOs in
https://github.com/nodejs/node/issues/4641 regarding making the
AtExit callback's per environment, instead of the current global.

bnoordhuis provided a few options for solving this, and one was to
use a thread-local which is what this commit attempts to do.

PR-URL: https://github.com/nodejs/node/pull/9163
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: bnoordhuis - Ben Noordhuis <info@bnoordhuis.nl>

4 files changed, 41 insertions, 14 deletions

diff --git a/src/env.cc b/src/env.cc
index 40f0c9ebd66..56d7e28a03a 100644
--- a/src/env.cc
+++ b/src/env.cc
@@ -153,4 +153,15 @@ void Environment::PrintSyncTrace() const {
   fflush(stderr);
 }
 
+void Environment::RunAtExitCallbacks() {
+  for (AtExitCallback at_exit : at_exit_functions_) {
+    at_exit.cb_(at_exit.arg_);
+  }
+  at_exit_functions_.clear();
+}
+
+void Environment::AtExit(void (*cb)(void* arg), void* arg) {
+  at_exit_functions_.push_back(AtExitCallback{cb, arg});
+}
+
 }  // namespace node
diff --git a/src/env.h b/src/env.h
index 66032c999f8..a719cda2d40 100644
--- a/src/env.h
+++ b/src/env.h
@@ -36,6 +36,7 @@
 #include "uv.h"
 #include "v8.h"
 
+#include <list>
 #include <stdint.h>
 #include <vector>
 
@@ -530,6 +531,9 @@ class Environment {
 
   inline v8::Local<v8::Object> NewInternalFieldObject();
 
+  void AtExit(void (*cb)(void* arg), void* arg);
+  void RunAtExitCallbacks();
+
   // Strings and private symbols are shared across shared contexts
   // The getters simply proxy to the per-isolate primitive.
 #define VP(PropertyName, StringValue) V(v8::Private, PropertyName)
@@ -609,6 +613,12 @@ class Environment {
 
   double* fs_stats_field_array_;
 
+  struct AtExitCallback {
+    void (*cb_)(void* arg);
+    void* arg_;
+  };
+  std::list<AtExitCallback> at_exit_functions_;
+
 #define V(PropertyName, TypeName)                                             \
   v8::Persistent<TypeName> PropertyName ## _;
   ENVIRONMENT_STRONG_PERSISTENT_PROPERTIES(V)
diff --git a/src/node.cc b/src/node.cc
index ba146241a23..c241f734b28 100644
--- a/src/node.cc
+++ b/src/node.cc
@@ -84,7 +84,6 @@
 
 #include <string>
 #include <vector>
-#include <list>
 
 #if defined(NODE_HAVE_I18N_SUPPORT)
 #include <unicode/uvernum.h>
@@ -4255,25 +4254,23 @@ void Init(int* argc,
 }
 
 
-struct AtExitCallback {
-  void (*cb_)(void* arg);
-  void* arg_;
-};
+void RunAtExit(Environment* env) {
+  env->RunAtExitCallbacks();
+}
 
-static std::list<AtExitCallback> at_exit_functions;
 
+static uv_key_t thread_local_env;
 
-// TODO(bnoordhuis) Turn into per-context event.
-void RunAtExit(Environment* env) {
-  for (AtExitCallback at_exit : at_exit_functions) {
-    at_exit.cb_(at_exit.arg_);
-  }
-  at_exit_functions.clear();
+
+void AtExit(void (*cb)(void* arg), void* arg) {
+  auto env = static_cast<Environment*>(uv_key_get(&thread_local_env));
+  AtExit(env, cb, arg);
 }
 
 
-void AtExit(void (*cb)(void* arg), void* arg) {
-  at_exit_functions.push_back(AtExitCallback{cb, arg});
+void AtExit(Environment* env, void (*cb)(void* arg), void* arg) {
+  CHECK_NE(env, nullptr);
+  env->AtExit(cb, arg);
 }
 
 
@@ -4349,6 +4346,8 @@ inline int Start(Isolate* isolate, IsolateData* isolate_data,
   Local<Context> context = Context::New(isolate);
   Context::Scope context_scope(context);
   Environment env(isolate_data, context);
+  CHECK_EQ(0, uv_key_create(&thread_local_env));
+  uv_key_set(&thread_local_env, &env);
   env.Start(argc, argv, exec_argc, exec_argv, v8_is_profiling);
 
   const char* path = argc > 1 ? argv[1] : nullptr;
@@ -4399,6 +4398,7 @@ inline int Start(Isolate* isolate, IsolateData* isolate_data,
 
   const int exit_code = EmitExit(&env);
   RunAtExit(&env);
+  uv_key_delete(&thread_local_env);
 
   WaitForInspectorDisconnect(&env);
 #if defined(LEAK_SANITIZER)
diff --git a/src/node.h b/src/node.h
index e0988649f40..4452b9d578b 100644
--- a/src/node.h
+++ b/src/node.h
@@ -511,6 +511,12 @@ extern "C" NODE_EXTERN void node_module_register(void* mod);
  */
 NODE_EXTERN void AtExit(void (*cb)(void* arg), void* arg = 0);
 
+/* Registers a callback with the passed-in Environment instance. The callback
+ * is called after the event loop exits, but before the VM is disposed.
+ * Callbacks are run in reverse order of registration, i.e. newest first.
+ */
+NODE_EXTERN void AtExit(Environment* env, void (*cb)(void* arg), void* arg = 0);
+
 }  // namespace node
 
 #endif  // SRC_NODE_H_