embedding: provide hook for custom process.exit() behaviour

src/api/environment.cc

@@ -724,4 +724,16 @@ ThreadId AllocateEnvironmentThreadId() {
   return ThreadId { next_thread_id++ };
 }
 
+void DefaultProcessExitHandler(Environment* env, int exit_code) {
+  Stop(env);
+  DisposePlatform();
+  exit(exit_code);
+}
+
+
+void SetProcessExitHandler(Environment* env,
+                           std::function<void(Environment*, int)>&& handler) {
+  env->set_process_exit_handler(std::move(handler));
+}
+
 }  // namespace node

src/env-inl.h

@@ -1279,6 +1279,11 @@ void Environment::set_main_utf16(std::unique_ptr<v8::String::Value> str) {
   main_utf16_ = std::move(str);
 }
 
+void Environment::set_process_exit_handler(
+    std::function<void(Environment*, int)>&& handler) {
+  process_exit_handler_ = std::move(handler);
+}
+
 #define VP(PropertyName, StringValue) V(v8::Private, PropertyName)
 #define VY(PropertyName, StringValue) V(v8::Symbol, PropertyName)
 #define VS(PropertyName, StringValue) V(v8::String, PropertyName)

src/env.cc

@@ -501,9 +501,10 @@ void Environment::InitializeLibuv(bool start_profiler_idle_notifier) {
   }
 }
 
-void Environment::ExitEnv() {
+void Environment::Stop() {
   set_can_call_into_js(false);
   set_stopping(true);
+  stop_sub_worker_contexts();
   isolate_->TerminateExecution();
   SetImmediateThreadsafe([](Environment* env) { uv_stop(env->event_loop()); });
 }
@@ -984,14 +985,7 @@ void Environment::Exit(int exit_code) {
                     StackTrace::CurrentStackTrace(
                         isolate(), stack_trace_limit(), StackTrace::kDetailed));
   }
-  if (is_main_thread()) {
-    set_can_call_into_js(false);
-    stop_sub_worker_contexts();
-    DisposePlatform();
-    exit(exit_code);
-  } else {
-    worker_context()->Exit(exit_code);
-  }
+  process_exit_handler_(this, exit_code);
 }
 
 void Environment::stop_sub_worker_contexts() {

src/env.h

@@ -897,7 +897,7 @@ class Environment : public MemoryRetainer {
   void RegisterHandleCleanups();
   void CleanupHandles();
   void Exit(int code);
-  void ExitEnv();
+  void Stop();
 
   // Register clean-up cb to be called on environment destruction.
   inline void RegisterHandleCleanup(uv_handle_t* handle,
@@ -1257,6 +1257,8 @@ class Environment : public MemoryRetainer {
 #endif  // HAVE_INSPECTOR
 
   inline void set_main_utf16(std::unique_ptr<v8::String::Value>);
+  inline void set_process_exit_handler(
+      std::function<void(Environment*, int)>&& handler);
 
  private:
   template <typename Fn>
@@ -1459,6 +1461,9 @@ class Environment : public MemoryRetainer {
   int64_t base_object_count_ = 0;
   std::atomic_bool is_stopping_ { false };
 
+  std::function<void(Environment*, int)> process_exit_handler_ {
+      DefaultProcessExitHandler };
+
   template <typename T>
   void ForEachBaseObject(T&& iterator);
 

src/node.cc

@@ -1059,7 +1059,7 @@ int Start(int argc, char** argv) {
 }
 
 int Stop(Environment* env) {
-  env->ExitEnv();
+  env->Stop();
   return 0;
 }
 

src/node.h

@@ -224,7 +224,8 @@ class Environment;
 NODE_EXTERN int Start(int argc, char* argv[]);
 
 // Tear down Node.js while it is running (there are active handles
-// in the loop and / or actively executing JavaScript code).
+// in the loop and / or actively executing JavaScript code). This also stops
+// all Workers that may have been started earlier.
 NODE_EXTERN int Stop(Environment* env);
 
 // TODO(addaleax): Officially deprecate this and replace it with something
@@ -473,6 +474,18 @@ NODE_EXTERN v8::MaybeLocal<v8::Value> LoadEnvironment(
     std::unique_ptr<InspectorParentHandle> inspector_parent_handle = {});
 NODE_EXTERN void FreeEnvironment(Environment* env);
 
+// Set a callback that is called when process.exit() is called from JS,
+// overriding the default handler.
+// It receives the Environment* instance and the exit code as arguments.
+// This could e.g. call Stop(env); in order to terminate execution and stop
+// the event loop.
+// The default handler calls Stop(), disposes of the global V8 platform
+// instance, if one is being used, and calls exit().
+NODE_EXTERN void SetProcessExitHandler(
+    Environment* env,
+    std::function<void(Environment*, int)>&& handler);
+NODE_EXTERN void DefaultProcessExitHandler(Environment* env, int exit_code);
+
 // This may return nullptr if context is not associated with a Node instance.
 NODE_EXTERN Environment* GetCurrentEnvironment(v8::Local<v8::Context> context);
 

src/node_worker.cc

@@ -311,6 +311,9 @@ void Worker::Run() {
         if (is_stopped()) return;
         CHECK_NOT_NULL(env_);
         env_->set_env_vars(std::move(env_vars_));
+        SetProcessExitHandler(env_.get(), [this](Environment*, int exit_code) {
+          Exit(exit_code);
+        });
       }
       {
         Mutex::ScopedLock lock(mutex_);
@@ -420,9 +423,10 @@ void Worker::JoinThread() {
     MakeCallback(env()->onexit_string(), arraysize(args), args);
   }
 
-  // We cleared all libuv handles bound to this Worker above,
-  // the C++ object is no longer needed for anything now.
-  MakeWeak();
+  // If we get here, the !thread_joined_ condition at the top of the function
+  // implies that the thread was running. In that case, its final action will
+  // be to schedule a callback on the parent thread which will delete this
+  // object, so there's nothing more to do here.
 }
 
 Worker::~Worker() {

test/cctest/test_environment.cc

@@ -447,3 +447,20 @@ TEST_F(EnvironmentTest, InspectorMultipleEmbeddedEnvironments) {
   CHECK_EQ(from_inspector->IntegerValue(context).FromJust(), 42);
 }
 #endif  // HAVE_INSPECTOR
+
+TEST_F(EnvironmentTest, ExitHandlerTest) {
+  const v8::HandleScope handle_scope(isolate_);
+  const Argv argv;
+
+  int callback_calls = 0;
+
+  Env env {handle_scope, argv};
+  SetProcessExitHandler(*env, [&](node::Environment* env_, int exit_code) {
+    EXPECT_EQ(*env, env_);
+    EXPECT_EQ(exit_code, 42);
+    callback_calls++;
+    node::Stop(*env);
+  });
+  node::LoadEnvironment(*env, "process.exit(42)").ToLocalChecked();
+  EXPECT_EQ(callback_calls, 1);
+}