async_hooks: add SymbolDispose support to AsyncLocalStorage

This introduces explicit resource management support to
AsyncLocalStorage. In order to avoid unexpected leakage of an async
local storage value to outer scope, `asyncLocalStorage.run` requires a
callback to nest all sub-procedures in a new function scope. However,
when using `AsyncLocalStorage` with different function types, like
async generator, moving statements and expression to a new function
body could be evaluated in different ways.

For example, given an async generator:

```js
class C {
  async* foo() {
    await a();
    yield b();
    await c();
  }
}
```

Then, if we want to modify the async local storage value for `b()` and
`c()`, simply moving statements would not work:

```js
const storage = new AsyncLocalStorage();

class C {
  async* foo() {
    await a();
    storage.run(value, () => {
      yield b(); // syntax error, arrow function is not a generator
      await this.c(); // syntax error, arrow function is not async
    });
  }
}
```

Making the arrow function to be an async generator as well still
requires significant refactoring:

```js
const storage = new AsyncLocalStorage();

class C {
  async* foo() {
    await a();
    yield* storage.run(value, async function*() => {
      yield b(); // ok
      await this.c(); // reference error, this is undefined
    });
  }
}
```

This could be potentially more convenient with `using` declarations:

```js
const storage = new AsyncLocalStorage();
class C {
  async* foo() {
    await a();
    using _ = storage.disposableStore(value);
    yield b();
    await this.c();
  }
}
```

However, creating a disposable without a `using` declaration could
still be a problem leading to leakage of the async local storage
value. To help identifying such mis-use, an
`ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED` error is thrown if an
`AsyncLocalStorageDisposableStore` instance is not disposed at the end
of the synchronous execution.

Author: legendecas

doc/api/async_context.md

@@ -361,6 +361,57 @@ try {
 }
 ```
 
+### `asyncLocalStorage.disposableStore(store)`
+
+<!-- YAML
+added:
+ - REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+* `store` {any}
+* Returns: {AsyncLocalStorageDisposableStore}
+
+Transitions into the context until the returned
+[`AsyncLocalStorageDisposableStore`][] instance is disposed, and then persists
+the store through any following asynchronous calls.
+
+Example:
+
+```js
+const store = { id: 1 };
+{
+  // Replaces previous store with the given store object
+  using _ = asyncLocalStorage.disposableStore(store);
+  asyncLocalStorage.getStore(); // Returns the store object
+  someAsyncOperation(() => {
+    asyncLocalStorage.getStore(); // Returns the same object
+  });
+  // The store goes out of scope
+}
+asyncLocalStorage.getStore(); // Returns undefined
+```
+
+This transition will continue for the _entire_ synchronous execution scope
+until the returned `AsyncLocalStorageDisposableStore` instance is disposed.
+
+```js
+const store = { id: 1 };
+
+emitter.on('my-event', () => {
+  // Declare a disposable store with using statement.
+  using _ = asyncLocalStorage.disposableStore(store);
+});
+emitter.on('my-event', () => {
+  asyncLocalStorage.getStore(); // Returns undefined
+});
+
+asyncLocalStorage.getStore(); // Returns undefined
+emitter.emit('my-event');
+asyncLocalStorage.getStore(); // Returns undefined
+```
+
 ### Usage with `async/await`
 
 If, within an async function, only one `await` call is to run within a context,
@@ -395,6 +446,89 @@ 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: `AsyncLocalStorageDisposableStore`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+A utility class that enables `using` statement to dispose of a store value for
+an `AsyncLocalStorage` store when the current scope was exited.
+
+### `asyncLocalStorageDisposableStore[Symbol.dispose]()`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+Restores the previous store for the creation `AsyncLocalStorage`.
+
+### Usage with `AsyncLocalStorage`
+
+```js
+const http = require('node:http');
+const { AsyncLocalStorage } = require('node:async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+function logWithId(msg) {
+  const id = asyncLocalStorage.getStore();
+  console.log(`${id !== undefined ? id : '-'}:`, msg);
+}
+
+let idSeq = 0;
+http.createServer((req, res) => {
+  using _ = asyncLocalStorage.disposableStore(idSeq++);
+  logWithId('start');
+  // Imagine any chain of async operations here
+  setImmediate(() => {
+    logWithId('finish');
+    res.end();
+  });
+}).listen(8080);
+
+http.get('http://localhost:8080');
+http.get('http://localhost:8080');
+// Prints:
+//   0: start
+//   1: start
+//   0: finish
+//   1: finish
+```
+
+### Troubleshooting: Instance not disposed
+
+The instantiation of the [`AsyncLocalStorageDisposableStore`][] can be used
+without the using statement. This can be problematic that if, for example, the
+context is entered within an event handler and the
+`AsyncLocalStorageDisposableStore` instance is not disposed, 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('leaking-event', () => {
+  // Declare a disposable store with variable declaration and did not dispose it.
+  const _ = asyncLocalStorage.disposableStore(store);
+});
+emitter.on('leaking-event', () => {
+  asyncLocalStorage.getStore(); // Returns the store object
+});
+
+asyncLocalStorage.getStore(); // Returns undefined
+emitter.emit('my-event');
+asyncLocalStorage.getStore(); // Returns the store object
+```
+
+To hint such a misuse, an [`ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED`][] error is
+thrown if an `AsyncLocalStorageDisposableStore` instance is not disposed at the
+end of the current async resource scope
+
 ## Class: `AsyncResource`
 
 <!-- YAML
