123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332
// Copyright 2018-2023 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
import { validateFunction } from "./internal/validators.mjs";import { core } from "./_core.ts";
function assert(cond: boolean) {  if (!cond) throw new Error("Assertion failed");}const asyncContextStack: AsyncContextFrame[] = [];
function pushAsyncFrame(frame: AsyncContextFrame) {  asyncContextStack.push(frame);}
function popAsyncFrame() {  assert(asyncContextStack.length > 0);  asyncContextStack.pop();}
let rootAsyncFrame: AsyncContextFrame | undefined = undefined;let promiseHooksSet = false;
const asyncContext = Symbol("asyncContext");function isRejected(promise: Promise<unknown>) {  const [state] = core.getPromiseDetails(promise);  return state == 2;}
function setPromiseHooks() {  if (promiseHooksSet) {    return;  }  promiseHooksSet = true;
  const init = (promise: Promise<unknown>) => {    const currentFrame = AsyncContextFrame.current();    if (!currentFrame.isRoot()) {      assert(AsyncContextFrame.tryGetContext(promise) == null);      AsyncContextFrame.attachContext(promise);    }  };  const before = (promise: Promise<unknown>) => {    const maybeFrame = AsyncContextFrame.tryGetContext(promise);    if (maybeFrame) {      pushAsyncFrame(maybeFrame);    } else {      pushAsyncFrame(AsyncContextFrame.getRootAsyncContext());    }  };  const after = (promise: Promise<unknown>) => {    popAsyncFrame();    if (!isRejected(promise)) {      // @ts-ignore promise async context      delete promise[asyncContext];    }  };  const resolve = (promise: Promise<unknown>) => {    const currentFrame = AsyncContextFrame.current();    if (      !currentFrame.isRoot() && isRejected(promise) &&      AsyncContextFrame.tryGetContext(promise) == null    ) {      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));
      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>) {    assert(!(asyncContext in promise));    // @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;  constructor(type: string) {    this.type = type;    this.frame = AsyncContextFrame.current();  }
  runInAsyncScope(    fn: (...args: unknown[]) => unknown,    thisArg: unknown,    ...args: unknown[]  ) {    Scope.enter(this.frame);
    try {      return fn.apply(thisArg, args);    } finally {      Scope.exit();    }  }
  bind(fn: (...args: unknown[]) => unknown, thisArg = this) {    validateFunction(fn, "fn");    const frame = AsyncContextFrame.current();    const bound = AsyncContextFrame.wrap(fn, frame, thisArg);
    Object.defineProperties(bound, {      "length": {        configurable: true,        enumerable: false,        value: fn.length,        writable: false,      },      "asyncResource": {        configurable: true,        enumerable: true,        value: this,        writable: true,      },    });    return bound;  }
  static bind(    fn: (...args: unknown[]) => unknown,    type?: string,    thisArg?: AsyncResource,  ) {    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;  }}
class StorageKey {  #dead = false;
  reset() {    this.#dead = true;  }
  isDead() {    return this.#dead;  }}
const fnReg = new FinalizationRegistry((key: StorageKey) => {  key.reset();});
export class AsyncLocalStorage {  #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 {    const frame = AsyncContextFrame.create(      null,      new StorageEntry(this.#key, store),    );    Scope.enter(frame);    let res;    try {      res = callback(...args);    } finally {      Scope.exit();    }    return res;  }
  // deno-lint-ignore no-explicit-any  exit(callback: (...args: unknown[]) => any, ...args: any[]): any {    return this.run(undefined, callback, args);  }
  // deno-lint-ignore no-explicit-any  getStore(): any {    const currentFrame = AsyncContextFrame.current();    return currentFrame.get(this.#key);  }}
export function executionAsyncId() {  return 1;}
class AsyncHook {  enable() {  }
  disable() {  }}
export function createHook() {  return new AsyncHook();}
// Placing all exports down here because the exported classes won't export// otherwise.export default {  // Embedder API  AsyncResource,  executionAsyncId,  createHook,  AsyncLocalStorage,};