async_hooks: improve resource stack performance

Removes some of the performance overhead that came with
`executionAsyncResource()` by storing async resources that
are managed by JS and those managed by C++ separately, and
instead caching the result of `executionAsyncResource()` with
low overhead to avoid multiple calls into C++. In particular,
this is useful when using the async_hooks callback trampoline.

This particularly improves performance when async hooks are not
being used.

(This is continuation of nodejs#33575.)

Refs: nodejs#33575

lib/internal/async_hooks.js

@@ -8,6 +8,7 @@ const {
   ObjectDefineProperty,
   Promise,
   ReflectApply,
+  SafeMap,
   Symbol,
 } = primordials;
 
@@ -41,7 +42,6 @@ const { setCallbackTrampoline } = async_wrap;
 const {
   async_hook_fields,
   async_id_fields,
-  execution_async_resources
 } = async_wrap;
 // Store the pair executionAsyncId and triggerAsyncId in a AliasedFloat64Array
 // in Environment::AsyncHooks::async_ids_stack_ which tracks the resource
@@ -50,7 +50,9 @@ const {
 // each hook's after() callback.
 const {
   pushAsyncContext: pushAsyncContext_,
-  popAsyncContext: popAsyncContext_
+  popAsyncContext: popAsyncContext_,
+  clearAsyncIdStack,
+  executionAsyncResource: executionAsyncResource_,
 } = async_wrap;
 // For performance reasons, only track Promises when a hook is enabled.
 const { enablePromiseHook, disablePromiseHook } = async_wrap;
@@ -87,9 +89,12 @@ const { resource_symbol, owner_symbol } = internalBinding('symbols');
 // Each constant tracks how many callbacks there are for any given step of
 // async execution. These are tracked so if the user didn't include callbacks
 // for a given step, that step can bail out early.
-const { kInit, kBefore, kAfter, kDestroy, kTotals, kPromiseResolve,
-        kCheck, kExecutionAsyncId, kAsyncIdCounter, kTriggerAsyncId,
-        kDefaultTriggerAsyncId, kStackLength } = async_wrap.constants;
+const {
+  kInit, kBefore, kAfter, kDestroy, kTotals, kPromiseResolve,
+  kCheck, kExecutionAsyncId, kAsyncIdCounter, kTriggerAsyncId,
+  kDefaultTriggerAsyncId, kCachedResourceIsValid, kStackLength,
+  kUsesExecutionAsyncResource
+} = async_wrap.constants;
 
 const { async_id_symbol,
         trigger_async_id_symbol } = internalBinding('symbols');
@@ -111,7 +116,10 @@ function useDomainTrampoline(fn) {
   domain_cb = fn;
 }
 
-function callbackTrampoline(asyncId, cb, ...args) {
+function callbackTrampoline(asyncId, resource, cb, ...args) {
+  const index = async_hook_fields[kStackLength] - 1;
+  cacheExecutionAsyncResource(resource);
+
   if (asyncId !== 0 && hasHooks(kBefore))
     emitBeforeNative(asyncId);
 
@@ -126,18 +134,41 @@ function callbackTrampoline(asyncId, cb, ...args) {
   if (asyncId !== 0 && hasHooks(kAfter))
     emitAfterNative(asyncId);
 
+  cachedResourceHolder.clear();
+
   return result;
 }
 
 setCallbackTrampoline(callbackTrampoline);
 
 const topLevelResource = {};
+// Contains a [stack height] -> [resource] map for things pushed onto the
+// async resource stack from JS.
+const jsResourceStack = new SafeMap();
+// Contains either a single key (null) or nothing. If the key is present,
+// this points to the current async resource.
+const cachedResourceHolder = new SafeMap();
+
+function cacheExecutionAsyncResource(resource) {
+  const publicResource = lookupPublicResource(resource);
+  async_hook_fields[kCachedResourceIsValid] = 1;
+  cachedResourceHolder.set(null, publicResource);
+  return publicResource;
+}
 
 function executionAsyncResource() {
+  // If we use this function once, we can expect it to be used again. In that
+  // case, this tells the native code to use the callbackTrampoline for entering
+  // JS, which also passed the native resource information to us.
+  async_hook_fields[kUsesExecutionAsyncResource] = 1;
+
   const index = async_hook_fields[kStackLength] - 1;
   if (index === -1) return topLevelResource;
-  const resource = execution_async_resources[index];
-  return lookupPublicResource(resource);
+
+  if (async_hook_fields[kCachedResourceIsValid])
+    return cachedResourceHolder.get(null);
+  const resource = jsResourceStack.get(index) ?? executionAsyncResource_();
+  return cacheExecutionAsyncResource(resource);
 }
 
 // Used to fatally abort the process if a callback throws.
@@ -478,16 +509,6 @@ function emitDestroyScript(asyncId) {
 }
 
 
-// Keep in sync with Environment::AsyncHooks::clear_async_id_stack
-// in src/env-inl.h.
-function clearAsyncIdStack() {
-  async_id_fields[kExecutionAsyncId] = 0;
-  async_id_fields[kTriggerAsyncId] = 0;
-  async_hook_fields[kStackLength] = 0;
-  execution_async_resources.splice(0, execution_async_resources.length);
-}
-
-
 function hasAsyncIdStack() {
   return hasHooks(kStackLength);
 }
@@ -500,15 +521,22 @@ function pushAsyncContext(asyncId, triggerAsyncId, resource) {
     return pushAsyncContext_(asyncId, triggerAsyncId, resource);
   async_wrap.async_ids_stack[offset * 2] = async_id_fields[kExecutionAsyncId];
   async_wrap.async_ids_stack[offset * 2 + 1] = async_id_fields[kTriggerAsyncId];
-  execution_async_resources[offset] = resource;
   async_hook_fields[kStackLength]++;
   async_id_fields[kExecutionAsyncId] = asyncId;
   async_id_fields[kTriggerAsyncId] = triggerAsyncId;
+  jsResourceStack.set(offset, resource);
+  if (async_hook_fields[kCachedResourceIsValid])
+    cachedResourceHolder.clear();
+  async_hook_fields[kCachedResourceIsValid] = 0;
 }
 
 
 // This is the equivalent of the native pop_async_ids() call.
 function popAsyncContext(asyncId) {
+  async_hook_fields[kCachedResourceIsValid] = 0;
+  if (async_hook_fields[kCachedResourceIsValid])
+    cachedResourceHolder.clear();
+
   const stackLength = async_hook_fields[kStackLength];
   if (stackLength === 0) return false;
 
@@ -518,9 +546,17 @@ function popAsyncContext(asyncId) {
   }
 
   const offset = stackLength - 1;
+
+  if (!jsResourceStack.has(offset)) {
+    // For some reason this popAsyncContext() call removes a resource for a
+    // corresponding push() call from C++, so let the C++ code handle this
+    // pop operation since it's on the native the resource stack.
+    return popAsyncContext_(asyncId);
+  }
+
   async_id_fields[kExecutionAsyncId] = async_wrap.async_ids_stack[2 * offset];
   async_id_fields[kTriggerAsyncId] = async_wrap.async_ids_stack[2 * offset + 1];
-  execution_async_resources.pop();
+  jsResourceStack.delete(offset);
   async_hook_fields[kStackLength] = offset;
   return offset > 0;
 }
@@ -576,5 +612,7 @@ module.exports = {
     after: emitAfterNative,
     destroy: emitDestroyNative,
     promise_resolve: emitPromiseResolveNative
-  }
+  },
+  jsResourceStack,
+  cachedResourceHolder
 };

lib/internal/bootstrap/node.js

@@ -106,8 +106,11 @@ if (credentials.implementsPosixCredentials) {
 // process. They use the same functions as the JS embedder API. These callbacks
 // are setup immediately to prevent async_wrap.setupHooks() from being hijacked
 // and the cost of doing so is negligible.
-const { nativeHooks } = require('internal/async_hooks');
-internalBinding('async_wrap').setupHooks(nativeHooks);
+const {
+  nativeHooks, jsResourceStack, cachedResourceHolder
+} = require('internal/async_hooks');
+internalBinding('async_wrap').setupHooks(
+  nativeHooks, jsResourceStack, cachedResourceHolder);
 
 const {
   setupTaskQueue,

src/api/callback.cc

@@ -161,11 +161,12 @@ MaybeLocal<Value> InternalMakeCallback(Environment* env,
   Local<Function> hook_cb = env->async_hooks_callback_trampoline();
   int flags = InternalCallbackScope::kNoFlags;
   int hook_count = 0;
+  AsyncHooks* async_hooks = env->async_hooks();
   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];
+                 async_hooks->fields()[AsyncHooks::kAfter] +
+                 async_hooks->fields()[AsyncHooks::kUsesExecutionAsyncResource];
   }
 
   InternalCallbackScope scope(env, resource, asyncContext, flags);
@@ -176,16 +177,18 @@ MaybeLocal<Value> InternalMakeCallback(Environment* env,
   MaybeLocal<Value> ret;
 
   if (hook_count != 0) {
-    MaybeStackBuffer<Local<Value>, 16> args(2 + argc);
+    MaybeStackBuffer<Local<Value>, 16> args(3 + argc);
     args[0] = v8::Number::New(env->isolate(), asyncContext.async_id);
-    args[1] = callback;
+    args[1] = resource;
+    args[2] = callback;
     for (int i = 0; i < argc; i++) {
-      args[i + 2] = argv[i];
+      args[i + 3] = argv[i];
     }
     ret = hook_cb->Call(env->context(), recv, args.length(), &args[0]);
   } else {
     ret = callback->Call(env->context(), recv, argc, argv);
   }
+  async_hooks->fields()[AsyncHooks::kCachedResourceIsValid] = 0;
 
   if (ret.IsEmpty()) {
     scope.MarkAsFailed();

src/async_wrap.cc

@@ -39,6 +39,7 @@ using v8::HandleScope;
 using v8::Integer;
 using v8::Isolate;
 using v8::Local;
+using v8::Map;
 using v8::MaybeLocal;
 using v8::Name;
 using v8::Number;
@@ -414,6 +415,11 @@ static void SetupHooks(const FunctionCallbackInfo<Value>& args) {
   SET_HOOK_FN(destroy);
   SET_HOOK_FN(promise_resolve);
 #undef SET_HOOK_FN
+
+  CHECK(args[1]->IsMap());
+  CHECK(args[2]->IsMap());
+  env->async_hooks()->set_js_execution_async_resources(args[1].As<Map>());
+  env->async_hooks()->set_cached_resource_holder(args[2].As<Map>());
 }
 
 static void EnablePromiseHook(const FunctionCallbackInfo<Value>& args) {
@@ -574,6 +580,16 @@ Local<FunctionTemplate> AsyncWrap::GetConstructorTemplate(Environment* env) {
   return tmpl;
 }
 
+static void ExecutionAsyncResource(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  args.GetReturnValue().Set(env->async_hooks()->execution_async_resource());
+}
+
+static void ClearAsyncIdStack(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  env->async_hooks()->clear_async_id_stack();
+}
+
 void AsyncWrap::Initialize(Local<Object> target,
                            Local<Value> unused,
                            Local<Context> context,
@@ -590,6 +606,8 @@ void AsyncWrap::Initialize(Local<Object> target,
   env->SetMethod(target, "enablePromiseHook", EnablePromiseHook);
   env->SetMethod(target, "disablePromiseHook", DisablePromiseHook);
   env->SetMethod(target, "registerDestroyHook", RegisterDestroyHook);
+  env->SetMethod(target, "executionAsyncResource", ExecutionAsyncResource);
+  env->SetMethod(target, "clearAsyncIdStack", ClearAsyncIdStack);
 
   PropertyAttribute ReadOnlyDontDelete =
       static_cast<PropertyAttribute>(ReadOnly | DontDelete);
@@ -622,10 +640,6 @@ void AsyncWrap::Initialize(Local<Object> target,
                          "async_id_fields",
                          env->async_hooks()->async_id_fields().GetJSArray());
 
-  FORCE_SET_TARGET_FIELD(target,
-                         "execution_async_resources",
-                         env->async_hooks()->execution_async_resources());
-
   target->Set(context,
               env->async_ids_stack_string(),
               env->async_hooks()->async_ids_stack().GetJSArray()).Check();
@@ -646,6 +660,8 @@ void AsyncWrap::Initialize(Local<Object> target,
   SET_HOOKS_CONSTANT(kTriggerAsyncId);
   SET_HOOKS_CONSTANT(kAsyncIdCounter);
   SET_HOOKS_CONSTANT(kDefaultTriggerAsyncId);
+  SET_HOOKS_CONSTANT(kCachedResourceIsValid);
+  SET_HOOKS_CONSTANT(kUsesExecutionAsyncResource);
   SET_HOOKS_CONSTANT(kStackLength);
 #undef SET_HOOKS_CONSTANT
   FORCE_SET_TARGET_FIELD(target, "constants", constants);

src/env-inl.h

@@ -105,8 +105,32 @@ inline AliasedFloat64Array& AsyncHooks::async_ids_stack() {
   return async_ids_stack_;
 }
 
-inline v8::Local<v8::Array> AsyncHooks::execution_async_resources() {
-  return PersistentToLocal::Strong(execution_async_resources_);
+v8::Local<v8::Value> AsyncHooks::execution_async_resource() {
+  if (fields_[kStackLength] == 0) return {};
+  uint32_t offset = fields_[kStackLength] - 1;
+  if (LIKELY(offset < native_execution_async_resources_.size()))
+    return PersistentToLocal::Strong(native_execution_async_resources_[offset]);
+
+  // This async resource was stored in the JS async resource map.
+  v8::Local<v8::Map> js_map =
+      PersistentToLocal::Strong(js_execution_async_resources_);
+  v8::Local<v8::Integer> key =
+      v8::Integer::NewFromUnsigned(env()->isolate(), offset);
+  v8::Local<v8::Value> ret;
+  if (UNLIKELY(js_map.IsEmpty() ||
+               !js_map->Get(env()->context(), key).ToLocal(&ret) ||
+               ret->IsUndefined())) {
+    return {};
+  }
+  return ret;
+}
+
+void AsyncHooks::set_js_execution_async_resources(v8::Local<v8::Map> value) {
+  js_execution_async_resources_.Reset(env()->isolate(), value);
+}
+
+void AsyncHooks::set_cached_resource_holder(v8::Local<v8::Map> value) {
+  cached_resource_holder_.Reset(env()->isolate(), value);
 }
 
 inline v8::Local<v8::String> AsyncHooks::provider_string(int idx) {
@@ -143,12 +167,23 @@ inline void AsyncHooks::push_async_context(double async_id,
   async_id_fields_[kExecutionAsyncId] = async_id;
   async_id_fields_[kTriggerAsyncId] = trigger_async_id;
 
-  auto resources = execution_async_resources();
-  USE(resources->Set(env()->context(), offset, resource));
+#ifdef DEBUG
+  for (uint32_t i = offset; i < native_execution_async_resources_.size(); i++)
+    CHECK(native_execution_async_resources_[i].IsEmpty());
+#endif
+  native_execution_async_resources_.resize(offset + 1);
+  native_execution_async_resources_[offset].Reset(env()->isolate(), resource);
+  if (fields_[kCachedResourceIsValid])
+    PersistentToLocal::Strong(cached_resource_holder_)->Clear();
+  fields_[kCachedResourceIsValid] = 0;
 }
 
 // Remember to keep this code aligned with popAsyncContext() in JS.
 inline bool AsyncHooks::pop_async_context(double async_id) {
+  if (fields_[kCachedResourceIsValid])
+    PersistentToLocal::Strong(cached_resource_holder_)->Clear();
+  fields_[kCachedResourceIsValid] = 0;
+
   // In case of an exception then this may have already been reset, if the
   // stack was multiple MakeCallback()'s deep.
   if (fields_[kStackLength] == 0) return false;
@@ -177,21 +212,42 @@ inline bool AsyncHooks::pop_async_context(double async_id) {
   async_id_fields_[kTriggerAsyncId] = async_ids_stack_[2 * offset + 1];
   fields_[kStackLength] = offset;
 
-  auto resources = execution_async_resources();
-  USE(resources->Delete(env()->context(), offset));
+  if (LIKELY(offset < native_execution_async_resources_.size() &&
+             !native_execution_async_resources_[offset].IsEmpty())) {
+#ifdef DEBUG
+    for (uint32_t i = offset + 1;
+         i < native_execution_async_resources_.size();
+         i++) {
+      CHECK(native_execution_async_resources_[i].IsEmpty());
+    }
+#endif
+    native_execution_async_resources_[offset].Reset();
+    native_execution_async_resources_.resize(offset);
+    if (native_execution_async_resources_.size() <
+        native_execution_async_resources_.capacity() / 2) {
+      native_execution_async_resources_.shrink_to_fit();
+    }
+  } else {
+    USE(PersistentToLocal::Strong(js_execution_async_resources_)->Delete(
+        env()->context(),
+        v8::Integer::NewFromUnsigned(env()->isolate(), offset)));
+  }
 
   return fields_[kStackLength] > 0;
 }
 
-// Keep in sync with clearAsyncIdStack in lib/internal/async_hooks.js.
-inline void AsyncHooks::clear_async_id_stack() {
-  auto isolate = env()->isolate();
-  v8::HandleScope handle_scope(isolate);
-  execution_async_resources_.Reset(isolate, v8::Array::New(isolate));
+void AsyncHooks::clear_async_id_stack() {
+  if (!js_execution_async_resources_.IsEmpty()) {
+    PersistentToLocal::Strong(cached_resource_holder_)->Clear();
+    PersistentToLocal::Strong(js_execution_async_resources_)->Clear();
+  }
+  native_execution_async_resources_.clear();
+  native_execution_async_resources_.shrink_to_fit();
 
   async_id_fields_[kExecutionAsyncId] = 0;
   async_id_fields_[kTriggerAsyncId] = 0;
   fields_[kStackLength] = 0;
+  fields_[kCachedResourceIsValid] = 0;
 }
 
 // The DefaultTriggerAsyncIdScope(AsyncWrap*) constructor is defined in