@@ -881,7 +1015,9 @@ const server = createServer((req, res) => {
 }).listen(3000);
 ```
 
+[`AsyncLocalStorageDisposableStore`]: #class-asynclocalstoragedisposablestore
 [`AsyncResource`]: #class-asyncresource
+[`ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED`]: errors.md#err_async_local_storage_not_disposed
 [`EventEmitter`]: events.md#class-eventemitter
 [`Stream`]: stream.md#stream
 [`Worker`]: worker_threads.md#class-worker

doc/api/errors.md

@@ -698,6 +698,13 @@ by the `node:assert` module.
 An attempt was made to register something that is not a function as an
 `AsyncHooks` callback.
 
+<a id="ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED"></a>
+
+### `ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED`
+
+An AsyncLocalStorage disposable store was not disposed when the current
+async resource scope was exited.
+
 <a id="ERR_ASYNC_TYPE"></a>
 
 ### `ERR_ASYNC_TYPE`

lib/async_hooks.js

@@ -11,12 +11,15 @@ const {
   ObjectDefineProperties,
   ObjectIs,
   ReflectApply,
+  SafeSet,
   Symbol,
+  SymbolDispose,
   ObjectFreeze,
 } = primordials;
 
 const {
   ERR_ASYNC_CALLBACK,
+  ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED,
   ERR_ASYNC_TYPE,
   ERR_INVALID_ASYNC_ID,
 } = require('internal/errors').codes;
@@ -30,6 +33,8 @@ const {
 } = require('internal/validators');
 const internal_async_hooks = require('internal/async_hooks');
 
+const { triggerUncaughtException } = internalBinding('errors');
+
 // Get functions
 // For userland AsyncResources, make sure to emit a destroy event when the
 // resource gets gced.
@@ -45,6 +50,7 @@ const {
   disableHooks,
   updatePromiseHookMode,
   executionAsyncResource,
+  topLevelResource,
   // Internal Embedder API
   newAsyncId,
   getDefaultTriggerAsyncId,
@@ -281,6 +287,73 @@ const storageHook = createHook({
   },
 });
 
+let topLevelTickScheduled = false;
+let activeDisposableStores = new SafeSet();
+function checkDisposables() {
+  topLevelTickScheduled = false;
+  if (activeDisposableStores.size === 0) {
+    return;
+  }
+  const remainedStores = activeDisposableStores;
+  activeDisposableStores = new SafeSet();
+  const err = new ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED(remainedStores);
+  err.disposableStores = remainedStores;
+  triggerUncaughtException(err, false);
+}
+
+const disposableCheckHook = createHook({
+  after() {
+    checkDisposables();
+  },
+});
+
+function enableDisposablesCheck() {
+  if (activeDisposableStores.size !== 0) {
+    return;
+  }
+  disposableCheckHook.enable();
+  // There is no after hook for the top level resource. Run check in
+  // a `process.nextTick` callback.
+  if (executionAsyncResource() === topLevelResource && !topLevelTickScheduled) {
+    process.nextTick(checkDisposables);
+  }
+}
+
+function maybeDisableDisposablesCheck() {
+  if (activeDisposableStores.size !== 0) {
+    return;
+  }
+  disposableCheckHook.disable();
+}
+
+class AsyncLocalStorageDisposableStore {
+  #oldStore;
+  #asyncLocalStorage;
+  #entered = false;
+  constructor(asyncLocalStorage, store) {
+    this.store = store;
+    this.#asyncLocalStorage = asyncLocalStorage;
+    this.#oldStore = asyncLocalStorage.getStore();
+
+    asyncLocalStorage.enterWith(store);
+    this.#entered = true;
+
+    enableDisposablesCheck();
+    activeDisposableStores.add(this);
+  }
+
+  [SymbolDispose]() {
+    if (!this.#entered) {
+      return;
+    }
+    this.#asyncLocalStorage.enterWith(this.#oldStore);
+    this.#entered = false;
+
+    activeDisposableStores.delete(this);
+    maybeDisableDisposablesCheck();
+  }
+}
+
 class AsyncLocalStorage {
   constructor() {
     this.kResourceStore = Symbol('kResourceStore');
@@ -324,6 +397,11 @@ class AsyncLocalStorage {
   }
 
   enterWith(store) {
+    // Avoid creation of an AsyncResource if store is already active
+    if (ObjectIs(store, this.getStore())) {
+      return;
+    }
+
     this._enable();
     const resource = executionAsyncResource();
     resource[this.kResourceStore] = store;
@@ -367,6 +445,10 @@ class AsyncLocalStorage {
       return resource[this.kResourceStore];
     }
   }
+
+  disposableStore(store) {
+    return new AsyncLocalStorageDisposableStore(this, store);
+  }
 }
 
 // Placing all exports down here because the exported classes won't export

lib/internal/errors.js

@@ -1119,6 +1119,7 @@ E('ERR_AMBIGUOUS_ARGUMENT', 'The "%s" argument is ambiguous. %s', TypeError);
 E('ERR_ARG_NOT_ITERABLE', '%s must be iterable', TypeError);
 E('ERR_ASSERTION', '%s', Error);
 E('ERR_ASYNC_CALLBACK', '%s must be a function', TypeError);
+E('ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED', 'The AsyncLocalStorage store is not disposed of: %s', Error);
 E('ERR_ASYNC_TYPE', 'Invalid name for async "type": %s', TypeError);
 E('ERR_BROTLI_INVALID_PARAM', '%s is not a valid Brotli parameter', RangeError);
 E('ERR_BUFFER_OUT_OF_BOUNDS',

test/parallel/test-async-local-storage-disposable-check-top-level.js

@@ -0,0 +1,22 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+// This test verifies that an error is thrown when an
+// `asyncLocalStorage.disposableStore` created in the top level async scope is
+// not disposed.
+
+{
+  // If the disposable is not modified with `using` statement or being disposed after
+  // the current tick is finished, an error is thrown.
+  const disposable = asyncLocalStorage.disposableStore(123);
+  process.on('uncaughtException', common.mustCall((err) => {
+    assert.strictEqual(err.code, 'ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED');
+    assert.strictEqual(err.disposableStores.size, 1);
+    assert.ok(err.disposableStores.has(disposable));
+  }));
+}

test/parallel/test-async-local-storage-disposable-check.js

@@ -0,0 +1,28 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage, AsyncResource } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+const resource = new AsyncResource('test');
+
+// This test verifies that an error is thrown when an
+// `asyncLocalStorage.disposableStore` created in an async resource scope is
+// not disposed.
+
+try {
+  resource.runInAsyncScope(() => {
+    // If the disposable is not modified with `using` statement or being disposed after
+    // the current tick is finished, an error is thrown.
+    const disposable = asyncLocalStorage.disposableStore(123);
+
+    process.on('uncaughtException', common.mustCall((err) => {
+      assert.strictEqual(err.code, 'ERR_ASYNC_LOCAL_STORAGE_NOT_DISPOSED');
+      assert.strictEqual(err.disposableStores.size, 1);
+      assert.ok(err.disposableStores.has(disposable));
+    }));
+  });
+} catch {
+  assert.fail('unreachable');
+}

test/parallel/test-async-local-storage-disposable.js

@@ -0,0 +1,30 @@
+'use strict';
+
+require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+{
+  // This is equivalent to `using asyncLocalStorage.disposableStore(123)`.
+  const disposable = asyncLocalStorage.disposableStore(123);
+  assert.strictEqual(asyncLocalStorage.getStore(), 123);
+  disposable[Symbol.dispose]();
+}
+assert.strictEqual(asyncLocalStorage.getStore(), undefined);
+
+{
+  // [Symbol.dispose] should be a no-op once it has been called.
+  const disposable = asyncLocalStorage.disposableStore(123);
+  assert.strictEqual(asyncLocalStorage.getStore(), 123);
+  disposable[Symbol.dispose]();
+
+  const disposable2 = asyncLocalStorage.disposableStore(456);
+  // Call [Symbol.dispose] again
+  disposable[Symbol.dispose]();
+  assert.strictEqual(asyncLocalStorage.getStore(), 456);
+
+  disposable2[Symbol.dispose]();
+}
+assert.strictEqual(asyncLocalStorage.getStore(), undefined);