lib: rewrite AsyncLocalStorage without async_hooks

PR-URL: #48528 Reviewed-By: Matteo Collina <[email protected]> Reviewed-By: Yagiz Nizipli <[email protected]> Reviewed-By: James M Snell <[email protected]> Reviewed-By: Benjamin Gruenbaum <[email protected]> Reviewed-By: Rafael Gonzaga <[email protected]> Reviewed-By: Chengzhong Wu <[email protected]> Reviewed-By: Santiago Gimeno <[email protected]> Reviewed-By: Gerhard Stöbich <[email protected]>
nodejs · Aug 14, 2024 · 9ee4b16 · 9ee4b16
1 parent 5dbff81
commit 9ee4b16
Show file tree

Hide file tree

Showing 29 changed files with 658 additions and 173 deletions.
diff --git a/benchmark/async_hooks/async-local-storage-getstore-nested-resources.js b/benchmark/async_hooks/async-local-storage-getstore-nested-resources.js
@@ -38,9 +38,11 @@ function runInAsyncScopes(resourceCount, cb, i = 0) {
 
 function main({ n, resourceCount }) {
   const store = new AsyncLocalStorage();
-  runInAsyncScopes(resourceCount, () => {
-    bench.start();
-    runBenchmark(store, n);
-    bench.end(n);
+  store.run({}, () => {
+    runInAsyncScopes(resourceCount, () => {
+      bench.start();
+      runBenchmark(store, n);
+      bench.end(n);
+    });
   });
 }
diff --git a/benchmark/async_hooks/async-local-storage-getstore-nested-run.js b/benchmark/async_hooks/async-local-storage-getstore-nested-run.js
@@ -14,7 +14,7 @@ const { AsyncLocalStorage } = require('async_hooks');
  *        - AsyncLocalStorage1.getStore()
  */
 const bench = common.createBenchmark(main, {
-  sotrageCount: [1, 10, 100],
+  storageCount: [1, 10, 100],
   n: [1e4],
 });
 
@@ -34,8 +34,8 @@ function runStores(stores, value, cb, idx = 0) {
   }
 }
 
-function main({ n, sotrageCount }) {
-  const stores = new Array(sotrageCount).fill(0).map(() => new AsyncLocalStorage());
+function main({ n, storageCount }) {
+  const stores = new Array(storageCount).fill(0).map(() => new AsyncLocalStorage());
   const contextValue = {};
 
   runStores(stores, contextValue, () => {

diff --git a/doc/api/cli.md b/doc/api/cli.md
@@ -876,6 +876,21 @@ and `"` are usable.
 It is possible to run code containing inline types by passing
 [`--experimental-strip-types`][].
 
+### `--experimental-async-context-frame`
+
+<!-- YAML
+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
+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
+production.
+
 ### `--experimental-default-type=type`
 
 <!-- YAML
@@ -2942,6 +2957,7 @@ one is included in the list below.
 * `--enable-network-family-autoselection`
 * `--enable-source-maps`
 * `--experimental-abortcontroller`
+* `--experimental-async-context-frame`
 * `--experimental-default-type`
 * `--experimental-detect-module`
 * `--experimental-eventsource`

diff --git a/lib/async_hooks.js b/lib/async_hooks.js
@@ -10,7 +10,6 @@ const {
   NumberIsSafeInteger,
   ObjectDefineProperties,
   ObjectFreeze,
-  ObjectIs,
   ReflectApply,
   Symbol,
 } = primordials;
@@ -30,6 +29,8 @@ const {
 } = require('internal/validators');
 const internal_async_hooks = require('internal/async_hooks');
 
+const AsyncContextFrame = require('internal/async_context_frame');
+
 // Get functions
 // For userland AsyncResources, make sure to emit a destroy event when the
 // resource gets gced.
@@ -158,6 +159,7 @@ function createHook(fns) {
 // Embedder API //
 
 const destroyedSymbol = Symbol('destroyed');
+const contextFrameSymbol = Symbol('context_frame');
 
 class AsyncResource {
   constructor(type, opts = kEmptyObject) {
@@ -177,6 +179,8 @@ class AsyncResource {
       throw new ERR_INVALID_ASYNC_ID('triggerAsyncId', triggerAsyncId);
     }
 
+    this[contextFrameSymbol] = AsyncContextFrame.current();
+
     const asyncId = newAsyncId();
     this[async_id_symbol] = asyncId;
     this[trigger_async_id_symbol] = triggerAsyncId;
@@ -201,12 +205,12 @@ class AsyncResource {
     const asyncId = this[async_id_symbol];
     emitBefore(asyncId, this[trigger_async_id_symbol], this);
 
+    const contextFrame = this[contextFrameSymbol];
+    const prior = AsyncContextFrame.exchange(contextFrame);
     try {
-      const ret =
-        ReflectApply(fn, thisArg, args);
-
-      return ret;
+      return ReflectApply(fn, thisArg, args);
     } finally {
+      AsyncContextFrame.set(prior);
       if (hasAsyncIdStack())
         emitAfter(asyncId);
     }
@@ -270,110 +274,15 @@ class AsyncResource {
   }
 }
 
-const storageList = [];
-const storageHook = createHook({
-  init(asyncId, type, triggerAsyncId, resource) {
-    const currentResource = executionAsyncResource();
-    // Value of currentResource is always a non null object
-    for (let i = 0; i < storageList.length; ++i) {
-      storageList[i]._propagate(resource, currentResource, type);
-    }
-  },
-});
-
-class AsyncLocalStorage {
-  constructor() {
-    this.kResourceStore = Symbol('kResourceStore');
-    this.enabled = false;
-  }
-
-  static bind(fn) {
-    return AsyncResource.bind(fn);
-  }
-
-  static snapshot() {
-    return AsyncLocalStorage.bind((cb, ...args) => cb(...args));
-  }
-
-  disable() {
-    if (this.enabled) {
-      this.enabled = false;
-      // If this.enabled, the instance must be in storageList
-      ArrayPrototypeSplice(storageList,
-                           ArrayPrototypeIndexOf(storageList, this), 1);
-      if (storageList.length === 0) {
-        storageHook.disable();
-      }
-    }
-  }
-
-  _enable() {
-    if (!this.enabled) {
-      this.enabled = true;
-      ArrayPrototypePush(storageList, this);
-      storageHook.enable();
-    }
-  }
-
-  // Propagate the context from a parent resource to a child one
-  _propagate(resource, triggerResource, type) {
-    const store = triggerResource[this.kResourceStore];
-    if (this.enabled) {
-      resource[this.kResourceStore] = store;
-    }
-  }
-
-  enterWith(store) {
-    this._enable();
-    const resource = executionAsyncResource();
-    resource[this.kResourceStore] = store;
-  }
-
-  run(store, callback, ...args) {
-    // Avoid creation of an AsyncResource if store is already active
-    if (ObjectIs(store, this.getStore())) {
-      return ReflectApply(callback, null, args);
-    }
-
-    this._enable();
-
-    const resource = executionAsyncResource();
-    const oldStore = resource[this.kResourceStore];
-
-    resource[this.kResourceStore] = store;
-
-    try {
-      return ReflectApply(callback, null, args);
-    } finally {
-      resource[this.kResourceStore] = oldStore;
-    }
-  }
-
-  exit(callback, ...args) {
-    if (!this.enabled) {
-      return ReflectApply(callback, null, args);
-    }
-    this.disable();
-    try {
-      return ReflectApply(callback, null, args);
-    } finally {
-      this._enable();
-    }
-  }
-
-  getStore() {
-    if (this.enabled) {
-      const resource = executionAsyncResource();
-      return resource[this.kResourceStore];
-    }
-  }
-}
-
 // Placing all exports down here because the exported classes won't export
 // otherwise.
 module.exports = {
   // Public API
-  AsyncLocalStorage,
+  get AsyncLocalStorage() {
+    return AsyncContextFrame.enabled ?
+      require('internal/async_local_storage/native') :
+      require('internal/async_local_storage/async_hooks');
+  },
   createHook,
   executionAsyncId,
   triggerAsyncId,

diff --git a/lib/internal/async_context_frame.js b/lib/internal/async_context_frame.js
@@ -0,0 +1,50 @@
+'use strict';
+
+const {
+  getContinuationPreservedEmbedderData,
+  setContinuationPreservedEmbedderData,
+} = internalBinding('async_context_frame');
+
+let enabled_;
+
+class AsyncContextFrame extends Map {
+  constructor(store, data) {
+    super(AsyncContextFrame.current());
+    this.set(store, data);
+  }
+
+  static get enabled() {
+    enabled_ ??= require('internal/options')
+      .getOptionValue('--experimental-async-context-frame');
+    return enabled_;
+  }
+
+  static current() {
+    if (this.enabled) {
+      return getContinuationPreservedEmbedderData();
+    }
+  }
+
+  static set(frame) {
+    if (this.enabled) {
+      setContinuationPreservedEmbedderData(frame);
+    }
+  }
+
+  static exchange(frame) {
+    const prior = this.current();
+    this.set(frame);
+    return prior;
+  }
+
+  static disable(store) {
+    const frame = this.current();
+    frame?.disable(store);
+  }
+
+  disable(store) {
+    this.delete(store);
+  }
+}
+
+module.exports = AsyncContextFrame;