src: track contexts in the Environment instead of AsyncHooks

This makes it easier to support the vm contexts in the startup
snapshot. We now manage the promise hooks using references to
the contexts from the Environment, and AsyncHooks only hold references
to the hooks.

PR-URL: https://github.com/nodejs/node/pull/45282
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Chengzhong Wu <legendecas@gmail.com>
Reviewed-By: Santiago Gimeno <santiago.gimeno@gmail.com>
Reviewed-By: Minwoo Jung <nodecorelab@gmail.com>

4 files changed, 42 insertions, 25 deletions

diff --git a/src/async_wrap.cc b/src/async_wrap.cc
index 9e76ad63ca6..e22a5a6bcff 100644
--- a/src/async_wrap.cc
+++ b/src/async_wrap.cc
@@ -185,11 +185,11 @@ static void SetupHooks(const FunctionCallbackInfo<Value>& args) {
 static void SetPromiseHooks(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);
 
-  env->async_hooks()->SetJSPromiseHooks(
-    args[0]->IsFunction() ? args[0].As<Function>() : Local<Function>(),
-    args[1]->IsFunction() ? args[1].As<Function>() : Local<Function>(),
-    args[2]->IsFunction() ? args[2].As<Function>() : Local<Function>(),
-    args[3]->IsFunction() ? args[3].As<Function>() : Local<Function>());
+  env->ResetPromiseHooks(
+      args[0]->IsFunction() ? args[0].As<Function>() : Local<Function>(),
+      args[1]->IsFunction() ? args[1].As<Function>() : Local<Function>(),
+      args[2]->IsFunction() ? args[2].As<Function>() : Local<Function>(),
+      args[3]->IsFunction() ? args[3].As<Function>() : Local<Function>());
 }
 
 class DestroyParam {
diff --git a/src/env.cc b/src/env.cc
index d5552ce9133..01cce4650b6 100644
--- a/src/env.cc
+++ b/src/env.cc
@@ -63,7 +63,7 @@ int const ContextEmbedderTag::kNodeContextTag = 0x6e6f64;
 void* const ContextEmbedderTag::kNodeContextTagPtr = const_cast<void*>(
     static_cast<const void*>(&ContextEmbedderTag::kNodeContextTag));
 
-void AsyncHooks::SetJSPromiseHooks(Local<Function> init,
+void AsyncHooks::ResetPromiseHooks(Local<Function> init,
                                    Local<Function> before,
                                    Local<Function> after,
                                    Local<Function> resolve) {
@@ -71,12 +71,20 @@ void AsyncHooks::SetJSPromiseHooks(Local<Function> init,
   js_promise_hooks_[1].Reset(env()->isolate(), before);
   js_promise_hooks_[2].Reset(env()->isolate(), after);
   js_promise_hooks_[3].Reset(env()->isolate(), resolve);
+}
+
+void Environment::ResetPromiseHooks(Local<Function> init,
+                                    Local<Function> before,
+                                    Local<Function> after,
+                                    Local<Function> resolve) {
+  async_hooks()->ResetPromiseHooks(init, before, after, resolve);
+
   for (auto it = contexts_.begin(); it != contexts_.end(); it++) {
     if (it->IsEmpty()) {
       contexts_.erase(it--);
       continue;
     }
-    PersistentToLocal::Weak(env()->isolate(), *it)
+    PersistentToLocal::Weak(isolate_, *it)
         ->SetPromiseHooks(init, before, after, resolve);
   }
 }
@@ -179,7 +187,7 @@ void AsyncHooks::clear_async_id_stack() {
   fields_[kStackLength] = 0;
 }
 
-void AsyncHooks::AddContext(Local<Context> ctx) {
+void AsyncHooks::InstallPromiseHooks(Local<Context> ctx) {
   ctx->SetPromiseHooks(js_promise_hooks_[0].IsEmpty()
                            ? Local<Function>()
                            : PersistentToLocal::Strong(js_promise_hooks_[0]),
@@ -192,23 +200,24 @@ void AsyncHooks::AddContext(Local<Context> ctx) {
                        js_promise_hooks_[3].IsEmpty()
                            ? Local<Function>()
                            : PersistentToLocal::Strong(js_promise_hooks_[3]));
+}
 
+void Environment::TrackContext(Local<Context> context) {
   size_t id = contexts_.size();
   contexts_.resize(id + 1);
-  contexts_[id].Reset(env()->isolate(), ctx);
+  contexts_[id].Reset(isolate_, context);
   contexts_[id].SetWeak();
 }
 
-void AsyncHooks::RemoveContext(Local<Context> ctx) {
-  Isolate* isolate = env()->isolate();
-  HandleScope handle_scope(isolate);
+void Environment::UntrackContext(Local<Context> context) {
+  HandleScope handle_scope(isolate_);
   contexts_.erase(std::remove_if(contexts_.begin(),
                                  contexts_.end(),
                                  [&](auto&& el) { return el.IsEmpty(); }),
                   contexts_.end());
   for (auto it = contexts_.begin(); it != contexts_.end(); it++) {
-    Local<Context> saved_context = PersistentToLocal::Weak(isolate, *it);
-    if (saved_context == ctx) {
+    Local<Context> saved_context = PersistentToLocal::Weak(isolate_, *it);
+    if (saved_context == context) {
       it->Reset();
       contexts_.erase(it);
       break;
@@ -543,7 +552,8 @@ void Environment::AssignToContext(Local<v8::Context> context,
   inspector_agent()->ContextCreated(context, info);
 #endif  // HAVE_INSPECTOR
 
-  this->async_hooks()->AddContext(context);
+  this->async_hooks()->InstallPromiseHooks(context);
+  TrackContext(context);
 }
 
 void Environment::TryLoadAddon(
@@ -1466,8 +1476,9 @@ AsyncHooks::SerializeInfo AsyncHooks::Serialize(Local<Context> context,
                 context,
                 native_execution_async_resources_[i]);
   }
-  CHECK_EQ(contexts_.size(), 1);
-  CHECK_EQ(contexts_[0], env()->context());
+
+  // At the moment, promise hooks are not supported in the startup snapshot.
+  // TODO(joyeecheung): support promise hooks in the startup snapshot.
   CHECK(js_promise_hooks_[0].IsEmpty());
   CHECK(js_promise_hooks_[1].IsEmpty());
   CHECK(js_promise_hooks_[2].IsEmpty());
@@ -1602,6 +1613,10 @@ EnvSerializeInfo Environment::Serialize(SnapshotCreator* creator) {
       should_abort_on_uncaught_toggle_.Serialize(ctx, creator);
 
   info.principal_realm = principal_realm_->Serialize(creator);
+  // For now we only support serialization of the main context.
+  // TODO(joyeecheung): support de/serialization of vm contexts.
+  CHECK_EQ(contexts_.size(), 1);
+  CHECK_EQ(contexts_[0], context());
   return info;
 }
 
diff --git a/src/env.h b/src/env.h
index a5b15863574..542d1cf6e56 100644
--- a/src/env.h
+++ b/src/env.h
@@ -303,7 +303,8 @@ class AsyncHooks : public MemoryRetainer {
   // The `js_execution_async_resources` array contains the value in that case.
   inline v8::Local<v8::Object> native_execution_async_resource(size_t index);
 
-  void SetJSPromiseHooks(v8::Local<v8::Function> init,
+  void InstallPromiseHooks(v8::Local<v8::Context> ctx);
+  void ResetPromiseHooks(v8::Local<v8::Function> init,
                          v8::Local<v8::Function> before,
                          v8::Local<v8::Function> after,
                          v8::Local<v8::Function> resolve);
@@ -322,9 +323,6 @@ class AsyncHooks : public MemoryRetainer {
   bool pop_async_context(double async_id);
   void clear_async_id_stack();  // Used in fatal exceptions.
 
-  void AddContext(v8::Local<v8::Context> ctx);
-  void RemoveContext(v8::Local<v8::Context> ctx);
-
   AsyncHooks(const AsyncHooks&) = delete;
   AsyncHooks& operator=(const AsyncHooks&) = delete;
   AsyncHooks(AsyncHooks&&) = delete;
@@ -387,8 +385,6 @@ class AsyncHooks : public MemoryRetainer {
   // Non-empty during deserialization
   const SerializeInfo* info_ = nullptr;
 
-  std::vector<v8::Global<v8::Context>> contexts_;
-
   std::array<v8::Global<v8::Function>, 4> js_promise_hooks_;
 };
 
@@ -701,9 +697,15 @@ class Environment : public MemoryRetainer {
   template <typename T, typename OnCloseCallback>
   inline void CloseHandle(T* handle, OnCloseCallback callback);
 
+  void ResetPromiseHooks(v8::Local<v8::Function> init,
+                         v8::Local<v8::Function> before,
+                         v8::Local<v8::Function> after,
+                         v8::Local<v8::Function> resolve);
   void AssignToContext(v8::Local<v8::Context> context,
                        Realm* realm,
                        const ContextInfo& info);
+  void TrackContext(v8::Local<v8::Context> context);
+  void UntrackContext(v8::Local<v8::Context> context);
 
   void StartProfilerIdleNotifier();
 
@@ -1145,6 +1147,7 @@ class Environment : public MemoryRetainer {
 
   EnabledDebugList enabled_debug_list_;
 
+  std::vector<v8::Global<v8::Context>> contexts_;
   std::list<node_module> extra_linked_bindings_;
   Mutex extra_linked_bindings_mutex_;
 
diff --git a/src/node_contextify.cc b/src/node_contextify.cc
index 1bfa01c1905..45c5b33b634 100644
--- a/src/node_contextify.cc
+++ b/src/node_contextify.cc
@@ -164,8 +164,7 @@ ContextifyContext::~ContextifyContext() {
   Isolate* isolate = env()->isolate();
   HandleScope scope(isolate);
 
-  env()->async_hooks()
-    ->RemoveContext(PersistentToLocal::Weak(isolate, context_));
+  env()->UntrackContext(PersistentToLocal::Weak(isolate, context_));
   context_.Reset();
 }