src: set PromiseHooks by Environment

The new JS PromiseHooks introduced in the referenced PR are per
v8::Context. This meant that code depending on them, such as
AsyncLocalStorage, wouldn't behave correctly across vm.Context
instances.

PromiseHooks are now synchronized across the main Context and any
Context created via vm.Context.

Refs: https://github.com/nodejs/node/pull/36394
Fixes: https://github.com/nodejs/node/issues/38781
Signed-off-by: Bryan English <bryan@bryanenglish.com>

PR-URL: https://github.com/nodejs/node/pull/38821
Reviewed-By: Stephen Belanger <admin@stephenbelanger.com>
Reviewed-By: Vladimir de Turckheim <vlad2t@hotmail.com>
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: Gerhard Stöbich <deb2001-github@yahoo.de>
Reviewed-By: Joyee Cheung <joyeec9h3@gmail.com>

5 files changed, 77 insertions, 2 deletions

diff --git a/src/async_wrap.cc b/src/async_wrap.cc
index a1c76b94138..1bca27d6a0f 100644
--- a/src/async_wrap.cc
+++ b/src/async_wrap.cc
@@ -454,8 +454,8 @@ static void EnablePromiseHook(const FunctionCallbackInfo<Value>& args) {
 
 static void SetPromiseHooks(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);
-  Local<Context> ctx = env->context();
-  ctx->SetPromiseHooks(
+
+  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>(),
diff --git a/src/env-inl.h b/src/env-inl.h
index f88e7648155..e0345b0daf2 100644
--- a/src/env-inl.h
+++ b/src/env-inl.h
@@ -95,6 +95,20 @@ v8::Local<v8::Object> AsyncHooks::native_execution_async_resource(size_t i) {
   return PersistentToLocal::Strong(native_execution_async_resources_[i]);
 }
 
+inline void AsyncHooks::SetJSPromiseHooks(v8::Local<v8::Function> init,
+                                          v8::Local<v8::Function> before,
+                                          v8::Local<v8::Function> after,
+                                          v8::Local<v8::Function> resolve) {
+  js_promise_hooks_[0].Reset(env()->isolate(), 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);
+  for (auto it = contexts_.begin(); it != contexts_.end(); it++) {
+    PersistentToLocal::Weak(env()->isolate(), *it)
+        ->SetPromiseHooks(init, before, after, resolve);
+  }
+}
+
 inline v8::Local<v8::String> AsyncHooks::provider_string(int idx) {
   return env()->isolate_data()->async_wrap_provider(idx);
 }
@@ -217,6 +231,41 @@ void AsyncHooks::clear_async_id_stack() {
   fields_[kStackLength] = 0;
 }
 
+inline void AsyncHooks::AddContext(v8::Local<v8::Context> ctx) {
+  ctx->SetPromiseHooks(
+    js_promise_hooks_[0].IsEmpty() ?
+      v8::Local<v8::Function>() :
+      PersistentToLocal::Strong(js_promise_hooks_[0]),
+    js_promise_hooks_[1].IsEmpty() ?
+      v8::Local<v8::Function>() :
+      PersistentToLocal::Strong(js_promise_hooks_[1]),
+    js_promise_hooks_[2].IsEmpty() ?
+      v8::Local<v8::Function>() :
+      PersistentToLocal::Strong(js_promise_hooks_[2]),
+    js_promise_hooks_[3].IsEmpty() ?
+      v8::Local<v8::Function>() :
+      PersistentToLocal::Strong(js_promise_hooks_[3]));
+
+  size_t id = contexts_.size();
+  contexts_.resize(id + 1);
+  contexts_[id].Reset(env()->isolate(), ctx);
+  contexts_[id].SetWeak();
+}
+
+inline void AsyncHooks::RemoveContext(v8::Local<v8::Context> ctx) {
+  v8::Isolate* isolate = env()->isolate();
+  v8::HandleScope handle_scope(isolate);
+  for (auto it = contexts_.begin(); it != contexts_.end(); it++) {
+    v8::Local<v8::Context> saved_context =
+      PersistentToLocal::Weak(env()->isolate(), *it);
+    if (saved_context == ctx) {
+      it->Reset();
+      contexts_.erase(it);
+      break;
+    }
+  }
+}
+
 // The DefaultTriggerAsyncIdScope(AsyncWrap*) constructor is defined in
 // async_wrap-inl.h to avoid a circular dependency.
 
@@ -304,6 +353,8 @@ inline void Environment::AssignToContext(v8::Local<v8::Context> context,
 #if HAVE_INSPECTOR
   inspector_agent()->ContextCreated(context, info);
 #endif  // HAVE_INSPECTOR
+
+  this->async_hooks()->AddContext(context);
 }
 
 inline Environment* Environment::GetCurrent(v8::Isolate* isolate) {
diff --git a/src/env.cc b/src/env.cc
index 0952d863b7a..9f6172de82b 100644
--- a/src/env.cc
+++ b/src/env.cc
@@ -1156,6 +1156,12 @@ AsyncHooks::SerializeInfo AsyncHooks::Serialize(Local<Context> context,
         context,
         native_execution_async_resources_[i].Get(context->GetIsolate()));
   }
+  CHECK_EQ(contexts_.size(), 1);
+  CHECK_EQ(contexts_[0], env()->context());
+  CHECK(js_promise_hooks_[0].IsEmpty());
+  CHECK(js_promise_hooks_[1].IsEmpty());
+  CHECK(js_promise_hooks_[2].IsEmpty());
+  CHECK(js_promise_hooks_[3].IsEmpty());
 
   return info;
 }
@@ -1164,6 +1170,7 @@ void AsyncHooks::MemoryInfo(MemoryTracker* tracker) const {
   tracker->TrackField("async_ids_stack", async_ids_stack_);
   tracker->TrackField("fields", fields_);
   tracker->TrackField("async_id_fields", async_id_fields_);
+  tracker->TrackField("js_promise_hooks", js_promise_hooks_);
 }
 
 void AsyncHooks::grow_async_ids_stack() {
diff --git a/src/env.h b/src/env.h
index ce88d74355c..4a50227ee8d 100644
--- a/src/env.h
+++ b/src/env.h
@@ -701,6 +701,11 @@ 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);
 
+  inline void SetJSPromiseHooks(v8::Local<v8::Function> init,
+                                v8::Local<v8::Function> before,
+                                v8::Local<v8::Function> after,
+                                v8::Local<v8::Function> resolve);
+
   inline v8::Local<v8::String> provider_string(int idx);
 
   inline void no_force_checks();
@@ -711,6 +716,9 @@ class AsyncHooks : public MemoryRetainer {
   inline bool pop_async_context(double async_id);
   inline void clear_async_id_stack();  // Used in fatal exceptions.
 
+  inline void AddContext(v8::Local<v8::Context> ctx);
+  inline void RemoveContext(v8::Local<v8::Context> ctx);
+
   AsyncHooks(const AsyncHooks&) = delete;
   AsyncHooks& operator=(const AsyncHooks&) = delete;
   AsyncHooks(AsyncHooks&&) = delete;
@@ -770,6 +778,10 @@ 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_;
 };
 
 class ImmediateInfo : public MemoryRetainer {
diff --git a/src/node_contextify.cc b/src/node_contextify.cc
index 984569ea6b5..ae5fc2fbf1a 100644
--- a/src/node_contextify.cc
+++ b/src/node_contextify.cc
@@ -127,6 +127,11 @@ ContextifyContext::ContextifyContext(
 
 ContextifyContext::~ContextifyContext() {
   env()->RemoveCleanupHook(CleanupHook, this);
+  Isolate* isolate = env()->isolate();
+  HandleScope scope(isolate);
+
+  env()->async_hooks()
+    ->RemoveContext(PersistentToLocal::Weak(isolate, context_));
 }