async_hooks: add using scopes to AsyncLocalStorage

Adds support for using scope = storage.withScope(data) to do
the equivalent of a storage.run(data, fn) with using syntax.
This enables avoiding unnecessary closures.

PR-URL: #61674
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: Gerhard Stöbich <deb2001-github@yahoo.de>
Reviewed-By: Bryan English <bryan@bryanenglish.com>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Chengzhong Wu <legendecas@gmail.com>

Author: Qard

doc/api/async_context.md

@@ -386,6 +386,110 @@ try {
 }
 ```
 
+### `asyncLocalStorage.withScope(store)`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+* `store` {any}
+* Returns: {RunScope}
+
+Creates a disposable scope that enters the given store and automatically
+restores the previous store value when the scope is disposed. This method is
+designed to work with JavaScript's explicit resource management (`using` syntax).
+
+Example:
+
+```mjs
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+{
+  using _ = asyncLocalStorage.withScope('my-store');
+  console.log(asyncLocalStorage.getStore()); // Prints: my-store
+}
+
+console.log(asyncLocalStorage.getStore()); // Prints: undefined
+```
+
+```cjs
+const { AsyncLocalStorage } = require('node:async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+{
+  using _ = asyncLocalStorage.withScope('my-store');
+  console.log(asyncLocalStorage.getStore()); // Prints: my-store
+}
+
+console.log(asyncLocalStorage.getStore()); // Prints: undefined
+```
+
+The `withScope()` method is particularly useful for managing context in
+synchronous code where you want to ensure the previous store value is restored
+when exiting a block, even if an error is thrown.
+
+```mjs
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+try {
+  using _ = asyncLocalStorage.withScope('my-store');
+  console.log(asyncLocalStorage.getStore()); // Prints: my-store
+  throw new Error('test');
+} catch (e) {
+  // Store is automatically restored even after error
+  console.log(asyncLocalStorage.getStore()); // Prints: undefined
+}
+```
+
+```cjs
+const { AsyncLocalStorage } = require('node:async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+try {
+  using _ = asyncLocalStorage.withScope('my-store');
+  console.log(asyncLocalStorage.getStore()); // Prints: my-store
+  throw new Error('test');
+} catch (e) {
+  // Store is automatically restored even after error
+  console.log(asyncLocalStorage.getStore()); // Prints: undefined
+}
+```
+
+**Important:** When using `withScope()` in async functions before the first
+`await`, be aware that the scope change will affect the caller's context. The
+synchronous portion of an async function (before the first `await`) runs
+immediately when called, and when it reaches the first `await`, it returns the
+promise to the caller. At that point, the scope change becomes visible in the
+caller's context and will persist in subsequent synchronous code until something
+else changes the scope value. For async operations, prefer using `run()` which
+properly isolates context across async boundaries.
+
+```mjs
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+async function example() {
+  using _ = asyncLocalStorage.withScope('my-store');
+  console.log(asyncLocalStorage.getStore()); // Prints: my-store
+  await someAsyncOperation(); // Function pauses here and returns promise
+  console.log(asyncLocalStorage.getStore()); // Prints: my-store
+}
+
+// Calling without await
+example(); // Synchronous portion runs, then pauses at first await
+// After the promise is returned, the scope 'my-store' is now active in caller!
+console.log(asyncLocalStorage.getStore()); // Prints: my-store (unexpected!)
+```
+
 ### Usage with `async/await`
 
 If, within an async function, only one `await` call is to run within a context,
@@ -420,6 +524,64 @@ of `asyncLocalStorage.getStore()` after the calls you suspect are responsible
 for the loss. When the code logs `undefined`, the last callback called is
 probably responsible for the context loss.
 
+## Class: `RunScope`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+A disposable scope returned by [`asyncLocalStorage.withScope()`][] that
+automatically restores the previous store value when disposed. This class
+implements the [Explicit Resource Management][] protocol and is designed to work
+with JavaScript's `using` syntax.
+
+The scope automatically restores the previous store value when the `using` block
+exits, whether through normal completion or by throwing an error.
+
+### `scope.dispose()`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+Explicitly ends the scope and restores the previous store value. This method
+is idempotent: calling it multiple times has the same effect as calling it once.
+
+The `[Symbol.dispose]()` method defers to `dispose()`.
+
+If `withScope()` is called without the `using` keyword, `dispose()` must be
+called manually to restore the previous store value. Forgetting to call
+`dispose()` will cause the store value to persist for the remainder of the
+current execution context:
+
+```mjs
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+const storage = new AsyncLocalStorage();
+
+// Without using, the scope must be disposed manually
+const scope = storage.withScope('my-store');
+// storage.getStore() === 'my-store' here
+
+scope.dispose(); // Restore previous value
+// storage.getStore() === undefined here
+```
+
+```cjs
+const { AsyncLocalStorage } = require('node:async_hooks');
+
+const storage = new AsyncLocalStorage();
+
+// Without using, the scope must be disposed manually
+const scope = storage.withScope('my-store');
+// storage.getStore() === 'my-store' here
+
+scope.dispose(); // Restore previous value
+// storage.getStore() === undefined here
+```
+
 ## Class: `AsyncResource`
 
 <!-- YAML
@@ -905,8 +1067,10 @@ const server = createServer((req, res) => {
 }).listen(3000);
 ```
 
