async_hooks: callback trampoline for MakeCallback

nodejs · Jun 10, 2020 · 34d91e5 · 34d91e5
1 parent 502f53f
commit 34d91e5
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 7 deletions.
diff --git a/lib/internal/async_hooks.js b/lib/internal/async_hooks.js
@@ -1,15 +1,18 @@
 'use strict';
 
 const {
+  ArrayPrototypeUnshift,
   Error,
   FunctionPrototypeBind,
   ObjectPrototypeHasOwnProperty,
   ObjectDefineProperty,
   Promise,
+  ReflectApply,
   Symbol,
 } = primordials;
 
 const async_wrap = internalBinding('async_wrap');
+const { setCallbackTrampoline } = async_wrap;
 /* async_hook_fields is a Uint32Array wrapping the uint32_t array of
  * Environment::AsyncHooks::fields_[]. Each index tracks the number of active
  * hooks for each type.
@@ -103,6 +106,26 @@ const emitDestroyNative = emitHookFactory(destroy_symbol, 'emitDestroyNative');
 const emitPromiseResolveNative =
     emitHookFactory(promise_resolve_symbol, 'emitPromiseResolveNative');
 
+function callbackTrampoline(asyncId, cb, domain_cb, ...args) {
+  if (hasHooks(kBefore))
+    emitBeforeNative(asyncId);
+
+  let result;
+  if (typeof domain_cb === 'function') {
+    ArrayPrototypeUnshift(args, cb);
+    result = ReflectApply(domain_cb, this, args);
+  } else {
+    result = ReflectApply(cb, this, args);
+  }
+
+  if (hasHooks(kAfter))
+    emitAfterNative(asyncId);
+
+  return result;
+}
+
+setCallbackTrampoline(callbackTrampoline);
+
 const topLevelResource = {};
 
 function executionAsyncResource() {

diff --git a/src/api/callback.cc b/src/api/callback.cc
@@ -158,20 +158,46 @@ MaybeLocal<Value> InternalMakeCallback(Environment* env,
     CHECK(!argv[i].IsEmpty());
 #endif
 
-  InternalCallbackScope scope(env, resource, asyncContext);
+  Local<Function> hook_cb = env->async_hooks_callback_trampoline();
+  int flags = InternalCallbackScope::kNoFlags;
+  int hook_count = 0;
+  if (!hook_cb.IsEmpty()) {
+    flags = InternalCallbackScope::kSkipAsyncHooks;
+    AsyncHooks* async_hooks = env->async_hooks();
+    hook_count = async_hooks->fields()[AsyncHooks::kBefore] +
+                 async_hooks->fields()[AsyncHooks::kAfter];
+  }
+
+  InternalCallbackScope scope(env, resource, asyncContext, flags);
   if (scope.Failed()) {
     return MaybeLocal<Value>();
   }
 
   Local<Function> domain_cb = env->domain_callback();
   MaybeLocal<Value> ret;
-  if (asyncContext.async_id != 0 || domain_cb.IsEmpty()) {
-    ret = callback->Call(env->context(), recv, argc, argv);
-  } else {
-    std::vector<Local<Value>> args(1 + argc);
+
+  if (asyncContext.async_id != 0 && hook_count != 0) {
+    MaybeStackBuffer<Local<Value>, 16> args(3 + argc);
+    args[0] = v8::Number::New(env->isolate(), asyncContext.async_id);
+    args[1] = callback;
+    if (domain_cb.IsEmpty()) {
+      args[2] = Undefined(env->isolate());
+    } else {
+      args[2] = domain_cb;
+    }
+    for (int i = 0; i < argc; i++) {
+      args[i + 3] = argv[i];
+    }
+    ret = hook_cb->Call(env->context(), recv, args.length(), &args[0]);
+  } else if (asyncContext.async_id == 0 && !domain_cb.IsEmpty()) {
+    MaybeStackBuffer<Local<Value>, 16> args(1 + argc);
     args[0] = callback;
-    std::copy(&argv[0], &argv[argc], args.begin() + 1);
-    ret = domain_cb->Call(env->context(), recv, args.size(), &args[0]);
+    for (int i = 0; i < argc; i++) {
+      args[i + 1] = argv[i];
+    }
+    ret = domain_cb->Call(env->context(), recv, args.length(), &args[0]);
+  } else {
+    ret = callback->Call(env->context(), recv, argc, argv);
   }
 
   if (ret.IsEmpty()) {

diff --git a/src/async_wrap.cc b/src/async_wrap.cc
@@ -552,6 +552,14 @@ void AsyncWrap::QueueDestroyAsyncId(const FunctionCallbackInfo<Value>& args) {
       args[0].As<Number>()->Value());
 }
 
+void AsyncWrap::SetCallbackTrampoline(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+
+  CHECK(args[0]->IsFunction());
+
+  env->set_async_hooks_callback_trampoline(args[0].As<Function>());
+}
+
 Local<FunctionTemplate> AsyncWrap::GetConstructorTemplate(Environment* env) {
   Local<FunctionTemplate> tmpl = env->async_wrap_ctor_template();
   if (tmpl.IsEmpty()) {
@@ -574,6 +582,7 @@ void AsyncWrap::Initialize(Local<Object> target,
   HandleScope scope(isolate);
 
   env->SetMethod(target, "setupHooks", SetupHooks);
+  env->SetMethod(target, "setCallbackTrampoline", SetCallbackTrampoline);
   env->SetMethod(target, "pushAsyncContext", PushAsyncContext);
   env->SetMethod(target, "popAsyncContext", PopAsyncContext);
   env->SetMethod(target, "queueDestroyAsyncId", QueueDestroyAsyncId);

diff --git a/src/async_wrap.h b/src/async_wrap.h
@@ -140,6 +140,8 @@ class AsyncWrap : public BaseObject {
   static void GetProviderType(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void QueueDestroyAsyncId(
     const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void SetCallbackTrampoline(
+    const v8::FunctionCallbackInfo<v8::Value>& args);
 
   static void EmitAsyncInit(Environment* env,
                             v8::Local<v8::Object> object,

diff --git a/src/env.h b/src/env.h
@@ -425,6 +425,7 @@ constexpr size_t kFsStatsBufferLength =
 #define ENVIRONMENT_STRONG_PERSISTENT_VALUES(V)                                \
   V(async_hooks_after_function, v8::Function)                                  \
   V(async_hooks_before_function, v8::Function)                                 \
+  V(async_hooks_callback_trampoline, v8::Function)                             \
   V(async_hooks_binding, v8::Object)                                           \
   V(async_hooks_destroy_function, v8::Function)                                \
   V(async_hooks_init_function, v8::Function)                                   \