async_hooks: introduce async-context API

Adding AsyncLocalStorage class to async_hooks
 module.
This API provide a simple CLS-like set
of features.

Co-authored-by: Andrey Pechkurov <apechkurov@gmail.com>

PR-URL: https://github.com/nodejs/node/pull/26540
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Stephen Belanger <admin@stephenbelanger.com>
Reviewed-By: Gireesh Punathil <gpunathi@in.ibm.com>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Michaël Zasso <targos@protonmail.com>

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

@@ -8,7 +8,8 @@ const common = require('../common.js');
 const {
   createHook,
   executionAsyncResource,
-  executionAsyncId
+  executionAsyncId,
+  AsyncLocalStorage
 } = require('async_hooks');
 const { createServer } = require('http');
 
@@ -18,7 +19,7 @@ const connections = 500;
 const path = '/';
 
 const bench = common.createBenchmark(main, {
-  type: ['async-resource', 'destroy'],
+  type: ['async-resource', 'destroy', 'async-local-storage'],
   asyncMethod: ['callbacks', 'async'],
   n: [1e6]
 });
@@ -102,6 +103,35 @@ function buildDestroy(getServe) {
   }
 }
 
+function buildAsyncLocalStorage(getServe) {
+  const asyncLocalStorage = new AsyncLocalStorage();
+  const server = createServer((req, res) => {
+    asyncLocalStorage.runSyncAndReturn(() => {
+      getServe(getCLS, setCLS)(req, res);
+    });
+  });
+
+  return {
+    server,
+    close
+  };
+
+  function getCLS() {
+    const store = asyncLocalStorage.getStore();
+    return store.get('store');
+  }
+
+  function setCLS(state) {
+    const store = asyncLocalStorage.getStore();
+    store.set('store', state);
+  }
+
+  function close() {
+    asyncLocalStorage.disable();
+    server.close();
+  }
+}
+
 function getServeAwait(getCLS, setCLS) {
   return async function serve(req, res) {
     setCLS(Math.random());
@@ -126,7 +156,8 @@ function getServeCallbacks(getCLS, setCLS) {
 
 const types = {
   'async-resource': buildCurrentResource,
-  'destroy': buildDestroy
+  'destroy': buildDestroy,
+  'async-local-storage': buildAsyncLocalStorage
 };
 
 const asyncMethods = {

doc/api/async_hooks.md

@@ -859,6 +859,293 @@ for (let i = 0; i < 10; i++) {
 }
 ```
 
+## Class: `AsyncLocalStorage`
+<!-- YAML
+added: REPLACEME
+-->
+
+This class is used to create asynchronous state within callbacks and promise
+chains. It allows storing data throughout the lifetime of a web request
+or any other asynchronous duration. It is similar to thread-local storage
+in other languages.
+
+The following example builds a logger that will always know the current HTTP
+request and uses it to display enhanced logs without needing to explicitly
+provide the current HTTP request to it.
+
+```js
+const { AsyncLocalStorage } = require('async_hooks');
+const http = require('http');
+
+const kReq = 'CURRENT_REQUEST';
+const asyncLocalStorage = new AsyncLocalStorage();
+
+function log(...args) {
+  const store = asyncLocalStorage.getStore();
+  // Make sure the store exists and it contains a request.
+  if (store && store.has(kReq)) {
+    const req = store.get(kReq);
+    // Prints `GET /items ERR could not do something
+    console.log(req.method, req.url, ...args);
+  } else {
+    console.log(...args);
+  }
+}
+
+http.createServer((request, response) => {
+  asyncLocalStorage.run(() => {
+    const store = asyncLocalStorage.getStore();
+    store.set(kReq, request);
+    someAsyncOperation((err, result) => {
+      if (err) {
+        log('ERR', err.message);
+      }
+    });
+  });
+})
+.listen(8080);
+```
+
+When having multiple instances of `AsyncLocalStorage`, they are independent
+from each other. It is safe to instantiate this class multiple times.
+
+### `new AsyncLocalStorage()`
+<!-- YAML
+added: REPLACEME
+-->
+
+Creates a new instance of `AsyncLocalStorage`. Store is only provided within a
+`run` or a `runSyncAndReturn` method call.
+
+### `asyncLocalStorage.disable()`
+<!-- YAML
+added: REPLACEME
+-->
+
+This method disables the instance of `AsyncLocalStorage`. All subsequent calls
+to `asyncLocalStorage.getStore()` will return `undefined` until
+`asyncLocalStorage.run()` or `asyncLocalStorage.runSyncAndReturn()`
+is called again.
+
+When calling `asyncLocalStorage.disable()`, all current contexts linked to the
+instance will be exited.
+
+Calling `asyncLocalStorage.disable()` is required before the
+`asyncLocalStorage` can be garbage collected. This does not apply to stores
+provided by the `asyncLocalStorage`, as those objects are garbage collected
+along with the corresponding async resources.
+
+This method is to be used when the `asyncLocalStorage` is not in use anymore
+in the current process.
+
+### `asyncLocalStorage.getStore()`
+<!-- YAML
+added: REPLACEME
+-->
+
+* Returns: {Map}
+
+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])`
+<!-- YAML
+added: REPLACEME
+-->
+
+* `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.
+
+The callback will be ran asynchronously. Optionally, arguments can be passed
+to the function. They will be passed to the callback function.
+
+If an error is thrown by the callback function, it will not be caught by
+a `try/catch` block as the callback is ran in a new asynchronous resource.
+Also, the stacktrace will be impacted by the asynchronous call.
+
+Example:
+
+```js
+asyncLocalStorage.run(() => {
+  asyncLocalStorage.getStore(); // Returns a Map
+  someAsyncOperation(() => {
+    asyncLocalStorage.getStore(); // Returns the same Map
+  });
+});
+asyncLocalStorage.getStore(); // Returns undefined
+```
+
+### `asyncLocalStorage.exit(callback[, ...args])`
+<!-- YAML
+added: REPLACEME
+-->
+
+* `callback` {Function}
+* `...args` {any}
+
+Calling `asyncLocalStorage.exit(callback)` will create a new asynchronous
+context.
+Within the callback function and the asynchronous operations from the callback,
+`asyncLocalStorage.getStore()` will return `undefined`.
+
+The callback will be ran asynchronously. Optionally, arguments can be passed
+to the function. They will be passed to the callback function.
+
+If an error is thrown by the callback function, it will not be caught by
+a `try/catch` block as the callback is ran in a new asynchronous resource.
+Also, the stacktrace will be impacted by the asynchronous call.
+
+Example:
+
+```js
+asyncLocalStorage.run(() => {
+  asyncLocalStorage.getStore(); // Returns a Map
+  asyncLocalStorage.exit(() => {
+    asyncLocalStorage.getStore(); // Returns undefined
+  });
+  asyncLocalStorage.getStore(); // Returns the same Map
+});
+```
+
+### `asyncLocalStorage.runSyncAndReturn(callback[, ...args])`
+<!-- YAML
+added: REPLACEME
+-->
+
+* `callback` {Function}
+* `...args` {any}
+
+This methods runs a function synchronously within a context and return its
+return value. The store is not accessible outside of the callback function or
+the asynchronous operations created within the callback.
+
+Optionally, arguments can be passed to the function. They will be passed to
+the callback function.
+
+If the callback function throws an error, it will be thrown by
+`runSyncAndReturn` too. The stacktrace will not be impacted by this call and
+the context will be exited.
+
+Example:
+
+```js
+try {
+  asyncLocalStorage.runSyncAndReturn(() => {
+    asyncLocalStorage.getStore(); // Returns a Map
+    throw new Error();
+  });
+} catch (e) {
+  asyncLocalStorage.getStore(); // Returns undefined
+  // The error will be caught here
+}
+```
+
+### `asyncLocalStorage.exitSyncAndReturn(callback[, ...args])`
+<!-- YAML
+added: REPLACEME
+-->
+
+* `callback` {Function}
+* `...args` {any}
+
+This methods runs a function synchronously outside of a context and return its
+return value. The store is not accessible within the callback function or
+the asynchronous operations created within the callback.
+
+Optionally, arguments can be passed to the function. They will be passed to
+the callback function.
+
+If the callback function throws an error, it will be thrown by
+`exitSyncAndReturn` too. The stacktrace will not be impacted by this call and
+the context will be re-entered.
+
+Example:
+
+```js
+// Within a call to run or runSyncAndReturn
+try {
+  asyncLocalStorage.getStore(); // Returns a Map
+  asyncLocalStorage.exitSyncAndReturn(() => {
+    asyncLocalStorage.getStore(); // Returns undefined
+    throw new Error();
+  });
+} catch (e) {
+  asyncLocalStorage.getStore(); // Returns the same Map
+  // The error will be caught here
+}
+```
+
+### Choosing between `run` and `runSyncAndReturn`
+
+#### When to choose `run`
+
+`run` is asynchronous. It is called with a callback function that
+runs within a new asynchronous call. This is the most explicit behavior as
+everything that is executed within the callback of `run` (including further
+asynchronous operations) will have access to the store.
+
+If an instance of `AsyncLocalStorage` is used for error management (for
+instance, with `process.setUncaughtExceptionCaptureCallback`), only
+exceptions thrown in the scope of the callback function will be associated
+with the context.
+
+This method is the safest as it provides strong scoping and consistent
+behavior.
+
+It cannot be promisified using `util.promisify`. If needed, the `Promise`
+constructor can be used:
+
+```js
+new Promise((resolve, reject) => {
+  asyncLocalStorage.run(() => {
+    someFunction((err, result) => {
+      if (err) {
+        return reject(err);
+      }
+      return resolve(result);
+    });
+  });
+});
+```
+
+#### When to choose `runSyncAndReturn`
+
+`runSyncAndReturn` is synchronous. The callback function will be executed
+synchronously and its return value will be returned by `runSyncAndReturn`.
+The store will only be accessible from within the callback
+function and the asynchronous operations created within this scope.
+If the callback throws an error, `runSyncAndReturn` will throw it and it will
+not be associated with the context.
+
+This method provides good scoping while being synchronous.
+
+#### Usage with `async/await`
+
+If, within an async function, only one `await` call is to run within a context,
+the following pattern should be used:
+
+```js
+async function fn() {
+  await asyncLocalStorage.runSyncAndReturn(() => {
+    asyncLocalStorage.getStore().set('key', value);
+    return foo(); // The return value of foo will be awaited
+  });
+}
+```
+
+In this example, the store is only available in the callback function and the
+functions called by `foo`. Outside of `runSyncAndReturn`, calling `getStore`
+will return `undefined`.
+
 [`after` callback]: #async_hooks_after_asyncid
 [`before` callback]: #async_hooks_before_asyncid
 [`destroy` callback]: #async_hooks_destroy_asyncid

lib/async_hooks.js

@@ -1,9 +1,11 @@
 'use strict';
 
 const {
+  Map,
   NumberIsSafeInteger,
   ReflectApply,
   Symbol,
+
 } = primordials;
 
 const {
@@ -209,11 +211,102 @@ 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);
+    }
+  }
+});
+
+class AsyncLocalStorage {
+  constructor() {
+    this.kResourceStore = Symbol('kResourceStore');
+    this.enabled = false;
+  }
+
+  disable() {
+    if (this.enabled) {
+      this.enabled = false;
+      // If this.enabled, the instance must be in storageList
+      storageList.splice(storageList.indexOf(this), 1);
+      if (storageList.length === 0) {
+        storageHook.disable();
+      }
+    }
+  }
+
+  // Propagate the context from a parent resource to a child one
+  _propagate(resource, triggerResource) {
+    const store = triggerResource[this.kResourceStore];
+    if (this.enabled) {
+      resource[this.kResourceStore] = store;
+    }
+  }
+
+  _enter() {
+    if (!this.enabled) {
+      this.enabled = true;
+      storageList.push(this);
+      storageHook.enable();
+    }
+    const resource = executionAsyncResource();
+    resource[this.kResourceStore] = new Map();
+  }
+
+  _exit() {
+    const resource = executionAsyncResource();
+    if (resource) {
+      resource[this.kResourceStore] = undefined;
+    }
+  }
+
+  runSyncAndReturn(callback, ...args) {
+    this._enter();
+    try {
+      return callback(...args);
+    } finally {
+      this._exit();
+    }
+  }
+
+  exitSyncAndReturn(callback, ...args) {
+    this.enabled = false;
+    try {
+      return callback(...args);
+    } finally {
+      this.enabled = true;
+    }
+  }
+
+  getStore() {
+    const resource = executionAsyncResource();
+    if (this.enabled) {
+      return resource[this.kResourceStore];
+    }
+  }
+
+  run(callback, ...args) {
+    this._enter();
+    process.nextTick(callback, ...args);
+    this._exit();
+  }
+
+  exit(callback, ...args) {
+    this.enabled = false;
+    process.nextTick(callback, ...args);
+    this.enabled = true;
+  }
+}
 
 // Placing all exports down here because the exported classes won't export
 // otherwise.
 module.exports = {
   // Public API
+  AsyncLocalStorage,
   createHook,
   executionAsyncId,
   triggerAsyncId,

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

@@ -0,0 +1,20 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+asyncLocalStorage.run((runArg) => {
+  assert.strictEqual(runArg, 1);
+  asyncLocalStorage.exit((exitArg) => {
+    assert.strictEqual(exitArg, 2);
+  }, 2);
+}, 1);
+
+asyncLocalStorage.runSyncAndReturn((runArg) => {
+  assert.strictEqual(runArg, 'foo');
+  asyncLocalStorage.exitSyncAndReturn((exitArg) => {
+    assert.strictEqual(exitArg, 'bar');
+  }, 'bar');
+}, 'foo');

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

@@ -0,0 +1,19 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+async function test() {
+  asyncLocalStorage.getStore().set('foo', 'bar');
+  await Promise.resolve();
+  assert.strictEqual(asyncLocalStorage.getStore().get('foo'), 'bar');
+}
+
+async function main() {
+  await asyncLocalStorage.runSyncAndReturn(test);
+  assert.strictEqual(asyncLocalStorage.getStore(), undefined);
+}
+
+main();

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

@@ -0,0 +1,27 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+async function foo() {}
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+async function testOut() {
+  await foo();
+  assert.strictEqual(asyncLocalStorage.getStore(), undefined);
+}
+
+async function testAwait() {
+  await foo();
+  assert.notStrictEqual(asyncLocalStorage.getStore(), undefined);
+  assert.strictEqual(asyncLocalStorage.getStore().get('key'), 'value');
+  await asyncLocalStorage.exitSyncAndReturn(testOut);
+}
+
+asyncLocalStorage.run(() => {
+  const store = asyncLocalStorage.getStore();
+  store.set('key', 'value');
+  testAwait(); // should not reject
+});
+assert.strictEqual(asyncLocalStorage.getStore(), undefined);

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

@@ -0,0 +1,21 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+asyncLocalStorage.runSyncAndReturn(() => {
+  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(() => {
+        assert.notStrictEqual(asyncLocalStorage.getStore(), undefined);
+      });
+    });
+  });
+});

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

@@ -0,0 +1,26 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+// case 1 fully async APIS (safe)
+const asyncLocalStorage = new AsyncLocalStorage();
+
+let i = 0;
+process.setUncaughtExceptionCaptureCallback((err) => {
+  ++i;
+  assert.strictEqual(err.message, 'err' + i);
+  assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'node');
+});
+
+asyncLocalStorage.run(() => {
+  const store = asyncLocalStorage.getStore();
+  store.set('hello', 'node');
+  setTimeout(() => {
+    process.nextTick(() => {
+      assert.strictEqual(i, 2);
+    });
+    throw new Error('err2');
+  }, 0);
+  throw new Error('err1');
+});

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

@@ -0,0 +1,31 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+// case 2 using *AndReturn calls (dual behaviors)
+const asyncLocalStorage = new AsyncLocalStorage();
+
+let i = 0;
+process.setUncaughtExceptionCaptureCallback((err) => {
+  ++i;
+  assert.strictEqual(err.message, 'err2');
+  assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'node');
+});
+
+try {
+  asyncLocalStorage.runSyncAndReturn(() => {
+    const store = asyncLocalStorage.getStore();
+    store.set('hello', 'node');
+    setTimeout(() => {
+      process.nextTick(() => {
+        assert.strictEqual(i, 1);
+      });
+      throw new Error('err2');
+    }, 0);
+    throw new Error('err1');
+  });
+} catch (e) {
+  assert.strictEqual(e.message, 'err1');
+  assert.strictEqual(asyncLocalStorage.getStore(), undefined);
+}

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

@@ -0,0 +1,21 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+const http = require('http');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+const server = http.createServer((req, res) => {
+  res.end('ok');
+});
+
+server.listen(0, () => {
+  asyncLocalStorage.run(() => {
+    const store = asyncLocalStorage.getStore();
+    store.set('hello', 'world');
+    http.get({ host: 'localhost', port: server.address().port }, () => {
+      assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'world');
+      server.close();
+    });
+  });
+});

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

@@ -0,0 +1,22 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+setTimeout(() => {
+  asyncLocalStorage.run(() => {
+    const asyncLocalStorage2 = new AsyncLocalStorage();
+    asyncLocalStorage2.run(() => {
+      const store = asyncLocalStorage.getStore();
+      const store2 = asyncLocalStorage2.getStore();
+      store.set('hello', 'world');
+      store2.set('hello', 'foo');
+      setTimeout(() => {
+        assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'world');
+        assert.strictEqual(asyncLocalStorage2.getStore().get('hello'), 'foo');
+      }, 200);
+    });
+  });
+}, 100);

test/async-hooks/test-async-local-storage-no-mix-contexts.js

@@ -0,0 +1,38 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+const asyncLocalStorage2 = new AsyncLocalStorage();
+
+setTimeout(() => {
+  asyncLocalStorage.run(() => {
+    asyncLocalStorage2.run(() => {
+      const store = asyncLocalStorage.getStore();
+      const store2 = asyncLocalStorage2.getStore();
+      store.set('hello', 'world');
+      store2.set('hello', 'foo');
+      setTimeout(() => {
+        assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'world');
+        assert.strictEqual(asyncLocalStorage2.getStore().get('hello'), 'foo');
+        asyncLocalStorage.exit(() => {
+          assert.strictEqual(asyncLocalStorage.getStore(), undefined);
+          assert.strictEqual(asyncLocalStorage2.getStore().get('hello'), 'foo');
+        });
+        assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'world');
+        assert.strictEqual(asyncLocalStorage2.getStore().get('hello'), 'foo');
+      }, 200);
+    });
+  });
+}, 100);
+
+setTimeout(() => {
+  asyncLocalStorage.run(() => {
+    const store = asyncLocalStorage.getStore();
+    store.set('hello', 'earth');
+    setTimeout(() => {
+      assert.strictEqual(asyncLocalStorage.getStore().get('hello'), 'earth');
+    }, 100);
+  });
+}, 100);

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

@@ -0,0 +1,28 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+async function main() {
+  const asyncLocalStorage = new AsyncLocalStorage();
+  const err = new Error();
+  const next = () => Promise.resolve()
+    .then(() => {
+      assert.strictEqual(asyncLocalStorage.getStore().get('a'), 1);
+      throw err;
+    });
+  await new Promise((resolve, reject) => {
+    asyncLocalStorage.run(() => {
+      const store = asyncLocalStorage.getStore();
+      store.set('a', 1);
+      next().then(resolve, reject);
+    });
+  })
+    .catch((e) => {
+      assert.strictEqual(asyncLocalStorage.getStore(), undefined);
+      assert.strictEqual(e, err);
+    });
+  assert.strictEqual(asyncLocalStorage.getStore(), undefined);
+}
+
+main();