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>

node.gyp

@@ -642,6 +642,7 @@
 
       'sources': [
         'test/cctest/test_base64.cc',
+        'test/cctest/test_environment.cc',
         'test/cctest/test_util.cc',
         'test/cctest/test_url.cc'
       ],

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

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)

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)

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_

test/cctest/node_test_fixture.cc

@@ -1,2 +0,0 @@
-#include <stdlib.h>
-#include "node_test_fixture.h"

test/cctest/node_test_fixture.h

@@ -22,7 +22,7 @@ class ArrayBufferAllocator : public v8::ArrayBuffer::Allocator {
   }
 
   virtual void* AllocateUninitialized(size_t length) {
-    return calloc(length, sizeof(int));
+    return calloc(length, 1);
   }
 
   virtual void Free(void* data, size_t) {
@@ -35,12 +35,12 @@ struct Argv {
   Argv() : Argv({"node", "-p", "process.version"}) {}
 
   Argv(const std::initializer_list<const char*> &args) {
-    int nrArgs = args.size();
+    nr_args_ = args.size();
     int totalLen = 0;
     for (auto it = args.begin(); it != args.end(); ++it) {
       totalLen += strlen(*it) + 1;
     }
-    argv_ = static_cast<char**>(malloc(nrArgs * sizeof(char*)));
+    argv_ = static_cast<char**>(malloc(nr_args_ * sizeof(char*)));
     argv_[0] = static_cast<char*>(malloc(totalLen));
     int i = 0;
     int offset = 0;
@@ -60,12 +60,17 @@ struct Argv {
     free(argv_);
   }
 
-  char** operator *() const {
+  int nr_args() const {
+    return nr_args_;
+  }
+
+  char** operator*() const {
     return argv_;
   }
 
  private:
   char** argv_;
+  int nr_args_;
 };
 
 class NodeTestFixture : public ::testing::Test {

test/cctest/test_environment.cc

@@ -0,0 +1,112 @@
+#include "node.h"
+#include "env.h"
+#include "v8.h"
+#include "libplatform/libplatform.h"
+
+#include <string>
+#include "gtest/gtest.h"
+#include "node_test_fixture.h"
+
+using node::Environment;
+using node::IsolateData;
+using node::CreateIsolateData;
+using node::FreeIsolateData;
+using node::CreateEnvironment;
+using node::FreeEnvironment;
+using node::AtExit;
+using node::RunAtExit;
+
+static bool called_cb_1 = false;
+static bool called_cb_2 = false;
+static void at_exit_callback1(void* arg);
+static void at_exit_callback2(void* arg);
+static std::string cb_1_arg;  // NOLINT(runtime/string)
+
+class EnvironmentTest : public NodeTestFixture {
+ public:
+  class Env {
+   public:
+    Env(const v8::HandleScope& handle_scope,
+        v8::Isolate* isolate,
+        const Argv& argv) {
+      context_ = v8::Context::New(isolate);
+      CHECK(!context_.IsEmpty());
+      isolate_data_ = CreateIsolateData(isolate, uv_default_loop());
+      CHECK_NE(nullptr, isolate_data_);
+      environment_ = CreateEnvironment(isolate_data_,
+                                       context_,
+                                       1, *argv,
+                                       argv.nr_args(), *argv);
+      CHECK_NE(nullptr, environment_);
+    }
+
+    ~Env() {
+      FreeIsolateData(isolate_data_);
+      FreeEnvironment(environment_);
+    }
+
+    Environment* operator*() const {
+      return environment_;
+    }
+
+   private:
+    v8::Local<v8::Context> context_;
+    IsolateData* isolate_data_;
+    Environment* environment_;
+  };
+
+ private:
+  virtual void TearDown() {
+    NodeTestFixture::TearDown();
+    called_cb_1 = false;
+    called_cb_2 = false;
+  }
+};
+
+TEST_F(EnvironmentTest, AtExitWithEnvironment) {
+  const v8::HandleScope handle_scope(isolate_);
+  const Argv argv;
+  Env env {handle_scope, isolate_, argv};
+
+  AtExit(*env, at_exit_callback1);
+  RunAtExit(*env);
+  EXPECT_TRUE(called_cb_1);
+}
+
+TEST_F(EnvironmentTest, AtExitWithArgument) {
+  const v8::HandleScope handle_scope(isolate_);
+  const Argv argv;
+  Env env {handle_scope, isolate_, argv};
+
+  std::string arg{"some args"};
+  AtExit(*env, at_exit_callback1, static_cast<void*>(&arg));
+  RunAtExit(*env);
+  EXPECT_EQ(arg, cb_1_arg);
+}
+
+TEST_F(EnvironmentTest, MultipleEnvironmentsPerIsolate) {
+  const v8::HandleScope handle_scope(isolate_);
+  const Argv argv;
+  Env env1 {handle_scope, isolate_, argv};
+  Env env2 {handle_scope, isolate_, argv};
+
+  AtExit(*env1, at_exit_callback1);
+  AtExit(*env2, at_exit_callback2);
+  RunAtExit(*env1);
+  EXPECT_TRUE(called_cb_1);
+  EXPECT_FALSE(called_cb_2);
+
+  RunAtExit(*env2);
+  EXPECT_TRUE(called_cb_2);
+}
+
+static void at_exit_callback1(void* arg) {
+  called_cb_1 = true;
+  if (arg) {
+    cb_1_arg = *static_cast<std::string*>(arg);
+  }
+}
+
+static void at_exit_callback2(void* arg) {
+  called_cb_2 = true;
+}