From 7cff1e14bab190bf7632f740e3985b4afd57f1ee Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gerhard=20St=C3=B6bich?= <deb2001-github@yahoo.de>
Date: Tue, 15 Nov 2022 13:27:03 +0100
Subject: [PATCH] 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: https://github.com/nodejs/node/pull/45386
Reviewed-By: Stephen Belanger <admin@stephenbelanger.com>
Reviewed-By: Minwoo Jung <nodecorelab@gmail.com>
Reviewed-By: Chengzhong Wu <legendecas@gmail.com>
---
 doc/api/async_context.md                      | 21 +++++++-
 lib/async_hooks.js                            | 23 +++++++--
 ...st-async-local-storage-stop-propagation.js | 50 +++++++++++++++++++
 3 files changed, 89 insertions(+), 5 deletions(-)
 create mode 100644 test/async-hooks/test-async-local-storage-stop-propagation.js

diff --git a/doc/api/async_context.md b/doc/api/async_context.md
index a41b73ce67697e..2d76d41392a525 100644
--- a/doc/api/async_context.md
+++ b/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
@@ -816,4 +834,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
diff --git a/lib/async_hooks.js b/lib/async_hooks.js
index 8638bb4d7b3497..187d29281f9652 100644
--- a/lib/async_hooks.js
+++ b/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;
+      }
     }
   }
 
diff --git a/test/async-hooks/test-async-local-storage-stop-propagation.js b/test/async-hooks/test-async-local-storage-stop-propagation.js
new file mode 100644
index 00000000000000..3b63da71cb15b6
--- /dev/null
+++ b/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'
+});