Revert "feat: async context" (#24856)

Reverts denoland/deno#24402

deno_web can't depend on code in runtime

8 files changed, 294 insertions, 208 deletions

diff --git a/ext/node/lib.rs b/ext/node/lib.rs
index 74c528b0c..a4a757996 100644
--- a/ext/node/lib.rs
+++ b/ext/node/lib.rs
@@ -181,6 +181,15 @@ fn op_node_build_os() -> String {
   env!("TARGET").split('-').nth(2).unwrap().to_string()
 }
 
+#[op2(fast)]
+fn op_node_is_promise_rejected(value: v8::Local<v8::Value>) -> bool {
+  let Ok(promise) = v8::Local::<v8::Promise>::try_from(value) else {
+    return false;
+  };
+
+  promise.state() == v8::PromiseState::Rejected
+}
+
 #[op2]
 #[string]
 fn op_npm_process_state(state: &mut OpState) -> Result<String, AnyError> {
@@ -341,6 +350,7 @@ deno_core::extension!(deno_node,
     ops::os::op_cpus<P>,
     ops::os::op_homedir<P>,
     op_node_build_os,
+    op_node_is_promise_rejected,
     op_npm_process_state,
     ops::require::op_require_init_paths,
     ops::require::op_require_node_module_paths<P>,
diff --git a/ext/node/polyfills/_next_tick.ts b/ext/node/polyfills/_next_tick.ts
index 5d895012e..5915c750e 100644
--- a/ext/node/polyfills/_next_tick.ts
+++ b/ext/node/polyfills/_next_tick.ts
@@ -5,10 +5,6 @@
 // deno-lint-ignore-file prefer-primordials
 
 import { core } from "ext:core/mod.js";
-import {
-  getAsyncContext,
-  setAsyncContext,
-} from "ext:runtime/01_async_context.js";
 
 import { validateFunction } from "ext:deno_node/internal/validators.mjs";
 import { _exiting } from "ext:deno_node/_process/exiting.ts";
@@ -17,7 +13,6 @@ import { FixedQueue } from "ext:deno_node/internal/fixed_queue.ts";
 interface Tock {
   callback: (...args: Array<unknown>) => void;
   args: Array<unknown>;
-  snapshot: unknown;
 }
 
 let nextTickEnabled = false;
@@ -28,7 +23,7 @@ export function enableNextTick() {
 const queue = new FixedQueue();
 
 export function processTicksAndRejections() {
-  let tock: Tock;
+  let tock;
   do {
     // deno-lint-ignore no-cond-assign
     while (tock = queue.shift()) {
@@ -36,11 +31,9 @@ export function processTicksAndRejections() {
       // const asyncId = tock[async_id_symbol];
       // emitBefore(asyncId, tock[trigger_async_id_symbol], tock);
 
-      const oldContext = getAsyncContext();
       try {
-        setAsyncContext(tock.snapshot);
-        const callback = tock.callback;
-        if (tock.args === undefined) {
+        const callback = (tock as Tock).callback;
+        if ((tock as Tock).args === undefined) {
           callback();
         } else {
           const args = (tock as Tock).args;
@@ -65,7 +58,6 @@ export function processTicksAndRejections() {
         // FIXME(bartlomieju): Deno currently doesn't support async hooks
         // if (destroyHooksExist())
         // emitDestroy(asyncId);
-        setAsyncContext(oldContext);
       }
 
       // FIXME(bartlomieju): Deno currently doesn't support async hooks
@@ -151,7 +143,6 @@ export function nextTick<T extends Array<unknown>>(
     // FIXME(bartlomieju): Deno currently doesn't support async hooks
     // [async_id_symbol]: asyncId,
     // [trigger_async_id_symbol]: triggerAsyncId,
-    snapshot: getAsyncContext(),
     callback,
     args: args_,
   };
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,
 };
diff --git a/ext/web/02_timers.js b/ext/web/02_timers.js
index a651df5a5..559147861 100644
--- a/ext/web/02_timers.js
+++ b/ext/web/02_timers.js
@@ -2,10 +2,6 @@
 
 import { core, primordials } from "ext:core/mod.js";
 import { op_defer, op_now } from "ext:core/ops";
-import {
-  getAsyncContext,
-  setAsyncContext,
-} from "ext:runtime/01_async_context.js";
 const {
   Uint8Array,
   Uint32Array,
@@ -37,16 +33,14 @@ function checkThis(thisArg) {
  * Call a callback function immediately.
  */
 function setImmediate(callback, ...args) {
-  const asyncContext = getAsyncContext();
-  return core.queueImmediate(() => {
-    const oldContext = getAsyncContext();
-    try {
-      setAsyncContext(asyncContext);
-      return ReflectApply(callback, globalThis, args);
-    } finally {
-      setAsyncContext(oldContext);
-    }
-  });
+  if (args.length > 0) {
+    const unboundCallback = callback;
+    callback = () => ReflectApply(unboundCallback, globalThis, args);
+  }
+
+  return core.queueImmediate(
+    callback,
+  );
 }
 
 /**
@@ -59,17 +53,10 @@ function setTimeout(callback, timeout = 0, ...args) {
     const unboundCallback = webidl.converters.DOMString(callback);
     callback = () => indirectEval(unboundCallback);
   }
-  const unboundCallback = callback;
-  const asyncContext = getAsyncContext();
-  callback = () => {
-    const oldContext = getAsyncContext();
-    try {
-      setAsyncContext(asyncContext);
-      ReflectApply(unboundCallback, globalThis, args);
-    } finally {
-      setAsyncContext(oldContext);
-    }
-  };
+  if (args.length > 0) {
+    const unboundCallback = callback;
+    callback = () => ReflectApply(unboundCallback, globalThis, args);
+  }
   timeout = webidl.converters.long(timeout);
   return core.queueUserTimer(
     core.getTimerDepth() + 1,
@@ -88,17 +75,10 @@ function setInterval(callback, timeout = 0, ...args) {
     const unboundCallback = webidl.converters.DOMString(callback);
     callback = () => indirectEval(unboundCallback);
   }
-  const unboundCallback = callback;
-  const asyncContext = getAsyncContext();
-  callback = () => {
-    const oldContext = getAsyncContext(asyncContext);
-    try {
-      setAsyncContext(asyncContext);
-      ReflectApply(unboundCallback, globalThis, args);
-    } finally {
-      setAsyncContext(oldContext);
-    }
-  };
+  if (args.length > 0) {
+    const unboundCallback = callback;
+    callback = () => ReflectApply(unboundCallback, globalThis, args);
+  }
   timeout = webidl.converters.long(timeout);
   return core.queueUserTimer(
     core.getTimerDepth() + 1,
diff --git a/runtime/js/01_async_context.js b/runtime/js/01_async_context.js
deleted file mode 100644
index 9c0236fbe..000000000
--- a/runtime/js/01_async_context.js
+++ /dev/null
@@ -1,45 +0,0 @@
-// Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
-
-import { primordials } from "ext:core/mod.js";
-import { op_get_extras_binding_object } from "ext:core/ops";
-
-const {
-  SafeWeakMap,
-} = primordials;
-
-const {
-  getContinuationPreservedEmbedderData,
-  setContinuationPreservedEmbedderData,
-} = op_get_extras_binding_object();
-
-let counter = 0;
-
-export const getAsyncContext = getContinuationPreservedEmbedderData;
-export const setAsyncContext = setContinuationPreservedEmbedderData;
-
-export class AsyncVariable {
-  #id = counter++;
-  #data = new SafeWeakMap();
-
-  enter(value) {
-    const previousContextMapping = getAsyncContext();
-    const entry = { id: this.#id };
-    const asyncContextMapping = {
-      __proto__: null,
-      ...previousContextMapping,
-      [this.#id]: entry,
-    };
-    this.#data.set(entry, value);
-    setAsyncContext(asyncContextMapping);
-    return previousContextMapping;
-  }
-
-  get() {
-    const current = getAsyncContext();
-    const entry = current?.[this.#id];
-    if (entry) {
-      return this.#data.get(entry);
-    }
-    return undefined;
-  }
-}
diff --git a/runtime/ops/runtime.rs b/runtime/ops/runtime.rs
index b52a23f30..306e6ce8f 100644
--- a/runtime/ops/runtime.rs
+++ b/runtime/ops/runtime.rs
@@ -2,14 +2,13 @@
 
 use deno_core::error::AnyError;
 use deno_core::op2;
-use deno_core::v8;
 use deno_core::ModuleSpecifier;
 use deno_core::OpState;
 use deno_permissions::PermissionsContainer;
 
 deno_core::extension!(
   deno_runtime,
-  ops = [op_main_module, op_ppid, op_get_extras_binding_object],
+  ops = [op_main_module, op_ppid],
   options = { main_module: ModuleSpecifier },
   state = |state, options| {
     state.put::<ModuleSpecifier>(options.main_module);
@@ -95,11 +94,3 @@ pub fn op_ppid() -> i64 {
     parent_id().into()
   }
 }
-
-#[op2]
-pub fn op_get_extras_binding_object<'a>(
-  scope: &mut v8::HandleScope<'a>,
-) -> v8::Local<'a, v8::Value> {
-  let context = scope.get_current_context();
-  context.get_extras_binding_object(scope).into()
-}
diff --git a/runtime/shared.rs b/runtime/shared.rs
index 185cbc0a9..1b2136c63 100644
--- a/runtime/shared.rs
+++ b/runtime/shared.rs
@@ -38,7 +38,6 @@ extension!(runtime,
     dir "js",
     "01_errors.js",
     "01_version.ts",
-    "01_async_context.js",
     "06_util.js",
     "10_permissions.js",
     "11_workers.js",
diff --git a/tests/unit_node/async_hooks_test.ts b/tests/unit_node/async_hooks_test.ts
index 91130972c..f153f6753 100644
--- a/tests/unit_node/async_hooks_test.ts
+++ b/tests/unit_node/async_hooks_test.ts
@@ -1,7 +1,5 @@
 // Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
 import { AsyncLocalStorage, AsyncResource } from "node:async_hooks";
-import process from "node:process";
-import { setImmediate } from "node:timers";
 import { assert, assertEquals } from "@std/assert";
 
 Deno.test(async function foo() {
@@ -94,7 +92,7 @@ Deno.test(async function enterWith() {
   });
 
   assertEquals(await deferred.promise, { x: 2 });
-  assertEquals(await deferred1.promise, null);
+  assertEquals(await deferred1.promise, { x: 1 });
 });
 
 Deno.test(async function snapshot() {
@@ -137,26 +135,3 @@ Deno.test(function emitDestroyStub() {
   const resource = new AsyncResource("foo");
   assert(typeof resource.emitDestroy === "function");
 });
-
-Deno.test(async function worksWithAsyncAPIs() {
-  const store = new AsyncLocalStorage();
-  const test = () => assertEquals(store.getStore(), "data");
-  await store.run("data", async () => {
-    test();
-    queueMicrotask(() => test());
-    process.nextTick(() => test());
-    setImmediate(() => test());
-    setTimeout(() => test(), 0);
-    const intervalId = setInterval(() => {
-      test();
-      clearInterval(intervalId);
-    }, 0);
-
-    store.run("data2", () => {
-      assertEquals(store.getStore(), "data2");
-    });
-
-    await new Promise((r) => setTimeout(r, 50));
-    test();
-  });
-});