+[Explicit Resource Management]: https://github.com/tc39/proposal-explicit-resource-management
 [`AsyncResource`]: #class-asyncresource
 [`EventEmitter`]: events.md#class-eventemitter
 [`Stream`]: stream.md#stream
 [`Worker`]: worker_threads.md#class-worker
+[`asyncLocalStorage.withScope()`]: #asynclocalstoragewithscopestore
 [`util.promisify()`]: util.md#utilpromisifyoriginal

lib/diagnostics_channel.js

@@ -372,16 +372,16 @@ class TracingChannel {
 
     const { start, end, asyncStart, asyncEnd, error } = this;
 
-    function reject(err) {
+    function onReject(err) {
       context.error = err;
       error.publish(context);
       asyncStart.publish(context);
       // TODO: Is there a way to have asyncEnd _after_ the continuation?
       asyncEnd.publish(context);
     }
 
-    function rejectAndRethrow(err) {
-      reject(err);
+    function onRejectWithRethrow(err) {
+      onReject(err);
       throw err;
     }
 
@@ -403,14 +403,16 @@ class TracingChannel {
           context.result = result;
           return result;
         }
+        // isPromise() matches sub-classes, but we need to match only direct
+        // instances of the native Promise type to safely use PromisePrototypeThen.
         if (isPromise(result) && result.constructor === Promise) {
-          return PromisePrototypeThen(result, resolve, rejectAndRethrow);
+          return PromisePrototypeThen(result, resolve, onRejectWithRethrow);
         }
         // For non-native thenables, subscribe to the result but return the
         // original thenable so the consumer can continue handling it directly.
         // Non-native thenables don't have unhandledRejection tracking, so
         // swallowing the rejection here doesn't change existing behaviour.
-        result.then(resolve, reject);
+        result.then(resolve, onReject);
         return result;
       } catch (err) {
         context.error = err;

lib/internal/async_local_storage/async_context_frame.js

@@ -12,6 +12,8 @@ const {
 const AsyncContextFrame = require('internal/async_context_frame');
 const { AsyncResource } = require('async_hooks');
 
+const RunScope = require('internal/async_local_storage/run_scope');
+
 class AsyncLocalStorage {
   #defaultValue = undefined;
   #name = undefined;
@@ -77,6 +79,10 @@ class AsyncLocalStorage {
     }
     return frame?.get(this);
   }
+
+  withScope(store) {
+    return new RunScope(this, store);
+  }
 }
 
 module.exports = AsyncLocalStorage;

lib/internal/async_local_storage/async_hooks.js

@@ -19,6 +19,8 @@ const {
   executionAsyncResource,
 } = require('async_hooks');
 
+const RunScope = require('internal/async_local_storage/run_scope');
+
 const storageList = [];
 const storageHook = createHook({
   init(asyncId, type, triggerAsyncId, resource) {
@@ -142,6 +144,10 @@ class AsyncLocalStorage {
     }
     return this.#defaultValue;
   }
+
+  withScope(store) {
+    return new RunScope(this, store);
+  }
 }
 
 module.exports = AsyncLocalStorage;

lib/internal/async_local_storage/run_scope.js

@@ -0,0 +1,31 @@
+'use strict';
+
+const {
+  SymbolDispose,
+} = primordials;
+
+class RunScope {
+  #storage;
+  #previousStore;
+  #disposed = false;
+
+  constructor(storage, store) {
+    this.#storage = storage;
+    this.#previousStore = storage.getStore();
+    storage.enterWith(store);
+  }
+
+  dispose() {
+    if (this.#disposed) {
+      return;
+    }
+    this.#disposed = true;
+    this.#storage.enterWith(this.#previousStore);
+  }
+
+  [SymbolDispose]() {
+    this.dispose();
+  }
+}
+
+module.exports = RunScope;