4 files changed, 37 insertions, 2 deletions

diff --git a/src/async_wrap.cc b/src/async_wrap.cc
index f770348b9c..5258674ff3 100644
--- a/src/async_wrap.cc
+++ b/src/async_wrap.cc
@@ -159,6 +159,12 @@ static void DestroyAsyncIdsCallback(Environment* env, void* data) {
   } while (!env->destroy_async_id_list()->empty());
 }
 
+static void DestroyAsyncIdsCallback(void* arg) {
+  Environment* env = static_cast<Environment*>(arg);
+  if (!env->destroy_async_id_list()->empty())
+    DestroyAsyncIdsCallback(env, nullptr);
+}
+
 
 void AsyncWrap::EmitPromiseResolve(Environment* env, double async_id) {
   AsyncHooks* async_hooks = env->async_hooks();
@@ -502,6 +508,8 @@ void AsyncWrap::Initialize(Local<Object> target,
   Isolate* isolate = env->isolate();
   HandleScope scope(isolate);
 
+  env->BeforeExit(DestroyAsyncIdsCallback, env);
+
   env->SetMethod(target, "setupHooks", SetupHooks);
   env->SetMethod(target, "pushAsyncIds", PushAsyncIds);
   env->SetMethod(target, "popAsyncIds", PopAsyncIds);
@@ -663,7 +671,7 @@ void AsyncWrap::EmitDestroy(Environment* env, double async_id) {
     return;
 
   if (env->destroy_async_id_list()->empty()) {
-    env->SetImmediate(DestroyAsyncIdsCallback, nullptr);
+    env->SetUnrefImmediate(DestroyAsyncIdsCallback, nullptr);
   }
 
   env->destroy_async_id_list()->push_back(async_id);
diff --git a/src/env.cc b/src/env.cc
index 706af77454..b05f0bec81 100644
--- a/src/env.cc
+++ b/src/env.cc
@@ -211,6 +211,17 @@ void Environment::PrintSyncTrace() const {
   fflush(stderr);
 }
 
+void Environment::RunBeforeExitCallbacks() {
+  for (BeforeExitCallback before_exit : before_exit_functions_) {
+    before_exit.cb_(before_exit.arg_);
+  }
+  before_exit_functions_.clear();
+}
+
+void Environment::BeforeExit(void (*cb)(void* arg), void* arg) {
+  before_exit_functions_.push_back(BeforeExitCallback{cb, arg});
+}
+
 void Environment::RunAtExitCallbacks() {
   for (AtExitCallback at_exit : at_exit_functions_) {
     at_exit.cb_(at_exit.arg_);
diff --git a/src/env.h b/src/env.h
index 813b2f5a62..f5a11dca07 100644
--- a/src/env.h
+++ b/src/env.h
@@ -658,6 +658,8 @@ class Environment {
                                 const char* name,
                                 v8::FunctionCallback callback);
 
+  void BeforeExit(void (*cb)(void* arg), void* arg);
+  void RunBeforeExitCallbacks();
   void AtExit(void (*cb)(void* arg), void* arg);
   void RunAtExitCallbacks();
 
@@ -778,6 +780,12 @@ class Environment {
 
   double* fs_stats_field_array_;
 
+  struct BeforeExitCallback {
+    void (*cb_)(void* arg);
+    void* arg_;
+  };
+  std::list<BeforeExitCallback> before_exit_functions_;
+
   struct AtExitCallback {
     void (*cb_)(void* arg);
     void* arg_;
diff --git a/src/node.cc b/src/node.cc
index 0ceddaa64e..656d132ec9 100644
--- a/src/node.cc
+++ b/src/node.cc
@@ -4319,6 +4319,14 @@ void AtExit(Environment* env, void (*cb)(void* arg), void* arg) {
 }
 
 
+void RunBeforeExit(Environment* env) {
+  env->RunBeforeExitCallbacks();
+
+  if (!uv_loop_alive(env->event_loop()))
+    EmitBeforeExit(env);
+}
+
+
 void EmitBeforeExit(Environment* env) {
   HandleScope handle_scope(env->isolate());
   Context::Scope context_scope(env->context());
@@ -4469,7 +4477,7 @@ inline int Start(Isolate* isolate, IsolateData* isolate_data,
       if (more)
         continue;
 
-      EmitBeforeExit(&env);
+      RunBeforeExit(&env);
 
       // Emit `beforeExit` if the loop became alive either after emitting
       // event, or after running some callbacks.