async_hooks: add store arg in AsyncLocalStorage

This commit introduces store as the first argument in
AsyncLocalStorage's run methods. The change is motivated by the
following expectation: most users are going to use a custom object
as the store and an extra Map created by the previous implementation
is an overhead for their use case.

Important note. This is a backwards incompatible change.
It was discussed and agreed an incompatible change is ok
since the API is still experimental and the modified
methods were only added within the last week so usage
will be minimal to none.

PR-URL: #31930
Reviewed-By: Stephen Belanger <[email protected]>
Reviewed-By: Vladimir de Turckheim <[email protected]>
Reviewed-By: Matteo Collina <[email protected]>
Reviewed-By: Michaël Zasso <[email protected]>
Reviewed-By: Michael Dawson <[email protected]>

Author: puzpuzpuz

benchmark/async_hooks/async-resource-vs-destroy.js

@@ -106,7 +106,7 @@ function buildDestroy(getServe) {
 function buildAsyncLocalStorage(getServe) {
   const asyncLocalStorage = new AsyncLocalStorage();
   const server = createServer((req, res) => {
-    asyncLocalStorage.runSyncAndReturn(() => {
+    asyncLocalStorage.runSyncAndReturn({}, () => {
       getServe(getCLS, setCLS)(req, res);
     });
   });
@@ -118,12 +118,12 @@ function buildAsyncLocalStorage(getServe) {
 
   function getCLS() {
     const store = asyncLocalStorage.getStore();
-    return store.get('store');
+    return store.state;
   }
 
   function setCLS(state) {
     const store = asyncLocalStorage.getStore();
-    store.set('store', state);
+    store.state = state;
   }
 
   function close() {

doc/api/async_hooks.md

@@ -900,7 +900,7 @@ function log(...args) {
 }
 
 http.createServer((request, response) => {
-  asyncLocalStorage.run(() => {
+  asyncLocalStorage.run(new Map(), () => {
     const store = asyncLocalStorage.getStore();
     store.set(kReq, request);
     someAsyncOperation((err, result) => {
@@ -950,27 +950,27 @@ in the current process.
 added: REPLACEME
 -->
 
-* Returns: {Map}
+* Returns: {any}
 
 This method returns the current store.
 If this method is called outside of an asynchronous context initialized by
 calling `asyncLocalStorage.run` or `asyncLocalStorage.runAndReturn`, it will
 return `undefined`.
 
-### `asyncLocalStorage.run(callback[, ...args])`
+### `asyncLocalStorage.run(store, callback[, ...args])`
 <!-- YAML
 added: REPLACEME
 -->
 
+* `store` {any}
 * `callback` {Function}
 * `...args` {any}
 
 Calling `asyncLocalStorage.run(callback)` will create a new asynchronous
-context.
-Within the callback function and the asynchronous operations from the callback,
-`asyncLocalStorage.getStore()` will return an instance of `Map` known as
-"the store". This store will be persistent through the following
-asynchronous calls.
+context. Within the callback function and the asynchronous operations from
+the callback, `asyncLocalStorage.getStore()` will return the object or
+the primitive value passed into the `store` argument (known as "the store").
+This store will be persistent through the following asynchronous calls.
 
 The callback will be ran asynchronously. Optionally, arguments can be passed
 to the function. They will be passed to the callback function.
@@ -982,10 +982,11 @@ Also, the stacktrace will be impacted by the asynchronous call.
 Example:
 
 ```js
-asyncLocalStorage.run(() => {
-  asyncLocalStorage.getStore(); // Returns a Map
+const store = { id: 1 };
+asyncLocalStorage.run(store, () => {
+  asyncLocalStorage.getStore(); // Returns the store object
   someAsyncOperation(() => {
-    asyncLocalStorage.getStore(); // Returns the same Map
+    asyncLocalStorage.getStore(); // Returns the same object
   });
 });
 asyncLocalStorage.getStore(); // Returns undefined
@@ -1014,20 +1015,21 @@ Also, the stacktrace will be impacted by the asynchronous call.
 Example:
 
 ```js
-asyncLocalStorage.run(() => {
-  asyncLocalStorage.getStore(); // Returns a Map
+asyncLocalStorage.run('store value', () => {
+  asyncLocalStorage.getStore(); // Returns 'store value'
   asyncLocalStorage.exit(() => {
     asyncLocalStorage.getStore(); // Returns undefined
   });
-  asyncLocalStorage.getStore(); // Returns the same Map
+  asyncLocalStorage.getStore(); // Returns 'store value'
 });
 ```
 
-### `asyncLocalStorage.runSyncAndReturn(callback[, ...args])`
+### `asyncLocalStorage.runSyncAndReturn(store, callback[, ...args])`
 <!-- YAML
 added: REPLACEME
 -->
 
+* `store` {any}
 * `callback` {Function}
 * `...args` {any}
 
@@ -1045,9 +1047,10 @@ the context will be exited.
 Example:
 
 ```js
+const store = { id: 2 };
 try {
-  asyncLocalStorage.runSyncAndReturn(() => {
-    asyncLocalStorage.getStore(); // Returns a Map
+  asyncLocalStorage.runSyncAndReturn(store, () => {
+    asyncLocalStorage.getStore(); // Returns the store object
     throw new Error();
   });
 } catch (e) {
@@ -1080,13 +1083,13 @@ Example:
 ```js
 // Within a call to run or runSyncAndReturn
 try {
-  asyncLocalStorage.getStore(); // Returns a Map
+  asyncLocalStorage.getStore(); // Returns the store object or value
   asyncLocalStorage.exitSyncAndReturn(() => {
     asyncLocalStorage.getStore(); // Returns undefined
     throw new Error();
   });
 } catch (e) {
-  asyncLocalStorage.getStore(); // Returns the same Map
+  asyncLocalStorage.getStore(); // Returns the same object or value
   // The error will be caught here
 }
 ```
@@ -1112,8 +1115,9 @@ It cannot be promisified using `util.promisify`. If needed, the `Promise`
 constructor can be used:
 
 ```js
+const store = new Map(); // initialize the store
 new Promise((resolve, reject) => {
-  asyncLocalStorage.run(() => {
+  asyncLocalStorage.run(store, () => {
     someFunction((err, result) => {
       if (err) {
         return reject(err);
@@ -1142,7 +1146,7 @@ the following pattern should be used:
 
 ```js
 async function fn() {
-  await asyncLocalStorage.runSyncAndReturn(() => {
+  await asyncLocalStorage.runSyncAndReturn(new Map(), () => {
     asyncLocalStorage.getStore().set('key', value);
     return foo(); // The return value of foo will be awaited
   });

lib/async_hooks.js

@@ -1,11 +1,9 @@
 'use strict';
 
 const {
-  Map,
   NumberIsSafeInteger,
   ReflectApply,
   Symbol,
-
 } = primordials;
 
 const {
@@ -248,14 +246,14 @@ class AsyncLocalStorage {
     }
   }
 
-  _enter() {
+  _enter(store) {
     if (!this.enabled) {
       this.enabled = true;
       storageList.push(this);
       storageHook.enable();
     }
     const resource = executionAsyncResource();
-    resource[this.kResourceStore] = new Map();
+    resource[this.kResourceStore] = store;
   }
 
   _exit() {
@@ -265,8 +263,8 @@ class AsyncLocalStorage {
     }
   }
 
-  runSyncAndReturn(callback, ...args) {
-    this._enter();
+  runSyncAndReturn(store, callback, ...args) {
+    this._enter(store);
     try {
       return callback(...args);
     } finally {
@@ -290,8 +288,8 @@ class AsyncLocalStorage {
     }
   }
 
-  run(callback, ...args) {
-    this._enter();
+  run(store, callback, ...args) {
+    this._enter(store);
     process.nextTick(callback, ...args);
     this._exit();
   }

test/async-hooks/test-async-local-storage-args.js

@@ -5,14 +5,14 @@ const { AsyncLocalStorage } = require('async_hooks');
 
 const asyncLocalStorage = new AsyncLocalStorage();
 
-asyncLocalStorage.run((runArg) => {
+asyncLocalStorage.run({}, (runArg) => {
   assert.strictEqual(runArg, 1);
   asyncLocalStorage.exit((exitArg) => {
     assert.strictEqual(exitArg, 2);
   }, 2);
 }, 1);
 
-asyncLocalStorage.runSyncAndReturn((runArg) => {
+asyncLocalStorage.runSyncAndReturn({}, (runArg) => {
   assert.strictEqual(runArg, 'foo');
   asyncLocalStorage.exitSyncAndReturn((exitArg) => {
     assert.strictEqual(exitArg, 'bar');

test/async-hooks/test-async-local-storage-async-await.js

@@ -12,7 +12,7 @@ async function test() {
 }
 
 async function main() {
-  await asyncLocalStorage.runSyncAndReturn(test);
+  await asyncLocalStorage.runSyncAndReturn(new Map(), test);
   assert.strictEqual(asyncLocalStorage.getStore(), undefined);
 }
 

test/async-hooks/test-async-local-storage-async-functions.js

@@ -19,7 +19,7 @@ async function testAwait() {
   await asyncLocalStorage.exitSyncAndReturn(testOut);
 }
 
-asyncLocalStorage.run(() => {
+asyncLocalStorage.run(new Map(), () => {
   const store = asyncLocalStorage.getStore();
   store.set('key', 'value');
   testAwait(); // should not reject

test/async-hooks/test-async-local-storage-enable-disable.js

@@ -5,15 +5,15 @@ const { AsyncLocalStorage } = require('async_hooks');
 
 const asyncLocalStorage = new AsyncLocalStorage();
 
-asyncLocalStorage.runSyncAndReturn(() => {
+asyncLocalStorage.runSyncAndReturn(new Map(), () => {
   asyncLocalStorage.getStore().set('foo', 'bar');
   process.nextTick(() => {
     assert.strictEqual(asyncLocalStorage.getStore().get('foo'), 'bar');
     asyncLocalStorage.disable();
     assert.strictEqual(asyncLocalStorage.getStore(), undefined);
     process.nextTick(() => {
       assert.strictEqual(asyncLocalStorage.getStore(), undefined);
-      asyncLocalStorage.runSyncAndReturn(() => {
+      asyncLocalStorage.runSyncAndReturn(new Map(), () => {
         assert.notStrictEqual(asyncLocalStorage.getStore(), undefined);
       });
     });

test/async-hooks/test-async-local-storage-errors-async.js

@@ -13,7 +13,7 @@ process.setUncaughtExceptionCaptureCallback((err) => {
   assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'node');
 });
 
-asyncLocalStorage.run(() => {
+asyncLocalStorage.run(new Map(), () => {
   const store = asyncLocalStorage.getStore();
   store.set('hello', 'node');
   setTimeout(() => {

test/async-hooks/test-async-local-storage-errors-sync-ret.js

@@ -14,7 +14,7 @@ process.setUncaughtExceptionCaptureCallback((err) => {
 });
 
 try {
-  asyncLocalStorage.runSyncAndReturn(() => {
+  asyncLocalStorage.runSyncAndReturn(new Map(), () => {
     const store = asyncLocalStorage.getStore();
     store.set('hello', 'node');
     setTimeout(() => {

test/async-hooks/test-async-local-storage-http.js

@@ -10,7 +10,7 @@ const server = http.createServer((req, res) => {
 });
 
 server.listen(0, () => {
-  asyncLocalStorage.run(() => {
+  asyncLocalStorage.run(new Map(), () => {
     const store = asyncLocalStorage.getStore();
     store.set('hello', 'world');
     http.get({ host: 'localhost', port: server.address().port }, () => {