From a4782127eb266ebab350c09e2faa9d7358d8e2dd Mon Sep 17 00:00:00 2001
From: Stephen Belanger <admin@stephenbelanger.com>
Date: Tue, 25 Feb 2020 13:46:46 +0800
Subject: [PATCH] async_hooks: add sync enterWith to ALS

This allows transitioning the entire following sync and async execution
sub-tree to the given async storage context. With this one can be sure
the context binding will remain for any following sync activity and all
descending async execution whereas the `run*(...)` methods must wrap
everything that is intended to exist within the context. This is helpful
for scenarios such as prepending a `'connection'` event to an http
server which binds everything that occurs within each request to
the given context. This is helpful for APMs to minimize the need
for patching and especially adding closures.

PR-URL: https://github.com/nodejs/node/pull/31945
Reviewed-By: Vladimir de Turckheim <vlad2t@hotmail.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Michael Dawson <michael_dawson@ca.ibm.com>
---
 doc/api/async_hooks.md                        | 42 +++++++++++++++++++
 lib/async_hooks.js                            |  6 +--
 .../test-async-local-storage-enter-with.js    | 20 +++++++++
 3 files changed, 65 insertions(+), 3 deletions(-)
 create mode 100644 test/async-hooks/test-async-local-storage-enter-with.js

diff --git a/doc/api/async_hooks.md b/doc/api/async_hooks.md
index ac432a4865d137..f0c2373448711a 100644
--- a/doc/api/async_hooks.md
+++ b/doc/api/async_hooks.md
@@ -957,6 +957,48 @@ If this method is called outside of an asynchronous context initialized by
 calling `asyncLocalStorage.run` or `asyncLocalStorage.runAndReturn`, it will
 return `undefined`.
 
+### `asyncLocalStorage.enterWith(store)`
+<!-- YAML
+added: REPLACEME
+-->
+
+* `store` {any}
+
+Calling `asyncLocalStorage.enterWith(store)` will transition into the context
+for the remainder of the current synchronous execution and will persist
+through any following asynchronous calls.
+
+Example:
+
+```js
+const store = { id: 1 };
+asyncLocalStorage.enterWith(store);
+asyncLocalStorage.getStore(); // Returns the store object
+someAsyncOperation(() => {
+  asyncLocalStorage.getStore(); // Returns the same object
+});
+```
+
+This transition will continue for the _entire_ synchronous execution.
+This means that if, for example, the context is entered within an event
+handler subsequent event handlers will also run within that context unless
+specifically bound to another context with an `AsyncResource`.
+
+```js
+const store = { id: 1 };
+
+emitter.on('my-event', () => {
+  asyncLocalStorage.enterWith(store);
+});
+emitter.on('my-event', () => {
+  asyncLocalStorage.getStore(); // Returns the same object
+});
+
+asyncLocalStorage.getStore(); // Returns undefined
+emitter.emit('my-event');
+asyncLocalStorage.getStore(); // Returns the same object
+```
+
 ### `asyncLocalStorage.run(store, callback[, ...args])`
 <!-- YAML
 added: REPLACEME
diff --git a/lib/async_hooks.js b/lib/async_hooks.js
index 47deffe84e233f..cb017222a9bce2 100644
--- a/lib/async_hooks.js
+++ b/lib/async_hooks.js
@@ -246,7 +246,7 @@ class AsyncLocalStorage {
     }
   }
 
-  _enter(store) {
+  enterWith(store) {
     if (!this.enabled) {
       this.enabled = true;
       storageList.push(this);
@@ -259,7 +259,7 @@ class AsyncLocalStorage {
   runSyncAndReturn(store, callback, ...args) {
     const resource = executionAsyncResource();
     const outerStore = resource[this.kResourceStore];
-    this._enter(store);
+    this.enterWith(store);
     try {
       return callback(...args);
     } finally {
@@ -289,7 +289,7 @@ class AsyncLocalStorage {
   run(store, callback, ...args) {
     const resource = executionAsyncResource();
     const outerStore = resource[this.kResourceStore];
-    this._enter(store);
+    this.enterWith(store);
     process.nextTick(callback, ...args);
     resource[this.kResourceStore] = outerStore;
   }
diff --git a/test/async-hooks/test-async-local-storage-enter-with.js b/test/async-hooks/test-async-local-storage-enter-with.js
new file mode 100644
index 00000000000000..736dd83f853763
--- /dev/null
+++ b/test/async-hooks/test-async-local-storage-enter-with.js
@@ -0,0 +1,20 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+setImmediate(() => {
+  const store = { foo: 'bar' };
+  asyncLocalStorage.enterWith(store);
+
+  assert.strictEqual(asyncLocalStorage.getStore(), store);
+  setTimeout(() => {
+    assert.strictEqual(asyncLocalStorage.getStore(), store);
+  }, 10);
+});
+
+setTimeout(() => {
+  assert.strictEqual(asyncLocalStorage.getStore(), undefined);
+}, 10);