async_hooks: add AsyncLocal API mk2

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

@@ -8,7 +8,8 @@ const common = require('../common.js');
 const {
   createHook,
   executionAsyncResource,
-  executionAsyncId
+  executionAsyncId,
+  AsyncLocal
 } = 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'],
   asyncMethod: ['callbacks', 'async'],
   n: [1e6]
 });
@@ -102,6 +103,29 @@ function buildDestroy(getServe) {
   }
 }
 
+function buildAsyncLocal(getServe) {
+  const server = createServer(getServe(getCLS, setCLS));
+  const asyncLocal = new AsyncLocal();
+
+  return {
+    server,
+    close
+  };
+
+  function getCLS() {
+    return asyncLocal.get();
+  }
+
+  function setCLS(state) {
+    asyncLocal.set(state);
+  }
+
+  function close() {
+    asyncLocal.remove();
+    server.close();
+  }
+}
+
 function getServeAwait(getCLS, setCLS) {
   return async function serve(req, res) {
     setCLS(Math.random());
@@ -126,7 +150,8 @@ function getServeCallbacks(getCLS, setCLS) {
 
 const types = {
   'async-resource': buildCurrentResource,
-  'destroy': buildDestroy
+  'destroy': buildDestroy,
+  'async-local': buildAsyncLocal,
 };
 
 const asyncMethods = {

doc/api/async_hooks.md

@@ -579,6 +579,120 @@ const server = net.createServer((conn) => {
 Promise contexts may not get valid `triggerAsyncId`s by default. See
 the section on [promise execution tracking][].
 
+### Class: `AsyncLocal`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+This class can be used to store a value which follows asynchronous execution
+flow. Any value set on an `AsyncLocal` instance is propagated to any callback
+or promise executed within the flow. Because of that, a continuation local
+storage can be build with an `AsyncLocal` instance. This API is similar to
+thread local storage in other runtimes and languages.
+
+The implementation relies on async hooks to follow the execution flow.
+So, if an application or a library does not play nicely with async hooks,
+the same problems will be seen with the `AsyncLocal` API. In order to fix
+such issues the `AsyncResource` API should be used.
+
+The following example shows how to use `AsyncLocal` to build a simple logger
+that assignes ids to HTTP requests and includes them into messages logged
+within each request.
+
+```js
+const http = require('http');
+const { AsyncLocal } = require('async_hooks');
+
+const asyncLocal = new AsyncLocal();
+
+function print(msg) {
+  const id = asyncLocal.get();
+  console.log(`${id !== undefined ? id : '-'}:`, msg);
+}
+
+let idSeq = 0;
+http.createServer((req, res) => {
+  asyncLocal.set(idSeq++);
+  print('start');
+  setImmediate(() => {
+    print('finish');
+    res.end();
+  });
+}).listen(8080);
+
+http.get('http://localhost:8080');
+http.get('http://localhost:8080');
+// Prints:
+//   0: start
+//   1: start
+//   0: finish
+//   1: finish
+```
+
+#### `new AsyncLocal()`
+
+Creates a new instance of `AsyncLocal`.
+
+### `asyncLocal.get()`
+
+* Returns: {any}
+
+Returns the value of the `AsyncLocal` in current execution context,
+or `undefined` if the value is not set or the `AsyncLocal` was removed.
+
+### `asyncLocal.set(value)`
+
+* `value` {any}
+
+Sets the value for the `AsyncLocal` within current execution context.
+
+Once set, the value will be kept through the subsequent asynchronous calls,
+unless overridden by calling `asyncLocal.set(value)`:
+
+```js
+const asyncLocal = new AsyncLocal();
+
+setImmediate(() => {
+  asyncLocal.set('A');
+
+  setImmediate(() => {
+    console.log(asyncLocal.get());
+    // Prints: A
+
+    asyncLocal.set('B');
+    console.log(asyncLocal.get());
+    // Prints: B
+  });
+
+  console.log(asyncLocal.get());
+  // Prints: A
+
+  // Stop further value propagation
+  asyncLocal.set(undefined);
+
+  console.log(asyncLocal.get());
+  // Prints: undefined
+
+  setImmediate(() => {
+    console.log(asyncLocal.get());
+    // Prints: undefined
+  });
+});
+```
+
+If the `AsyncLocal` was removed before this call is made,
+the call will have no effect.
+
+### `asyncLocal.remove()`
+
+Disables value propagation for the `AsyncLocal` and releases all
+values stored by it. Calling `asyncLocal.remove()` multiple times will
+have no effect.
+
+Any subsequent `asyncLocal.get()` calls will return `undefined`.
+Any subsequent `asyncLocal.set(value)` calls will have no effect.
+
 ## Promise execution tracking
 
 By default, promise executions are not assigned `asyncId`s due to the relatively

lib/async_hooks.js

@@ -4,12 +4,13 @@ const {
   NumberIsSafeInteger,
   ReflectApply,
   Symbol,
+  WeakMap,
 } = primordials;
 
 const {
   ERR_ASYNC_CALLBACK,
   ERR_ASYNC_TYPE,
-  ERR_INVALID_ASYNC_ID
+  ERR_INVALID_ASYNC_ID,
 } = require('internal/errors').codes;
 const { validateString } = require('internal/validators');
 const internal_async_hooks = require('internal/async_hooks');
@@ -132,6 +133,64 @@ function createHook(fns) {
   return new AsyncHook(fns);
 }
 
+// AsyncLocal API //
+
+const locals = [];
+const localsHook = createHook({
+  init(asyncId, type, triggerAsyncId, resource) {
+    const execRes = executionAsyncResource();
+    // Using var here instead of let because "for (var ...)" is faster than let.
+    // Refs: https://github.com/nodejs/node/pull/30380#issuecomment-552948364
+    for (var i = 0; i < locals.length; i++) {
+      locals[i][kPropagateSymbol](execRes, resource);
+    }
+  }
+});
+
+const kResToValSymbol = Symbol('resToVal');
+const kPropagateSymbol = Symbol('propagate');
+
+class AsyncLocal {
+  constructor() {
+    this[kResToValSymbol] = new WeakMap();
+    locals.push(this);
+    localsHook.enable();
+  }
+
+  [kPropagateSymbol](execRes, initRes) {
+    const value = this[kResToValSymbol].get(execRes);
+    // Always overwrite value to prevent issues with reused resources.
+    this[kResToValSymbol].set(initRes, value);
+  }
+
+  get() {
+    if (this[kResToValSymbol]) {
+      const execRes = executionAsyncResource();
+      return this[kResToValSymbol].get(execRes);
+    }
+    return undefined;
+  }
+
+  set(value) {
+    if (this[kResToValSymbol]) {
+      const execRes = executionAsyncResource();
+      this[kResToValSymbol].set(execRes, value);
+    }
+  }
+
+  remove() {
+    const index = locals.indexOf(this);
+    if (index === -1)
+      return;
+
+    delete this[kResToValSymbol];
+    locals.splice(index, 1);
+    if (locals.size === 0) {
+      localsHook.disable();
+    }
+  }
+}
+
 
 // Embedder API //
 
@@ -213,6 +272,7 @@ module.exports = {
   executionAsyncId,
   triggerAsyncId,
   executionAsyncResource,
+  AsyncLocal,
   // Embedder API
   AsyncResource,
 };

test/async-hooks/test-async-local-isolation.js

@@ -0,0 +1,29 @@
+'use strict';
+
+require('../common');
+const assert = require('assert');
+const async_hooks = require('async_hooks');
+const { AsyncLocal } = async_hooks;
+
+// This test ensures isolation of `AsyncLocal`s
+// from each other in terms of stored values
+
+const asyncLocalOne = new AsyncLocal();
+const asyncLocalTwo = new AsyncLocal();
+
+setTimeout(() => {
+  assert.strictEqual(asyncLocalOne.get(), undefined);
+  assert.strictEqual(asyncLocalTwo.get(), undefined);
+
+  asyncLocalOne.set('foo');
+  asyncLocalTwo.set('bar');
+  assert.strictEqual(asyncLocalOne.get(), 'foo');
+  assert.strictEqual(asyncLocalTwo.get(), 'bar');
+
+  asyncLocalOne.set('baz');
+  asyncLocalTwo.set(42);
+  setTimeout(() => {
+    assert.strictEqual(asyncLocalOne.get(), 'baz');
+    assert.strictEqual(asyncLocalTwo.get(), 42);
+  }, 0);
+}, 0);

test/async-hooks/test-async-local-propagation.js

@@ -0,0 +1,30 @@
+'use strict';
+
+require('../common');
+const assert = require('assert');
+const async_hooks = require('async_hooks');
+const { AsyncLocal } = async_hooks;
+
+// This test ensures correct work of the global hook
+// that serves for propagation of all `AsyncLocal`s
+// in the context of `.get()`/`.set(value)` calls
+
+const asyncLocal = new AsyncLocal();
+
+setTimeout(() => {
+  assert.strictEqual(asyncLocal.get(), undefined);
+
+  asyncLocal.set('A');
+  setTimeout(() => {
+    assert.strictEqual(asyncLocal.get(), 'A');
+
+    asyncLocal.set('B');
+    setTimeout(() => {
+      assert.strictEqual(asyncLocal.get(), 'B');
+    }, 0);
+
+    assert.strictEqual(asyncLocal.get(), 'B');
+  }, 0);
+
+  assert.strictEqual(asyncLocal.get(), 'A');
+}, 0);

test/async-hooks/test-async-local-removal.js

@@ -0,0 +1,30 @@
+'use strict';
+
+require('../common');
+const assert = require('assert');
+const async_hooks = require('async_hooks');
+const { AsyncLocal } = async_hooks;
+
+// This test ensures correct work of the global hook
+// that serves for propagation of all `AsyncLocal`s
+// in the context of `.remove()` call
+
+const asyncLocalOne = new AsyncLocal();
+asyncLocalOne.set(1);
+const asyncLocalTwo = new AsyncLocal();
+asyncLocalTwo.set(2);
+
+setImmediate(() => {
+  // Removal of one local should not affect others
+  asyncLocalTwo.remove();
+  assert.strictEqual(asyncLocalOne.get(), 1);
+
+  // Removal of the last active local should not
+  // prevent propagation of locals created later
+  asyncLocalOne.remove();
+  const asyncLocalThree = new AsyncLocal();
+  asyncLocalThree.set(3);
+  setImmediate(() => {
+    assert.strictEqual(asyncLocalThree.get(), 3);
+  });
+});

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

@@ -0,0 +1,24 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const async_hooks = require('async_hooks');
+const { AsyncLocal } = async_hooks;
+
+const asyncLocal = new AsyncLocal();
+
+async function asyncFunc() {
+  return new Promise((resolve) => {
+    setTimeout(resolve, 0);
+  });
+}
+
+async function testAwait() {
+  asyncLocal.set('foo');
+  await asyncFunc();
+  assert.strictEqual(asyncLocal.get(), 'foo');
+}
+
+testAwait().then(common.mustCall(() =>
+  assert.strictEqual(asyncLocal.get(), 'foo')
+));