From a3d4c950d63f5a81fc802324a9aab42057343108 Mon Sep 17 00:00:00 2001
From: flakey5 <73616808+flakey5@users.noreply.github.com>
Date: Fri, 27 Jan 2023 17:12:00 -0800
Subject: [PATCH] lib: add AsyncLocalStorage.bind() and .snapshot()

---
 doc/api/async_context.md                      | 50 +++++++++++++++++++
 lib/async_hooks.js                            |  8 +++
 .../parallel/test-async-local-storage-bind.js | 17 +++++++
 .../test-async-local-storage-snapshot.js      | 16 ++++++
 4 files changed, 91 insertions(+)
 create mode 100644 test/parallel/test-async-local-storage-bind.js
 create mode 100644 test/parallel/test-async-local-storage-snapshot.js

diff --git a/doc/api/async_context.md b/doc/api/async_context.md
index c3f287e690955c..f9896c0a2c692c 100644
--- a/doc/api/async_context.md
+++ b/doc/api/async_context.md
@@ -136,6 +136,56 @@ changes:
 Creates a new instance of `AsyncLocalStorage`. Store is only provided within a
 `run()` call or after an `enterWith()` call.
 
+### Static method: `AsyncLocalStorage.bind(fn)`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+* `fn` {Function} The function to bind to the current execution context.
+* Returns: {Function} A new function that calls `fn` within the captured
+  execution context.
+
+Binds the given function to the current execution context.
+
+### Static method: `AsyncLocalStorage.snapshot()`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+* Returns: {Function} A new function with the signature
+  `(fn: (...args) : R, ...args) : R`.
+
+Captures the current execution context and returns a function that accepts a
+function as an argument. Whenever the returned function is called, it
+calls the function passed to it within the captured context.
+
+```js
+const asyncLocalStorage = new AsyncLocalStorage();
+const runInAsyncScope = asyncLocalStorage.run(123, () => asyncLocalStorage.snapshot());
+const result = asyncLocalStorage.run(321, () => runInAsyncScope(() => asyncLocalStorage.getStore()));
+console.log(result);  // returns 123
+```
+
+AsyncLocalStorage.snapshot() can replace the use of AsyncResource for simple
+async context tracking purposes, for example:
+
+```js
+class Foo {
+  #runInAsyncScope = AsyncLocalStorage.snapshot();
+
+  get() { return this.#runInAsyncScope(() => asyncLocalStorage.getStore()); }
+}
+
+const foo = asyncLocalStorage.run(123, () => new Foo());
+console.log(asyncLocalStorage.run(321, () => foo.get())); // returns 123
+```
+
 ### `asyncLocalStorage.disable()`
 
 <!-- YAML
diff --git a/lib/async_hooks.js b/lib/async_hooks.js
index 23efdfabdb8c6e..552eda0bb681be 100644
--- a/lib/async_hooks.js
+++ b/lib/async_hooks.js
@@ -287,6 +287,14 @@ class AsyncLocalStorage {
     this.enabled = false;
   }
 
+  static bind(fn) {
+    return AsyncResource.bind(fn);
+  }
+
+  static snapshot() {
+    return AsyncLocalStorage.bind((cb, ...args) => cb(...args));
+  }
+
   disable() {
     if (this.enabled) {
       this.enabled = false;
diff --git a/test/parallel/test-async-local-storage-bind.js b/test/parallel/test-async-local-storage-bind.js
new file mode 100644
index 00000000000000..d8d4c4599826f9
--- /dev/null
+++ b/test/parallel/test-async-local-storage-bind.js
@@ -0,0 +1,17 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+[1, false, '', {}, []].forEach((i) => {
+  assert.throws(() => AsyncLocalStorage.bind(i), {
+    code: 'ERR_INVALID_ARG_TYPE'
+  });
+});
+
+const fn = common.mustCall(AsyncLocalStorage.bind(() => 123));
+assert.strictEqual(fn(), 123);
+
+const fn2 = AsyncLocalStorage.bind(common.mustCall((arg) => assert.strictEqual(arg, 'test')));
+fn2('test');
diff --git a/test/parallel/test-async-local-storage-snapshot.js b/test/parallel/test-async-local-storage-snapshot.js
new file mode 100644
index 00000000000000..63e47ba3ce0ac0
--- /dev/null
+++ b/test/parallel/test-async-local-storage-snapshot.js
@@ -0,0 +1,16 @@
+'use strict';
+
+const common = require('../common');
+const { strictEqual } = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+const runInAsyncScope =
+  asyncLocalStorage.run(123, common.mustCall(() => AsyncLocalStorage.snapshot()));
+const result =
+  asyncLocalStorage.run(321, common.mustCall(() => {
+    return runInAsyncScope(() => {
+      return asyncLocalStorage.getStore();
+    });
+  }));
+strictEqual(result, 123);