AsyncContext polyfill

async-context-polyfill.js

// Copyright (C) 2024 Igalia, S.L.
// This code is licensed under the terms of the MIT license
// (https://spdx.org/licenses/MIT.html).

// This is a polyfill for AsyncContext, including the web integration, without
// requiring user code transpilation. However, it only works in Chrome, and only
// enabling the `UnexposedTaskIds` feature. This can be enabled through a
// command line flag: `google-chrome --enable-blink-features=UnexposedTaskIds`.
// 
// This polyfill has certain features that are known to be missing/incorrect:
//   - Propagation across generators
//   - Unhandled rejections
//   - Evaluating modules in an empty context
//   - Anything cross-realm
//
// Because of this, this polyfill should only be used for testing, not in
// production.

(() => {
  if ("AsyncContext" in globalThis) {
    return;
  }

  if (typeof scheduler.taskId !== "number") {
    throw new Error("Invalid environment");
  }

  let nextId = 1000000;

  const idMap /*: Map<TaskId, Map<AsyncContext.Variable, any>> */ = new Map();
  idMap.set(scheduler.taskId, new Map());

  globalThis.AsyncContext = {
    Variable: class Variable {
      #name;
      #defaultValue;

      constructor(options) {
        if (typeof options === "object" && options !== null) {
          this.#name = ("name" in options) ? String(options.name) : "";
          this.#defaultValue = options.defaultValue;
        } else {
          this.#name = "";
          this.#defaultValue = undefined;
        }
      }

      run(value, cb, ...args) {
        checkTaskId();

        const newId = nextId;
        nextId++;

        const newSnapshot = new Map(idMap.get(scheduler.taskId));
        newSnapshot.set(this, value);
        idMap.set(newId, newSnapshot);

        return runTaskId(newId, cb, args);
      }

      get name() {
        return this.#name;
      }

      get() {
        checkTaskId();
        const currentSnapshot = idMap.get(scheduler.taskId);
        if (currentSnapshot.has(this)) {
          return currentSnapshot.get(this);
        }
        return this.#defaultValue;
      }
    },

    Snapshot: class Snapshot {
      #id;

      constructor() {
        checkTaskId();
        this.#id = scheduler.taskId;
      }

      run(cb, ...args) {
        checkTaskId();
        return runTaskId(this.#id, cb, args);
      }

      static wrap(cb) {
        checkTaskId();
        const id = scheduler.taskId;
        return function (...args) {
          return runTaskId(id, cb.bind(this), args);
        }
      }
    }
  };

  function checkTaskId() {
    if (!idMap.has(scheduler.taskId)) {
      console.warn("Found an unknown taskId at: ", new Error().stack);
      idMap.set(scheduler.taskId, new Map());
    }
  }

  function runTaskId(newTaskId, cb, args) {
    const previousTaskId = scheduler.taskId;
    scheduler.taskId = newTaskId;
    try {
      return cb(...args);
    } finally {
      scheduler.taskId = previousTaskId;
    }
  }

  // -------------------------------------------------------------------------

  const fallbackTaskIdVar /*: AsyncContext.Variable<TaskId> */ = new AsyncContext.Variable();

  function wrapCbWithFallbackTask(cb) {
    // If the callback is called with an unknown taskId, we treat that as having
    // no dispatch context.
    //
    // We must wrap all event listener, even if there is no fallback context,
    // because if we have some listeners with a fallback context and some
    // without for the same event, we can lose track of whether the event has no
    // dispatch context. For that reason, we always do `runTaskId` if the
    // taskId is unknown.

    const fallbackTaskId = fallbackTaskIdVar.get();
    return (...args) => {
      if (!idMap.has(scheduler.taskId)) {
        let taskIdToRun = fallbackTaskId;
        if (!taskIdToRun) {
          taskIdToRun = nextId;
          nextId++;
          idMap.set(taskIdToRun, new Map());
        }
        return runTaskId(fallbackTaskId, cb, args);
      }
      return cb(...args);
    };
  }

  const originalAddEventListener = EventTarget.prototype.addEventListener;

  EventTarget.prototype.addEventListener = function (type, cb, options) {
    return originalAddEventListener.call(this, type, wrapCbWithFallbackTask(cb), options);
  };

  EventTarget.captureFallbackContext = (cb, ...args) => {
    checkTaskId();
    return fallbackTaskIdVar.run(scheduler.taskId, cb, ...args);
  };

  function patchEventHandlers(proto) {
    for (propertyName of Object.getOwnPropertyNames(proto)) {
      if (propertyName.startsWith("on")) {
        const descriptor = Object.getOwnPropertyDescriptor(proto, propertyName);
        if (descriptor?.get && descriptor?.set) {
          Object.defineProperty(proto, propertyName, {
            ...descriptor,
            set(cb) {
              descriptor.set.call(this, wrapCbWithFallbackTask(cb));
            }
          });
        }
      }
    }
  }

  const handledProtos = new Set();

  for (const name of Object.getOwnPropertyNames(globalThis)) {
    if (typeof globalThis[name] === "function" && globalThis[name].prototype) {
      const proto = globalThis[name].prototype;
      if (!handledProtos.has(proto)) {
        patchEventHandlers(proto);
        handledProtos.add(proto);
      }
    }
  }
  patchEventHandlers(globalThis);

})();