re-untangle domain/async_hooks promise support

nodejs · May 23, 2017 · e0c194b · e0c194b
1 parent 34ee31f
commit e0c194b
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 33 deletions.
diff --git a/src/async-wrap.cc b/src/async-wrap.cc
@@ -179,23 +179,50 @@ static void PushBackDestroyId(Environment* env, double id) {
 }
 
 
-static bool PreCallbackExecution(AsyncWrap* wrap) {
-  AsyncHooks* async_hooks = wrap->env()->async_hooks();
-  if (wrap->env()->using_domains()) {
-    Local<Value> domain_v = wrap->object()->Get(wrap->env()->domain_string());
-    if (domain_v->IsObject()) {
-      Local<Object> domain = domain_v.As<Object>();
-      if (domain->Get(wrap->env()->disposed_string())->IsTrue())
-        return false;
-      Local<Value> enter_v = domain->Get(wrap->env()->enter_string());
-      if (enter_v->IsFunction()) {
-        if (enter_v.As<Function>()->Call(domain, 0, nullptr).IsEmpty()) {
-          FatalError("node::AsyncWrap::MakeCallback",
-                     "domain enter callback threw, please report this");
-        }
+bool DomainEnter(Environment* env, Local<Object> object) {
+  Local<Value> domain_v = object->Get(env->domain_string());
+  if (domain_v->IsObject()) {
+    Local<Object> domain = domain_v.As<Object>();
+    if (domain->Get(env->disposed_string())->IsTrue())
+      return true;
+    Local<Value> enter_v = domain->Get(env->enter_string());
+    if (enter_v->IsFunction()) {
+      if (enter_v.As<Function>()->Call(domain, 0, nullptr).IsEmpty()) {
+        FatalError("node::AsyncWrap::MakeCallback",
+                   "domain enter callback threw, please report this");
       }
     }
   }
+  return false;
+}
+
+
+bool DomainExit(Environment* env, v8::Local<v8::Object> object) {
+  Local<Value> domain_v = object->Get(env->domain_string());
+  if (domain_v->IsObject()) {
+    Local<Object> domain = domain_v.As<Object>();
+    if (domain->Get(env->disposed_string())->IsTrue())
+      return true;
+    Local<Value> exit_v = domain->Get(env->exit_string());
+    if (exit_v->IsFunction()) {
+      if (exit_v.As<Function>()->Call(domain, 0, nullptr).IsEmpty()) {
+        FatalError("node::AsyncWrap::MakeCallback",
+                  "domain exit callback threw, please report this");
+      }
+    }
+  }
+  return false;
+}
+
+
+static bool PreCallbackExecution(AsyncWrap* wrap, bool run_domain_cbs) {
+  AsyncHooks* async_hooks = wrap->env()->async_hooks();
+
+  if (wrap->env()->using_domains() && run_domain_cbs) {
+    bool is_disposed = DomainEnter(wrap->env(), wrap->object());
+    if (is_disposed)
+      return false;
+  }
 
   if (async_hooks->fields()[AsyncHooks::kBefore] > 0) {
     Local<Value> uid = Number::New(wrap->env()->isolate(), wrap->get_id());
@@ -209,11 +236,12 @@ static bool PreCallbackExecution(AsyncWrap* wrap) {
       return false;
     }
   }
+
   return true;
 }
 
 
-static bool PostCallbackExecution(AsyncWrap* wrap) {
+static bool PostCallbackExecution(AsyncWrap* wrap, bool run_domain_cbs) {
   AsyncHooks* async_hooks = wrap->env()->async_hooks();
 
   // If the callback failed then the after() hooks will be called at the end
@@ -231,20 +259,10 @@ static bool PostCallbackExecution(AsyncWrap* wrap) {
     }
   }
 
-  if (wrap->env()->using_domains()) {
-    Local<Value> domain_v = wrap->object()->Get(wrap->env()->domain_string());
-    if (domain_v->IsObject()) {
-      Local<Object> domain = domain_v.As<Object>();
-      if (domain->Get(wrap->env()->disposed_string())->IsTrue())
-        return false;
-      Local<Value> exit_v = domain->Get(wrap->env()->exit_string());
-      if (exit_v->IsFunction()) {
-        if (exit_v.As<Function>()->Call(domain, 0, nullptr).IsEmpty()) {
-          FatalError("node::AsyncWrap::MakeCallback",
-                    "domain exit callback threw, please report this");
-        }
-      }
-    }
+  if (wrap->env()->using_domains() && run_domain_cbs) {
+    bool is_disposed = DomainExit(wrap->env(), wrap->object());
+    if (is_disposed)
+      return false;
   }
 
   return true;
@@ -301,9 +319,9 @@ static void PromiseHook(PromiseHookType type, Local<Promise> promise,
     static_cast<PromiseWrap*>(external_wrap.As<v8::External>()->Value());
   CHECK_NE(wrap, nullptr);
   if (type == PromiseHookType::kBefore) {
-    PreCallbackExecution(wrap);
+    PreCallbackExecution(wrap, false);
   } else if (type == PromiseHookType::kAfter) {
-    PostCallbackExecution(wrap);
+    PostCallbackExecution(wrap, false);
   }
 }
 
@@ -554,7 +572,7 @@ Local<Value> AsyncWrap::MakeCallback(const Local<Function> cb,
                                                 get_id(),
                                                 get_trigger_id());
 
-  if (!PreCallbackExecution(this)) {
+  if (!PreCallbackExecution(this, true)) {
     return Local<Value>();
   }
 
@@ -566,7 +584,7 @@ Local<Value> AsyncWrap::MakeCallback(const Local<Function> cb,
     return Local<Value>();
   }
 
-  if (!PostCallbackExecution(this)) {
+  if (!PostCallbackExecution(this, true)) {
     return Local<Value>();
   }
 

diff --git a/src/async-wrap.h b/src/async-wrap.h
@@ -124,6 +124,9 @@ class AsyncWrap : public BaseObject {
 
 void LoadAsyncWrapperInfo(Environment* env);
 
+bool DomainEnter(Environment* env, v8::Local<v8::Object> object);
+bool DomainExit(Environment* env, v8::Local<v8::Object> object);
+
 }  // namespace node
 
 #endif  // defined(NODE_WANT_INTERNALS) && NODE_WANT_INTERNALS

diff --git a/src/node.cc b/src/node.cc
@@ -143,6 +143,7 @@ using v8::Number;
 using v8::Object;
 using v8::ObjectTemplate;
 using v8::Promise;
+using v8::PromiseHookType;
 using v8::PromiseRejectMessage;
 using v8::PropertyCallbackInfo;
 using v8::ScriptOrigin;
@@ -1116,6 +1117,29 @@ bool ShouldAbortOnUncaughtException(Isolate* isolate) {
 }
 
 
+void DomainPromiseHook(PromiseHookType type,
+                       Local<Promise> promise,
+                       Local<Value> parent,
+                       void* arg) {
+  Environment* env = static_cast<Environment*>(arg);
+  Local<Context> context = env->context();
+
+  if (type == PromiseHookType::kInit && env->in_domain()) {
+    promise->Set(context,
+                 env->domain_string(),
+                 env->domain_array()->Get(context,
+                                          0).ToLocalChecked()).FromJust();
+    return;
+  }
+
+  if (type == PromiseHookType::kBefore) {
+    DomainEnter(env, promise);
+  } else if (type == PromiseHookType::kAfter) {
+    DomainExit(env, promise);
+  }
+}
+
+
 void SetupDomainUse(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);
 
@@ -1155,6 +1179,8 @@ void SetupDomainUse(const FunctionCallbackInfo<Value>& args) {
   Local<ArrayBuffer> array_buffer =
       ArrayBuffer::New(env->isolate(), fields, sizeof(*fields) * fields_count);
 
+  env->AddPromiseHook(DomainPromiseHook, static_cast<void*>(env));
+
   args.GetReturnValue().Set(Uint32Array::New(array_buffer, 0, fields_count));
 }