src: set PromiseHooks by Environment

bengl · bengl · commit 7986a9bee486 · 2021-05-26T18:02:59.000-04:00
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: #36394 Fixes: #38781 Signed-off-by: Bryan English <bryan@bryanenglish.com>
diff --git 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
@@ -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::Strong(*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,37 @@ 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);
+}
+
+inline void AsyncHooks::RemoveContext(v8::Local<v8::Context> ctx) {
+  for (auto it = contexts_.begin(); it != contexts_.end(); it++) {
+    v8::Local<v8::Context> saved_context = PersistentToLocal::Strong(*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 +349,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.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,10 @@ 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 +779,10 @@ class AsyncHooks : public MemoryRetainer {
 
   // Non-empty during deserialization
   const SerializeInfo* info_ = nullptr;
+
+  std::vector<v8::Global<v8::Context>> contexts_;
+
+  v8::Global<v8::Function> js_promise_hooks_[4];
 };
 
 class ImmediateInfo : public MemoryRetainer {
diff --git a/src/node_contextify.cc b/src/node_contextify.cc
@@ -127,6 +127,8 @@ ContextifyContext::ContextifyContext(
 
 ContextifyContext::~ContextifyContext() {
   env()->RemoveCleanupHook(CleanupHook, this);
+
+  env()->async_hooks()->RemoveContext(PersistentToLocal::Strong(context_));
 }
 
 
diff --git a/test/parallel/test-async-local-storage-contexts.js b/test/parallel/test-async-local-storage-contexts.js
@@ -0,0 +1,35 @@
+'use strict';
+
+require('../common');
+const assert = require('assert');
+const vm = require('vm');
+const { AsyncLocalStorage } = require('async_hooks');
+
+// Regression test for https://github.com/nodejs/node/issues/38781
+
+const context = vm.createContext({
+  AsyncLocalStorage,
+  assert
+});
+
+vm.runInContext(`
+  const storage = new AsyncLocalStorage()
+  async function test() {
+    return storage.run({ test: 'vm' }, async () => {
+      assert.strictEqual(storage.getStore().test, 'vm');
+      await 42;
+      assert.strictEqual(storage.getStore().test, 'vm');
+    });
+  }
+  test()
+`, context);
+
+const storage = new AsyncLocalStorage();
+async function test() {
+  return storage.run({ test: 'main context' }, async () => {
+    assert.strictEqual(storage.getStore().test, 'main context');
+    await 42;
+    assert.strictEqual(storage.getStore().test, 'main context');
+  });
+}
+test();

Original file line number	Diff line number	Diff line change
`@@ -127,6 +127,8 @@ ContextifyContext::ContextifyContext(`
`127`	`127`
`128`	`128`	`ContextifyContext::~ContextifyContext() {`
`129`	`129`	`env()->RemoveCleanupHook(CleanupHook, this);`
	`130`	`+`
	`131`	`+ env()->async_hooks()->RemoveContext(PersistentToLocal::Strong(context_));`
`130`	`132`	`}`
`131`	`133`
`132`	`134`