lib: add AsyncLocalStorage.bind() and .snapshot()

flakey5 · danielleadams · commit f1b79828bc42 · 2023-04-05T12:52:17.000-04:00
PR-URL: #46387 Reviewed-By: Stephen Belanger <admin@stephenbelanger.com> Reviewed-By: Chengzhong Wu <legendecas@gmail.com> Reviewed-By: Gerhard Stöbich <deb2001-github@yahoo.de>
diff --git a/doc/api/async_context.md b/doc/api/async_context.md
@@ -134,6 +134,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
@@ -279,6 +279,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
@@ -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
@@ -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);
