src: set PromiseHooks by Environment

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>(),

src/env-inl.h

@@ -95,6 +95,24 @@ 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);
+  if (!init.IsEmpty()) js_promise_hooks_[0].SetWeak();
+  js_promise_hooks_[1].Reset(env()->isolate(), before);
+  if (!before.IsEmpty()) js_promise_hooks_[1].SetWeak();
+  js_promise_hooks_[2].Reset(env()->isolate(), after);
+  if (!after.IsEmpty()) js_promise_hooks_[2].SetWeak();
+  js_promise_hooks_[3].Reset(env()->isolate(), resolve);
+  if (!resolve.IsEmpty()) js_promise_hooks_[3].SetWeak();
+  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 +235,39 @@ void AsyncHooks::clear_async_id_stack() {
   fields_[kStackLength] = 0;
 }
 
+inline void AsyncHooks::AddContext(v8::Local<v8::Context> ctx) {
+  v8::Isolate * isolate = ctx->GetIsolate();
+  ctx->SetPromiseHooks(
+    js_promise_hooks_[0].IsEmpty() ?
+      v8::Local<v8::Function>() :
+      PersistentToLocal::Weak(isolate, js_promise_hooks_[0]),
+    js_promise_hooks_[1].IsEmpty() ?
+      v8::Local<v8::Function>() :
+      PersistentToLocal::Weak(isolate, js_promise_hooks_[1]),
+    js_promise_hooks_[2].IsEmpty() ?
+      v8::Local<v8::Function>() :
+      PersistentToLocal::Weak(isolate, js_promise_hooks_[2]),
+    js_promise_hooks_[3].IsEmpty() ?
+      v8::Local<v8::Function>() :
+      PersistentToLocal::Weak(isolate, 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) {
+  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 +355,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) {

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 {

src/node_contextify.cc

@@ -127,6 +127,9 @@ ContextifyContext::ContextifyContext(
 
 ContextifyContext::~ContextifyContext() {
   env()->RemoveCleanupHook(CleanupHook, this);
+
+  env()->async_hooks()
+    ->RemoveContext(PersistentToLocal::Weak(env()->isolate(), context_));
 }
 
 

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();