1 files changed, 263 insertions, 78 deletions

diff --git a/ext/node/polyfills/async_hooks.ts b/ext/node/polyfills/async_hooks.ts
index ad720d936..f94b8d2c6 100644
--- a/ext/node/polyfills/async_hooks.ts
+++ b/ext/node/polyfills/async_hooks.ts
@@ -1,34 +1,191 @@
 // Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
 // Copyright Joyent and Node contributors. All rights reserved. MIT license.
 
+// This implementation is inspired by "workerd" AsyncLocalStorage implementation:
+// https://github.com/cloudflare/workerd/blob/77fd0ed6ddba184414f0216508fc62b06e716cab/src/workerd/api/node/async-hooks.c++#L9
+
 // TODO(petamoriken): enable prefer-primordials for node polyfills
 // deno-lint-ignore-file prefer-primordials
 
-import { primordials } from "ext:core/mod.js";
-import {
-  AsyncVariable,
-  getAsyncContext,
-  setAsyncContext,
-} from "ext:runtime/01_async_context.js";
+import { core } from "ext:core/mod.js";
+import { op_node_is_promise_rejected } from "ext:core/ops";
 import { validateFunction } from "ext:deno_node/internal/validators.mjs";
 import { newAsyncId } from "ext:deno_node/internal/async_hooks.ts";
 
