async_hooks: add hook to stop propagation

Add hook to AsyncLocalStorage to allow user to stop propagation.
This is needed to avoid leaking a store if e.g. the store indicates
that its operations are finished or it reached its time to live.

PR-URL: #45386
Reviewed-By: Stephen Belanger <[email protected]>
Reviewed-By: Minwoo Jung <[email protected]>
Reviewed-By: Chengzhong Wu <[email protected]>

doc/api/async_context.md

@@ -116,17 +116,35 @@ Each instance of `AsyncLocalStorage` maintains an independent storage context.
 Multiple instances can safely exist simultaneously without risk of interfering
 with each other's data.
 
-### `new AsyncLocalStorage()`
+### `new AsyncLocalStorage([options])`
 
 <!-- YAML
 added:
  - v13.10.0
  - v12.17.0
+changes:
+ - version: REPLACEME
+   pr-url: https://github.com/nodejs/node/pull/45386
+   description: Add option onPropagate.
 -->
 
+> Stability: 1 - `options.onPropagate` is experimental.
+
+* `options` {Object}
+  * `onPropagate` {Function} Optional callback invoked before a store is
+    propagated to a new async resource. Returning `true` allows propagation,
+    returning `false` avoids it. Default is to propagate always.
+
 Creates a new instance of `AsyncLocalStorage`. Store is only provided within a
 `run()` call or after an `enterWith()` call.
 
+The `onPropagate` is called during creation of an async resource. Throwing at
+this time will print the stack trace and exit. See
+[`async_hooks` Error handling][] for details.
+
+Creating an async resource within the `onPropagate` callback will result in
+a recursive call to `onPropagate`.
+
 ### `asyncLocalStorage.disable()`
 
 <!-- YAML
@@ -812,4 +830,5 @@ const server = createServer((req, res) => {
 [`EventEmitter`]: events.md#class-eventemitter
 [`Stream`]: stream.md#stream
 [`Worker`]: worker_threads.md#class-worker
+[`async_hooks` Error handling]: async_hooks.md#error-handling
 [`util.promisify()`]: util.md#utilpromisifyoriginal

lib/async_hooks.js

@@ -18,6 +18,7 @@ const {
 const {
   ERR_ASYNC_CALLBACK,
   ERR_ASYNC_TYPE,
+  ERR_INVALID_ARG_TYPE,
   ERR_INVALID_ASYNC_ID
 } = require('internal/errors').codes;
 const { kEmptyObject } = require('internal/util');
@@ -268,15 +269,27 @@ const storageHook = createHook({
     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);
+      storageList[i]._propagate(resource, currentResource, type);
     }
   }
 });
 
 class AsyncLocalStorage {
-  constructor() {
+  constructor(options = kEmptyObject) {
+    if (typeof options !== 'object' || options === null) {
+      throw new ERR_INVALID_ARG_TYPE('options', 'Object', options);
+    }
+
+    const { onPropagate = null } = options;
+    if (onPropagate !== null && typeof onPropagate !== 'function') {
+      throw new ERR_INVALID_ARG_TYPE('options.onPropagate',
+                                     'function',
+                                     onPropagate);
+    }
+
     this.kResourceStore = Symbol('kResourceStore');
     this.enabled = false;
+    this._onPropagate = onPropagate;
   }
 
   disable() {
@@ -300,10 +313,12 @@ class AsyncLocalStorage {
   }
 
   // Propagate the context from a parent resource to a child one
-  _propagate(resource, triggerResource) {
+  _propagate(resource, triggerResource, type) {
     const store = triggerResource[this.kResourceStore];
     if (this.enabled) {
-      resource[this.kResourceStore] = store;
+      if (this._onPropagate === null || this._onPropagate(type, store)) {
+        resource[this.kResourceStore] = store;
+      }
     }
   }
 

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

@@ -0,0 +1,50 @@
+'use strict';
+const common = require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage, AsyncResource } = require('async_hooks');
+
+let cnt = 0;
+function onPropagate(type, store) {
+  assert.strictEqual(als.getStore(), store);
+  cnt++;
+  if (cnt === 1) {
+    assert.strictEqual(type, 'r1');
+    return true;
+  }
+  if (cnt === 2) {
+    assert.strictEqual(type, 'r2');
+    return false;
+  }
+}
+
+const als = new AsyncLocalStorage({
+  onPropagate: common.mustCall(onPropagate, 2)
+});
+
+const myStore = {};
+
+als.run(myStore, common.mustCall(() => {
+  const r1 = new AsyncResource('r1');
+  const r2 = new AsyncResource('r2');
+  r1.runInAsyncScope(common.mustCall(() => {
+    assert.strictEqual(als.getStore(), myStore);
+  }));
+  r2.runInAsyncScope(common.mustCall(() => {
+    assert.strictEqual(als.getStore(), undefined);
+    r1.runInAsyncScope(common.mustCall(() => {
+      assert.strictEqual(als.getStore(), myStore);
+    }));
+  }));
+}));
+
+assert.throws(() => new AsyncLocalStorage(15), {
+  message: 'The "options" argument must be of type object. Received type number (15)',
+  code: 'ERR_INVALID_ARG_TYPE',
+  name: 'TypeError'
+});
+
+assert.throws(() => new AsyncLocalStorage({ onPropagate: 'bar' }), {
+  message: 'The "options.onPropagate" property must be of type function. Received type string (\'bar\')',
+  code: 'ERR_INVALID_ARG_TYPE',
+  name: 'TypeError'
+});