async_hooks: add SymbolDispose support to AsyncLocalStorage

.eslintrc.js

@@ -19,6 +19,7 @@ const hacks = [
   'eslint-plugin-markdown',
   '@babel/eslint-parser',
   '@babel/plugin-syntax-import-attributes',
+  '@babel/plugin-proposal-explicit-resource-management',
 ];
 Module._findPath = (request, paths, isMain) => {
   const r = ModuleFindPath(request, paths, isMain);
@@ -45,6 +46,7 @@ module.exports = {
     babelOptions: {
       plugins: [
         Module._findPath('@babel/plugin-syntax-import-attributes'),
+        Module._findPath('@babel/plugin-proposal-explicit-resource-management'),
       ],
     },
     requireConfigFile: false,

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,83 @@ 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]()`
+
+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 +1009,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,77 @@ 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, as async hooks are
+  // disabled in the InternalCallbackScope created in `node::StartExecution`.
+  // Run check in a `process.nextTick` callback instead. If a previously
+  // scheduled tick callback is not called yet, do not schedule a new one as it
+  // can not be canceled.
+  if (executionAsyncResource() === topLevelResource && !topLevelTickScheduled) {
+    process.nextTick(checkDisposables);
+    topLevelTickScheduled = true;
+  }
+}
+
+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 +401,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 +449,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/async_hooks.js

@@ -79,7 +79,10 @@ const active_hooks = {
   tmp_fields: null,
 };
 
-const { registerDestroyHook } = async_wrap;
+const {
+  registerDestroyHook,
+  topLevelResource,
+} = async_wrap;
 const { enqueueMicrotask } = internalBinding('task_queue');
 const { resource_symbol, owner_symbol } = internalBinding('symbols');
 
@@ -137,8 +140,6 @@ function callbackTrampoline(asyncId, resource, cb, ...args) {
   return result;
 }
 
-const topLevelResource = {};
-
 function executionAsyncResource() {
   // Indicate to the native layer that this function is likely to be used,
   // in which case it will inform JS about the current async resource via
@@ -597,6 +598,7 @@ module.exports = {
   clearAsyncIdStack,
   hasAsyncIdStack,
   executionAsyncResource,
+  topLevelResource,
   // Internal Embedder API
   newAsyncId,
   getOrSetAsyncId,

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',

src/async_wrap.cc

@@ -412,6 +412,9 @@ void AsyncWrap::CreatePerContextProperties(Local<Object> target,
                          "execution_async_resources",
                          env->async_hooks()->js_execution_async_resources());
 
+  FORCE_SET_TARGET_FIELD(
+      target, "topLevelResource", realm->async_hooks_top_level_resource());
+
   target->Set(context,
               env->async_ids_stack_string(),
               env->async_hooks()->async_ids_stack().GetJSArray()).Check();

src/env_properties.h

@@ -404,6 +404,7 @@
   V(async_hooks_destroy_function, v8::Function)                                \
   V(async_hooks_init_function, v8::Function)                                   \
   V(async_hooks_promise_resolve_function, v8::Function)                        \
+  V(async_hooks_top_level_resource, v8::Object)                                \
   V(buffer_prototype_object, v8::Object)                                       \
   V(crypto_key_object_constructor, v8::Function)                               \
   V(crypto_key_object_private_constructor, v8::Function)                       \

src/node.cc

@@ -311,8 +311,8 @@ std::optional<StartExecutionCallbackInfo> CallbackInfoFromArray(
 MaybeLocal<Value> StartExecution(Environment* env, StartExecutionCallback cb) {
   InternalCallbackScope callback_scope(
       env,
-      Object::New(env->isolate()),
-      { 1, 0 },
+      env->principal_realm()->async_hooks_top_level_resource(),
+      {1, 0},
       InternalCallbackScope::kSkipAsyncHooks);
 
   // Only snapshot builder or embedder applications set the

src/node_realm.cc

@@ -81,6 +81,9 @@ void Realm::CreateProperties() {
   Local<Object> process_object =
       node::CreateProcessObject(this).FromMaybe(Local<Object>());
   set_process_object(process_object);
+
+  set_async_hooks_top_level_resource(
+      Object::New(isolate_, v8::Null(isolate_), nullptr, nullptr, 0));
 }
 
 RealmSerializeInfo Realm::Serialize(SnapshotCreator* creator) {