-const {
-  ObjectDefineProperties,
-  ReflectApply,
-  FunctionPrototypeBind,
-  ArrayPrototypeUnshift,
-  ObjectFreeze,
-} = primordials;
+function assert(cond: boolean) {
+  if (!cond) throw new Error("Assertion failed");
+}
+const asyncContextStack: AsyncContextFrame[] = [];
+
+function pushAsyncFrame(frame: AsyncContextFrame) {
+  asyncContextStack.push(frame);
+}
+
+function popAsyncFrame() {
+  if (asyncContextStack.length > 0) {
+    asyncContextStack.pop();
+  }
+}
+
+let rootAsyncFrame: AsyncContextFrame | undefined = undefined;
+let promiseHooksSet = false;
+
+const asyncContext = Symbol("asyncContext");
+
+function setPromiseHooks() {
+  if (promiseHooksSet) {
+    return;
+  }
+  promiseHooksSet = true;
+
+  const init = (promise: Promise<unknown>) => {
+    const currentFrame = AsyncContextFrame.current();
+    if (!currentFrame.isRoot()) {
+      if (typeof promise[asyncContext] !== "undefined") {
+        throw new Error("Promise already has async context");
+      }
+      AsyncContextFrame.attachContext(promise);
+    }
+  };
+  const before = (promise: Promise<unknown>) => {
+    const maybeFrame = promise[asyncContext];
+    if (maybeFrame) {
+      pushAsyncFrame(maybeFrame);
+    } else {
+      pushAsyncFrame(AsyncContextFrame.getRootAsyncContext());
+    }
+  };
+  const after = (promise: Promise<unknown>) => {
+    popAsyncFrame();
+    if (!op_node_is_promise_rejected(promise)) {
+      // @ts-ignore promise async context
+      promise[asyncContext] = undefined;
+    }
+  };
+  const resolve = (promise: Promise<unknown>) => {
+    const currentFrame = AsyncContextFrame.current();
+    if (
+      !currentFrame.isRoot() && op_node_is_promise_rejected(promise) &&
+      typeof promise[asyncContext] === "undefined"
+    ) {
+      AsyncContextFrame.attachContext(promise);
+    }
+  };
+
+  core.setPromiseHooks(init, before, after, resolve);
+}
+
+class AsyncContextFrame {
+  storage: StorageEntry[];
+  constructor(
+    maybeParent?: AsyncContextFrame | null,
+    maybeStorageEntry?: StorageEntry | null,
+    isRoot = false,
+  ) {
+    this.storage = [];
+
+    setPromiseHooks();
+
+    const propagate = (parent: AsyncContextFrame) => {
+      parent.storage = parent.storage.filter((entry) => !entry.key.isDead());
+      parent.storage.forEach((entry) => this.storage.push(entry.clone()));
+
+      if (maybeStorageEntry) {
+        const existingEntry = this.storage.find((entry) =>
+          entry.key === maybeStorageEntry.key
+        );
+        if (existingEntry) {
+          existingEntry.value = maybeStorageEntry.value;
+        } else {
+          this.storage.push(maybeStorageEntry);
+        }
+      }
+    };
+
+    if (!isRoot) {
+      if (maybeParent) {
+        propagate(maybeParent);
+      } else {
+        propagate(AsyncContextFrame.current());
+      }
+    }
+  }
+
+  static tryGetContext(promise: Promise<unknown>) {
+    // @ts-ignore promise async context
+    return promise[asyncContext];
+  }
+
+  static attachContext(promise: Promise<unknown>) {
+    // @ts-ignore promise async context
+    promise[asyncContext] = AsyncContextFrame.current();
+  }
+
+  static getRootAsyncContext() {
+    if (typeof rootAsyncFrame !== "undefined") {
+      return rootAsyncFrame;
+    }
+
+    rootAsyncFrame = new AsyncContextFrame(null, null, true);
+    return rootAsyncFrame;
+  }
+
+  static current() {
+    if (asyncContextStack.length === 0) {
+      return AsyncContextFrame.getRootAsyncContext();
+    }
+
+    return asyncContextStack[asyncContextStack.length - 1];
+  }
+
+  static create(
+    maybeParent?: AsyncContextFrame | null,
+    maybeStorageEntry?: StorageEntry | null,
+  ) {
+    return new AsyncContextFrame(maybeParent, maybeStorageEntry);
+  }
+
+  static wrap(
+    fn: () => unknown,
+    maybeFrame: AsyncContextFrame | undefined,
+    // deno-lint-ignore no-explicit-any
+    thisArg: any,
+  ) {
+    // deno-lint-ignore no-explicit-any
+    return (...args: any) => {
+      const frame = maybeFrame || AsyncContextFrame.current();
+      Scope.enter(frame);
+      try {
+        return fn.apply(thisArg, args);
+      } finally {
+        Scope.exit();
+      }
+    };
+  }
+
+  get(key: StorageKey) {
+    assert(!key.isDead());
+    this.storage = this.storage.filter((entry) => !entry.key.isDead());
+    const entry = this.storage.find((entry) => entry.key === key);
+    if (entry) {
+      return entry.value;
+    }
+    return undefined;
+  }
+
+  isRoot() {
+    return AsyncContextFrame.getRootAsyncContext() == this;
+  }
+}
 
 export class AsyncResource {
+  frame: AsyncContextFrame;
   type: string;
-  #snapshot: unknown;
   #asyncId: number;
 
   constructor(type: string) {
     this.type = type;
-    this.#snapshot = getAsyncContext();
+    this.frame = AsyncContextFrame.current();
     this.#asyncId = newAsyncId();
   }
 
@@ -41,38 +198,35 @@ export class AsyncResource {
     thisArg: unknown,
     ...args: unknown[]
   ) {
-    const previousContext = getAsyncContext();
+    Scope.enter(this.frame);
+
     try {
-      setAsyncContext(this.#snapshot);
-      return ReflectApply(fn, thisArg, args);
+      return fn.apply(thisArg, args);
     } finally {
-      setAsyncContext(previousContext);
+      Scope.exit();
     }
   }
 
   emitDestroy() {}
 
-  bind(fn: (...args: unknown[]) => unknown, thisArg) {
+  bind(fn: (...args: unknown[]) => unknown, thisArg = this) {
     validateFunction(fn, "fn");
-    let bound;
-    if (thisArg === undefined) {
-      // deno-lint-ignore no-this-alias
-      const resource = this;
-      bound = function (...args) {
-        ArrayPrototypeUnshift(args, fn, this);
-        return ReflectApply(resource.runInAsyncScope, resource, args);
-      };
-    } else {
-      bound = FunctionPrototypeBind(this.runInAsyncScope, this, fn, thisArg);
-    }
-    ObjectDefineProperties(bound, {
+    const frame = AsyncContextFrame.current();
+    const bound = AsyncContextFrame.wrap(fn, frame, thisArg);
+
+    Object.defineProperties(bound, {
       "length": {
-        __proto__: null,
         configurable: true,
         enumerable: false,
         value: fn.length,
         writable: false,
       },
+      "asyncResource": {
+        configurable: true,
+        enumerable: true,
+        value: this,
+        writable: true,
+      },
     });
     return bound;
   }
@@ -82,54 +236,95 @@ export class AsyncResource {
     type?: string,
     thisArg?: AsyncResource,
   ) {
-    type = type || fn.name || "bound-anonymous-fn";
-    return (new AsyncResource(type)).bind(fn, thisArg);
+    type = type || fn.name;
+    return (new AsyncResource(type || "AsyncResource")).bind(fn, thisArg);
+  }
+}
+
+class Scope {
+  static enter(maybeFrame?: AsyncContextFrame) {
+    if (maybeFrame) {
+      pushAsyncFrame(maybeFrame);
+    } else {
+      pushAsyncFrame(AsyncContextFrame.getRootAsyncContext());
+    }
+  }
+
+  static exit() {
+    popAsyncFrame();
+  }
+}
+
+class StorageEntry {
+  key: StorageKey;
+  value: unknown;
+  constructor(key: StorageKey, value: unknown) {
+    this.key = key;
+    this.value = value;
+  }
+
+  clone() {
+    return new StorageEntry(this.key, this.value);
+  }
+}
+
+class StorageKey {
+  #dead = false;
+
+  reset() {
+    this.#dead = true;
+  }
+
+  isDead() {
+    return this.#dead;
   }
 }
 
+const fnReg = new FinalizationRegistry((key: StorageKey) => {
+  key.reset();
+});
+
 export class AsyncLocalStorage {
-  #variable = new AsyncVariable();
-  enabled = false;
+  #key;
+
+  constructor() {
+    this.#key = new StorageKey();
+    fnReg.register(this, this.#key);
+  }
 
   // deno-lint-ignore no-explicit-any
   run(store: any, callback: any, ...args: any[]): any {
-    this.enabled = true;
-    const previous = this.#variable.enter(store);
+    const frame = AsyncContextFrame.create(
+      null,
+      new StorageEntry(this.#key, store),
+    );
+    Scope.enter(frame);
+    let res;
     try {
-      return ReflectApply(callback, null, args);
+      res = callback(...args);
     } finally {
-      setAsyncContext(previous);
+      Scope.exit();
     }
+    return res;
   }
 
   // deno-lint-ignore no-explicit-any
   exit(callback: (...args: unknown[]) => any, ...args: any[]): any {
-    if (!this.enabled) {
-      return ReflectApply(callback, null, args);
-    }
-    this.enabled = false;
-    try {
-      return ReflectApply(callback, null, args);
-    } finally {
-      this.enabled = true;
-    }
+    return this.run(undefined, callback, args);
   }
 
   // deno-lint-ignore no-explicit-any
   getStore(): any {
-    if (!this.enabled) {
-      return undefined;
-    }
-    return this.#variable.get();
+    const currentFrame = AsyncContextFrame.current();
+    return currentFrame.get(this.#key);
   }
 
   enterWith(store: unknown) {
-    this.enabled = true;
-    this.#variable.enter(store);
-  }
-
-  disable() {
-    this.enabled = false;
+    const frame = AsyncContextFrame.create(
+      null,
+      new StorageEntry(this.#key, store),
+    );
+    Scope.enter(frame);
   }
 
   static bind(fn: (...args: unknown[]) => unknown) {
@@ -140,24 +335,14 @@ export class AsyncLocalStorage {
     return AsyncLocalStorage.bind((
       cb: (...args: unknown[]) => unknown,
       ...args: unknown[]
-    ) => ReflectApply(cb, null, args));
+    ) => cb(...args));
   }
 }
 
 export function executionAsyncId() {
-  return 0;
-}
-
-export function triggerAsyncId() {
-  return 0;
+  return 1;
 }
 
-export function executionAsyncResource() {
-  return {};
-}
-
-export const asyncWrapProviders = ObjectFreeze({ __proto__: null });
-
 class AsyncHook {
   enable() {
   }
@@ -170,12 +355,12 @@ export function createHook() {
   return new AsyncHook();
 }
 
+// Placing all exports down here because the exported classes won't export
+// otherwise.
 export default {
-  AsyncLocalStorage,
-  createHook,
-  executionAsyncId,
-  triggerAsyncId,
-  executionAsyncResource,
-  asyncWrapProviders,
+  // Embedder API
   AsyncResource,
+  executionAsyncId,
+  createHook,
+  AsyncLocalStorage,
 };