async_hooks: add use() method to AsyncLocalStorage

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

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);
  }

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);

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);