From b0bc07a0b8ae02ff3fc9746e6cc3ffc19df97ba5 Mon Sep 17 00:00:00 2001
From: Bryan English <bryan@bryanenglish.com>
Date: Thu, 1 May 2025 14:06:19 -0400
Subject: [PATCH] async_hooks: add use() method to AsyncLocalStorage

This provides a way to use the `using` syntax (when available) to manage
AsyncLocalStorage contexts, as an alternative to `run()`.
---
 doc/api/async_context.md                      | 37 +++++++++++++++++++
 .../async_context_frame.js                    | 18 +++++++++
 .../async_local_storage/async_hooks.js        | 31 ++++++++++++++++
 test/parallel/test-async-local-storage-use.js | 25 +++++++++++++
 4 files changed, 111 insertions(+)
 create mode 100644 test/parallel/test-async-local-storage-use.js

diff --git a/doc/api/async_context.md b/doc/api/async_context.md
index c60f76c746fed1..04d94d59473023 100644
--- a/doc/api/async_context.md
+++ b/doc/api/async_context.md
@@ -347,6 +347,43 @@ try {
 }
 ```
 
+### `asyncLocalStorage.use(store)`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+* `store` {any}
+* Returns: {Disposable} A disposable object.
+
+Transitions into the given context, and transitions back into the previous
+context when the returned disposable object is disposed. The store is
+accessible to any asynchronous operations created before disposal.
+
+Example:
+
+```js
+const store1 = { id: 1 };
+const store2 = { id: 2 };
+
+function inner() {
+  // Once `using` syntax is supported, you can use that here, and omit the
+  // dispose call at the end of this function.
+  const disposable = asyncLocalStorage.use(store);
+  asyncLocalStorage.getStore(); // Returns store2
+  setTimeout(() => {
+    asyncLocalStorage.getStore(); // Returns store2
+  }, 200);
+  disposable[Symbol.dispose]();
+}
+
+asyncLocalStorage.run(store1, () => {
+  asyncLocalStorage.getStore(); // Returns store1
+  inner();
+  asyncLocalStorage.getStore(); // Returns store1
+});
+```
+
 ### `asyncLocalStorage.exit(callback[, ...args])`
 
 <!-- YAML
diff --git a/lib/internal/async_local_storage/async_context_frame.js b/lib/internal/async_local_storage/async_context_frame.js
index 13b769a4642aa1..4a7214bdca5085 100644
--- a/lib/internal/async_local_storage/async_context_frame.js
+++ b/lib/internal/async_local_storage/async_context_frame.js
@@ -2,6 +2,7 @@
 
 const {
   ReflectApply,
+  SymbolDispose,
 } = primordials;
 
 const {
@@ -11,6 +12,19 @@ const {
 const AsyncContextFrame = require('internal/async_context_frame');
 const { AsyncResource } = require('async_hooks');
 
+class DisposableStore {
+  #oldFrame = undefined;
+
+  constructor(store, storage) {
+    this.#oldFrame = AsyncContextFrame.current();
+    storage.enterWith(store);
+  }
+
+  [SymbolDispose]() {
+    AsyncContextFrame.set(this.#oldFrame);
+  }
+}
+
 class AsyncLocalStorage {
   #defaultValue = undefined;
   #name = undefined;
@@ -62,6 +76,10 @@ class AsyncLocalStorage {
     }
   }
 
+  use(data) {
+    return new DisposableStore(data, this);
+  }
+
   exit(fn, ...args) {
     return this.run(undefined, fn, ...args);
   }
diff --git a/lib/internal/async_local_storage/async_hooks.js b/lib/internal/async_local_storage/async_hooks.js
index d227549412bf61..b1fdf7ad447efb 100644
--- a/lib/internal/async_local_storage/async_hooks.js
+++ b/lib/internal/async_local_storage/async_hooks.js
@@ -7,6 +7,7 @@ const {
   ObjectIs,
   ReflectApply,
   Symbol,
+  SymbolDispose,
 } = primordials;
 
 const {
@@ -30,6 +31,31 @@ const storageHook = createHook({
   },
 });
 
+class DisposableStore {
+  #oldStore = undefined;
+  #resource = undefined;
+  #kResourceStore = undefined;
+
+  constructor(store, storage) {
+    this.#oldStore = storage.getStore();
+
+    if (ObjectIs(store, this.#oldStore)) {
+      return;
+    }
+
+    this.#kResourceStore = storage.kResourceStore;
+    this.#resource = executionAsyncResource();
+    this.#resource[this.#kResourceStore] = store;
+  }
+
+  [SymbolDispose]() {
+    if (this.#resource === undefined) {
+      return;
+    }
+    this.#resource[this.#kResourceStore] = this.#oldStore;
+  }
+}
+
 class AsyncLocalStorage {
   #defaultValue = undefined;
   #name = undefined;
@@ -120,6 +146,11 @@ class AsyncLocalStorage {
     }
   }
 
+  use(store) {
+    this._enable();
+    return new DisposableStore(store, this);
+  }
+
   exit(callback, ...args) {
     if (!this.enabled) {
       return ReflectApply(callback, null, args);
diff --git a/test/parallel/test-async-local-storage-use.js b/test/parallel/test-async-local-storage-use.js
new file mode 100644
index 00000000000000..c355afa26cb3c4
--- /dev/null
+++ b/test/parallel/test-async-local-storage-use.js
@@ -0,0 +1,25 @@
+'use strict';
+
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const storage = new AsyncLocalStorage();
+
+const store1 = {};
+const store2 = {};
+const store3 = {};
+
+function inner() {
+  // TODO(bengl): Once `using` is supported use that here and don't call
+  // dispose manually later.
+  const disposable = storage.use(store2);
+  assert.strictEqual(storage.getStore(), store2);
+  storage.enterWith(store3);
+  disposable[Symbol.dispose]();
+}
+
+storage.enterWith(store1);
+assert.strictEqual(storage.getStore(), store1);
+inner();
+assert.strictEqual(storage.getStore(), store1);