lib: improve async_context_frame structure

PR-URL: #54239 Reviewed-By: Gerhard Stöbich <[email protected]> Reviewed-By: Chengzhong Wu <[email protected]>
nodejs · Aug 14, 2024 · b897057 · b897057
1 parent a476837
commit b897057
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 21 deletions.
diff --git a/doc/api/cli.md b/doc/api/cli.md
@@ -884,8 +884,8 @@ added: REPLACEME
 
 > Stability: 1 - Experimental
 
-Enables the use of AsyncLocalStorage backed by AsyncContextFrame rather than
-the default implementation which relies on async\_hooks. This new model is
+Enables the use of [`AsyncLocalStorage`][] backed by `AsyncContextFrame` rather
+than the default implementation which relies on async\_hooks. This new model is
 implemented very differently and so could have differences in how context data
 flows within the application. As such, it is presently recommended to be sure
 your application behaviour is unaffected by this change before using it in
@@ -3527,6 +3527,7 @@ node --stack-trace-limit=12 -p -e "Error.stackTraceLimit" # prints 12
 [`--print`]: #-p---print-script
 [`--redirect-warnings`]: #--redirect-warningsfile
 [`--require`]: #-r---require-module
+[`AsyncLocalStorage`]: async_context.md#class-asynclocalstorage
 [`Atomics.wait()`]: https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Atomics/wait
 [`Buffer`]: buffer.md#class-buffer
 [`CRYPTO_secure_malloc_init`]: https://www.openssl.org/docs/man3.0/man3/CRYPTO_secure_malloc_init.html

diff --git a/lib/async_hooks.js b/lib/async_hooks.js
@@ -280,7 +280,7 @@ module.exports = {
   // Public API
   get AsyncLocalStorage() {
     return AsyncContextFrame.enabled ?
-      require('internal/async_local_storage/native') :
+      require('internal/async_local_storage/async_context_frame') :
       require('internal/async_local_storage/async_hooks');
   },
   createHook,

diff --git a/lib/internal/async_context_frame.js b/lib/internal/async_context_frame.js
@@ -1,34 +1,27 @@
 'use strict';
 
+const {
+  ObjectSetPrototypeOf,
+} = primordials;
+
 const {
   getContinuationPreservedEmbedderData,
   setContinuationPreservedEmbedderData,
 } = internalBinding('async_context_frame');
 
 let enabled_;
 
-class AsyncContextFrame extends Map {
-  constructor(store, data) {
-    super(AsyncContextFrame.current());
-    this.set(store, data);
-  }
-
+class ActiveAsyncContextFrame {
   static get enabled() {
-    enabled_ ??= require('internal/options')
-      .getOptionValue('--experimental-async-context-frame');
-    return enabled_;
+    return true;
   }
 
   static current() {
-    if (this.enabled) {
-      return getContinuationPreservedEmbedderData();
-    }
+    return getContinuationPreservedEmbedderData();
   }
 
   static set(frame) {
-    if (this.enabled) {
-      setContinuationPreservedEmbedderData(frame);
-    }
+    setContinuationPreservedEmbedderData(frame);
   }
 
   static exchange(frame) {
@@ -41,6 +34,37 @@ class AsyncContextFrame extends Map {
     const frame = this.current();
     frame?.disable(store);
   }
+}
+
+function checkEnabled() {
+  const enabled = require('internal/options')
+    .getOptionValue('--experimental-async-context-frame');
+
+  // If enabled, swap to active prototype so we don't need to check status
+  // on every interaction with the async context frame.
+  if (enabled) {
+    // eslint-disable-next-line no-use-before-define
+    ObjectSetPrototypeOf(AsyncContextFrame, ActiveAsyncContextFrame);
+  }
+
+  return enabled;
+}
+
+class AsyncContextFrame extends Map {
+  constructor(store, data) {
+    super(AsyncContextFrame.current());
+    this.set(store, data);
+  }
+
+  static get enabled() {
+    enabled_ ??= checkEnabled();
+    return enabled_;
+  }
+
+  static current() {}
+  static set(frame) {}
+  static exchange(frame) {}
+  static disable(store) {}
 
   disable(store) {
     this.delete(store);

diff --git a/lib/internal/async_local_storage/native.js → ...sync_local_storage/async_context_frame.js b/lib/internal/async_local_storage/native.js → ...sync_local_storage/async_context_frame.js
diff --git a/lib/internal/process/task_queues.js b/lib/internal/process/task_queues.js
@@ -44,7 +44,7 @@ const { AsyncResource } = require('async_hooks');
 
 const AsyncContextFrame = require('internal/async_context_frame');
 
-const async_context_frame = Symbol('asyncContextFrame');
+const async_context_frame = Symbol('kAsyncContextFrame');
 
 // *Must* match Environment::TickInfo::Fields in src/env.h.
 const kHasTickScheduled = 0;

diff --git a/lib/internal/timers.js b/lib/internal/timers.js
@@ -123,7 +123,7 @@ let debug = require('internal/util/debuglog').debuglog('timer', (fn) => {
 
 const AsyncContextFrame = require('internal/async_context_frame');
 
-const async_context_frame = Symbol('asyncContextFrame');
+const async_context_frame = Symbol('kAsyncContextFrame');
 
 // *Must* match Environment::ImmediateInfo::Fields in src/env.h.
 const kCount = 0;

diff --git a/src/api/async_resource.cc b/src/api/async_resource.cc
@@ -27,8 +27,8 @@ AsyncResource::AsyncResource(Isolate* isolate,
 
 AsyncResource::~AsyncResource() {
   CHECK_NOT_NULL(env_);
-  env_->RemoveAsyncResourceContextFrame(reinterpret_cast<std::uintptr_t>(this));
   EmitAsyncDestroy(env_, async_context_);
+  env_->RemoveAsyncResourceContextFrame(reinterpret_cast<std::uintptr_t>(this));
 }
 
 MaybeLocal<Value> AsyncResource::MakeCallback(Local<Function> callback,