multithreading/dist/index.mjs

⚡ Multithreading functions in JavaScript to speedup heavy workloads, designed to feel like writing vanilla functions.

var _Promise;
(_Promise = Promise).withResolvers ?? (_Promise.withResolvers = function withResolvers() {
  var a,
    b,
    c = new this(function (resolve, reject) {
      a = resolve;
      b = reject;
    });
  return {
    resolve: a,
    reject: b,
    promise: c
  };
});

var _ref;
(_ref = import.meta).resolve ?? (_ref.resolve = async function (specifier, parentUrl) {
  const {
    Module: {
      createRequire
    }
  } = await import('node:module');
  const require = createRequire(import.meta.url);
  return require.resolve(specifier, {
    ...(parentUrl ? {
      paths: [parentUrl]
    } : {})
  });
});

const EVENTS = Symbol.for("events");
class EventTarget {
  addEventListener(type, callback, options) {
    let events = this[EVENTS].get(type);
    if (!events) this[EVENTS].set(type, events = []);
    events.push(callback);
  }
  dispatchEvent(event) {
    event.target = event.currentTarget = this;
    if (this["on" + event.type]) {
      try {
        this["on" + event.type](event);
      } catch (err) {
        console.error(err);
      }
    }
    const list = this[EVENTS].get(event.type);
    if (list == null) return false;
    list.forEach(handler => {
      try {
        handler.call(this, event);
      } catch (err) {
        console.error(err);
      }
    });
    return false;
  }
  removeEventListener(type, callback, options) {
    let events = this[EVENTS].get(type);
    if (events) {
      const index = events.indexOf(callback);
      if (index !== -1) events.splice(index, 1);
    }
  }
  constructor() {
    Object.defineProperty(this, EVENTS, {
      value: new Map()
    });
  }
}

class Event {
  composedPath() {
    throw new Error("Method not implemented.");
  }
  initEvent(type, bubbles, cancelable) {
    throw new Error("Method not implemented.");
  }
  preventDefault() {
    // throw new Error("Method not implemented.");
  }
  stopImmediatePropagation() {
    // throw new Error("Method not implemented.");
  }
  stopPropagation() {
    // throw new Error("Method not implemented.");
  }
  constructor(type, target = null) {
    this.bubbles = false;
    this.cancelBubble = false;
    this.cancelable = false;
    this.composed = false;
    this.currentTarget = null;
    this.defaultPrevented = false;
    this.eventPhase = 0;
    this.isTrusted = false;
    this.returnValue = false;
    this.srcElement = null;
    this.target = null;
    this.NONE = 0;
    this.CAPTURING_PHASE = 1;
    this.AT_TARGET = 2;
    this.BUBBLING_PHASE = 3;
    this.type = type;
    this.timeStamp = Date.now();
  }
}

var _globalThis;
(_globalThis = globalThis).Worker ?? (_globalThis.Worker = (async () => {
  const {
    default: threads
  } = await import('node:worker_threads');
  const {
    URL,
    pathToFileURL,
    fileURLToPath
  } = await import('node:url');
  const WORKER = Symbol.for("worker");
  return class Worker extends EventTarget {
    postMessage(message, options) {
      this[WORKER].postMessage(message, options);
    }
    terminate() {
      this[WORKER].terminate();
    }
    constructor(url, options = {}) {
      super();
      this.onmessage = null;
      this.onmessageerror = null;
      this.onerror = null;
      const {
        name,
        type
      } = options;
      url += "";
      let mod;
      if (/^data:/.test(url)) {
        mod = url;
      } else {
        const baseUrl = pathToFileURL(process.cwd() + "/");
        mod = fileURLToPath(new URL(url, baseUrl));
      }
      const worker = new threads.Worker(new URL(import.meta.url), {
        workerData: {
          mod,
          name,
          type
        }
      });
      Object.defineProperty(this, WORKER, {
        value: worker
      });
      worker.on("message", data => {
        const event = new Event("message");
        event.data = data;
        this.dispatchEvent(event);
      });
      worker.on("error", error => {
        error.type = "error";
        this.dispatchEvent(error);
      });
      worker.on("exit", () => {
        this.dispatchEvent(new Event("close"));
      });
    }
  };
})());

class ErrorEvent extends Event {
  constructor(init) {
    super("error");
    this.colno = 0;
    this.filename = "";
    this.lineno = 0;
    this.message = "";
    Object.assign(this, init);
  }
}

var WorkerGlobalScopePromise = Worker instanceof Promise ? (async () => {
  const {
    default: threads
  } = await import('node:worker_threads');
  return class WorkerGlobalScope extends EventTarget {
    postMessage(data, transferList) {
      threads.parentPort.postMessage(data, transferList);
    }
    // Emulates https://developer.mozilla.org/en-US/docs/Web/API/DedicatedWorkerGlobalScope/close
    close() {
      process.exit();
    }
  };
})() : undefined;

class PromiseRejectionEvent extends Event {
  constructor(init) {
    super("unhandledrejection");
    Object.assign(this, init);
  }
}

Worker instanceof Promise && (async () => {
  const {
    default: threads
  } = await import('node:worker_threads');
  const {
    default: VM
  } = await import('node:vm');
  const WorkerGlobalScope = await WorkerGlobalScopePromise;
  return threads.isMainThread || workerThread();
  async function workerThread() {
    let {
      mod,
      name,
      type
    } = threads.workerData;
    if (!mod) return await Worker;
    // turn global into a mock WorkerGlobalScope
    const self = global.self = global;
    // enqueue messages to dispatch after modules are loaded
    let queue = [];
    function flushQueue() {
      const buffered = queue;
      queue = null;
      buffered.forEach(event => {
        self.dispatchEvent(event);
      });
    }
    threads.parentPort.on("message", data => {
      const event = new Event("message");
      event.data = data;
      if (queue == null) self.dispatchEvent(event);else queue.push(event);
    });
    threads.parentPort.on("error", error => {
      error.type = "Error";
      self.dispatchEvent(new ErrorEvent({
        error
      }));
    });
    process.on("unhandledRejection", (reason, promise) => {
      self.dispatchEvent(new PromiseRejectionEvent({
        reason,
        promise
      }));
    });
    process.on("uncaughtException", (error, origin) => {
      self.dispatchEvent(new ErrorEvent({
        error
      }));
    });
    let proto = Object.getPrototypeOf(global);
    delete proto.constructor;
    Object.defineProperties(WorkerGlobalScope.prototype, proto);
    proto = Object.setPrototypeOf(global, new WorkerGlobalScope());
    ["postMessage", "addEventListener", "removeEventListener", "dispatchEvent"].forEach(fn => {
      proto[fn] = proto[fn].bind(global);
    });
    global.name = name;
    const isDataUrl = /^data:/.test(mod);
    if (type === "module") {
      import(mod).catch(err => {
        if (isDataUrl && err.message === "Not supported") {
          console.warn("Worker(): Importing data: URLs requires Node 12.10+. Falling back to classic worker.");
          return evaluateDataUrl(mod, name);
        }
        console.error(err);
      }).then(flushQueue);
    } else {
      try {
        if (isDataUrl) {
          evaluateDataUrl(mod, name);
        } else {
          require(mod);
        }
      } catch (err) {
        console.error(err);
      }
      Promise.resolve().then(flushQueue);
    }
  }
  function evaluateDataUrl(url, name) {
    const {
      data
    } = parseDataUrl(url);
    return VM.runInThisContext(data, {
      filename: "worker.<" + (name || "data:") + ">"
    });
  }
  function parseDataUrl(url) {
    let [m, type, encoding, data] = url.match(/^data: *([^;,]*)(?: *; *([^,]*))? *,(.*)$/) || [];
    if (!m) throw Error("Invalid Data URL.");
    if (encoding) switch (encoding.toLowerCase()) {
      case "base64":
        data = Buffer.from(data, "base64").toString();
        break;
      default:
        throw Error('Unknown Data URL encoding "' + encoding + '"');
    }
    return {
      type,
      data
    };
  }
})();

const Function = 2;
const Other = 3;
const Init = 4;
const Claim = 5;
const Unclaim = 6;
const Return = 7;
const ClaimAcceptance = 8;
const ClaimRejection = 9;
const Invocation = 12;
const Synchronization = 13;
const Error$1 = 14;
const Variables = "a";
const Args = "b";
const EventType = "c";
const EventValue = "d";
const WasType = "e";
const Name = "f";
const YieldList = "g";
const InvocationId = "h";
const Value = "i";
const ProcessId = "j";
const DebugEnabled = "k";
const Type = "l";
const AbsolutePath = "m";
const Code = "n";
const Pid = "o";
const UserFunction = "p";
const Internal = "q";
const ShareableNameMap = "r";
const ValueClaimMap = "s";
const ValueInUseCount = "t";

const serialize = variables => {
  const serializedVariables = {};
  for (const [key, value] of Object.entries(variables)) {
    if (typeof value === "function") {
      serializedVariables[key] = {
        [WasType]: Function,
        value: value.toString()
      };
    } else {
      serializedVariables[key] = value;
    }
  }
  return serializedVariables;
};
const deserialize = variables => {
  const deserializedVariables = {};
  for (const [key, value] of Object.entries(variables)) {
    if (typeof value === "object" && WasType in value) {
      switch (value[WasType]) {
        default:
          deserializedVariables[key] = value;
          break;
      }
    } else {
      deserializedVariables[key] = value;
    }
  }
  return deserializedVariables;
};

function replaceContents(originalObject, newValue) {
  if (Array.isArray(originalObject)) {
    // Clear the array and push new values
    originalObject.length = 0;
    newValue.forEach(item => originalObject.push(item));
  } else if (originalObject instanceof Map) {
    // Clear the map and set new key-value pairs
    originalObject.clear();
    newValue.forEach(([key, value]) => originalObject.set(key, value));
  } else if (originalObject instanceof Set) {
    // Clear the set and add new values
    originalObject.clear();
    newValue.forEach(item => originalObject.add(item));
  } else if (typeof originalObject === "object" && originalObject !== null) {
    // Clear the object and assign new properties
    for (const key in originalObject) {
      delete originalObject[key];
    }
    Object.assign(originalObject, newValue);
  } else {
    throw new Error("Unsupported object type");
  }
}

const colorGray = "\x1b[90m";
const colorRed = "\x1b[31m";
const colorCyan = "\x1b[36m";
const colorReset = "\x1b[39m";
function getErrorPreview(error, code, pid) {
  const [message, ...serializedStackFrames] = error.stack.split("\n");
  // Check if error originates from inside the user function
  const stackFrame = serializedStackFrames.find(frame => frame.includes("data:application/javascript;base64"));
  if (!stackFrame) {
    return error.stack;
  }
  // Split at the comma of data:application/javascript;base64,
  const [functionPart, tracePart] = stackFrame.split(",");
  const [encodedBodyPart, lineNumberStr, columnNumberStr] = tracePart.split(":");
  const lineNumber = parseInt(lineNumberStr);
  const columnNumber = parseInt(columnNumberStr);
  const codeLines = code.split(/\r?\n/);
  const amountOfPreviousLines = Math.min(3, lineNumber - 1);
  const amountOfNextLines = 2;
  const previewLines = codeLines.slice(lineNumber - (amountOfPreviousLines + 1), lineNumber + amountOfNextLines);
  const previousLineLength = codeLines[lineNumber - 1].trimEnd().length - columnNumber;
  previewLines.splice(amountOfPreviousLines + 1, 0, colorRed + " ".repeat(columnNumber - 1) + "^".repeat(previousLineLength) + " " +
  // "Error" +
  message + colorGray);
  const index = serializedStackFrames.indexOf(stackFrame);
  serializedStackFrames[index] = `    at ${colorCyan}<Thread_${pid}>${colorReset}\n` + colorGray + "    " + previewLines.join("\n    ") + colorReset;
  // return message + "\n" + serializedStackFrames.slice(0, index + 1).join("\n");
  return message.split(":").slice(1).join(":").trim() + "\n" + serializedStackFrames.slice(0, index + 1).join("\n");
}

function announceOwnership(queue, valueName, value) {
  // Get first worker in queue
  const worker = queue[0];
  worker.postMessage({
    [EventType]: ClaimAcceptance,
    [EventValue]: {
      [Name]: valueName,
      [Value]: value
    }
  });
}
function setupWorkerListeners(worker, context, valueOwnershipQueue, invocationQueue, workerPool, workerCodeString, pid) {
  worker.onmessage = e => {
    switch (e.data[EventType]) {
      case Return:
        const invocationId = e.data[EventValue][InvocationId];
        const value = e.data[EventValue][Value];
        const {
          resolve
        } = invocationQueue.get(invocationId);
        resolve(value);
        invocationQueue.delete(invocationId);
        break;
      case Claim:
        {
          const valueName = e.data[EventValue];
          const value = context[valueName];
          const queue = valueOwnershipQueue.get(value);
          queue.push(worker);
          if (queue.length === 1) {
            announceOwnership(queue, valueName, value);
          }
          break;
        }
      case Unclaim:
        {
          const data = e.data[EventValue];
          const valueName = data[Name];
          const value = context[valueName];
          const ownershipQueue = valueOwnershipQueue.get(value);
          // Check if worker is first in queue
          if (ownershipQueue[0] !== worker) break;
          const newValue = data[Value];
          // Update local value with new value
          replaceContents(value, newValue);
          // Synchronize all other workers with new value
          for (const otherWorker of workerPool) {
            if (otherWorker === worker) continue;
            worker.postMessage({
              [EventType]: Synchronization,
              [EventValue]: {
                [Name]: valueName,
                [Value]: newValue
              }
            });
          }
          ownershipQueue.shift();
          if (ownershipQueue.length > 0) {
            announceOwnership(ownershipQueue, valueName, value);
          }
          break;
        }
      case Error$1:
        {
          const error = e.data[EventValue];
          error.message = getErrorPreview(error, workerCodeString, pid);
          error.stack = ""; // Deno doesn't like custom stack traces, use message instead
          invocationQueue.forEach(({
            reject
          }) => reject(error));
        }
    }
  };
}

async function parseImport(name) {
  const resolved = await import.meta.resolve(name);
  if (resolved.startsWith("http://") || resolved.startsWith("https://") || resolved.startsWith("npm:") || resolved.startsWith("node:")) return resolved;
  // Check if running in browser
  const isBrowser = typeof window !== "undefined";
  if (isBrowser) {
    // If running in browser, return the resolved URL
    return resolved;
  }
  const {
    pathToFileURL
  } = await import('node:url');
  return pathToFileURL(resolved).toString();
}
async function parseTopLevelYieldStatements(fnStr) {
  const bodyStart = fnStr.indexOf("{") + 1;
  const code = fnStr.slice(bodyStart, -1).trim();
  const lines = code.split(/(?:\s*[;\r\n]+\s*)+/);
  const yieldList = [];
  // let insideCommentBlock = false;
  for (const line of lines) {
    // Skip comments
    // if (line.startsWith("/*")) insideCommentBlock = true;
    // if (line.endsWith("*/") || line.startsWith("*/")) {
    //   insideCommentBlock = false;
    //   continue;
    // }
    // if (insideCommentBlock || line.startsWith("//")) continue;
    // If line is not a yield statement, stop parsing
    if (!line.includes("yield ")) continue;
    const yielded = line.split("yield ")[1];
    if (/^["'`]/.test(yielded)) {
      const name = yielded.slice(1, -1);
      yieldList.push({
        [Type]: "import",
        [Name]: name,
        [AbsolutePath]: await parseImport(name)
      });
    } else {
      yieldList.push({
        [Type]: "variable",
        [Name]: yielded
      });
    }
  }
  return yieldList;
}

const inlineWorker = `Promise.withResolvers ??= function withResolvers() {
  var a,
    b,
    c = new this(function (resolve, reject) {
      a = resolve;
      b = reject;
    });
  return {
    resolve: a,
    reject: b,
    promise: c
  };
};

const Init = 4;
const Claim = 5;
const Unclaim = 6;
const Return = 7;
const ClaimAcceptance = 8;
const Invocation = 12;
const Synchronization = 13;
const Error\$1 = 14;
const Variables = "a";
const Args = "b";
const EventType = "c";
const EventValue = "d";
const WasType = "e";
const Name = "f";
const YieldList = "g";
const InvocationId = "h";
const Value = "i";
const ProcessId = "j";
const Type = "l";
const AbsolutePath = "m";
const Code = "n";
const UserFunction = "p";

const deserialize = variables => {
  const deserializedVariables = {};
  for (const [key, value] of Object.entries(variables)) {
    if (typeof value === "object" && WasType in value) {
      switch (value[WasType]) {
        default:
          deserializedVariables[key] = value;
          break;
      }
    } else {
      deserializedVariables[key] = value;
    }
  }
  return deserializedVariables;
};

function replaceContents(originalObject, newValue) {
  if (Array.isArray(originalObject)) {
    // Clear the array and push new values
    originalObject.length = 0;
    newValue.forEach(item => originalObject.push(item));
  } else if (originalObject instanceof Map) {
    // Clear the map and set new key-value pairs
    originalObject.clear();
    newValue.forEach(([key, value]) => originalObject.set(key, value));
  } else if (originalObject instanceof Set) {
    // Clear the set and add new values
    originalObject.clear();
    newValue.forEach(item => originalObject.add(item));
  } else if (typeof originalObject === "object" && originalObject !== null) {
    // Clear the object and assign new properties
    for (const key in originalObject) {
      delete originalObject[key];
    }
    Object.assign(originalObject, newValue);
  } else {
    throw new Error("Unsupported object type");
  }
}

// Wrap in self-invoking function to avoid polluting the global namespace
// and avoid name collisions with the user defined function
globalThis.__internal = function () {
  const state = {
    [UserFunction]: function* () {},
    [Code]: ""
  };

  // const originalLog = console.log;
  // console.log = (...args) => {
  //   originalLog(\`\${cyan}[Thread_\${pid}]\${reset}\`, ...args);
  // };
  // const originalError = console.error;
  // console.error = (...args) => {
  //   originalError(\`\${red}[Thread_\${pid}]\${reset}\`, ...args);
  // };

  globalThis.\$claim = async function \$claim(value) {
    const valueName = shareableNameMap.get(value);
    valueInUseCount[valueName]++;

    // First check if the variable is already (being) claimed
    if (valueClaimMap.has(valueName)) {
      return valueClaimMap.get(valueName).promise;
    }
    valueClaimMap.set(valueName, Promise.withResolvers());
    postMessage({
      [EventType]: Claim,
      [EventValue]: valueName
    });
    return valueClaimMap.get(valueName).promise;
  };
  globalThis.\$unclaim = function \$unclaim(value) {
    const valueName = shareableNameMap.get(value);
    if (--valueInUseCount[valueName] > 0) return;
    valueClaimMap.delete(valueName);
    postMessage({
      [EventType]: Unclaim,
      [EventValue]: {
        [Name]: valueName,
        [Value]: value
      }
    });
  };
  let yieldList = [];
  const shareableNameMap = new WeakMap();

  // ShareableValues that are currently (being) claimed
  const valueClaimMap = new Map();
  // ShareableValues that are currently in use by
  // one of the invokations of the user defined function
  const valueInUseCount = {};
  function handleClaimAcceptance(data) {
    const valueName = data[Name];
    replaceContents(globalThis[valueName], data[Value]);
    valueClaimMap.get(valueName).resolve();
  }
  async function handleInit(data) {
    data[ProcessId];
    yieldList = data[YieldList];
    state[Code] = data[Code];
    const variables = deserialize(data[Variables]);
    for (const key in variables) {
      const value = variables[key];
      if (value instanceof Object) {
        shareableNameMap.set(value, key);
        valueInUseCount[key] = 0;
      }
    }
    Object.assign(globalThis, variables);
  }
  async function handleInvocation(data) {
    const gen = state[UserFunction](...data[Args]);
    let isDone = false;
    let returnValue = undefined;
    let isFirstImport = true;
    for (const yieldItem of yieldList) {
      if (yieldItem[Type] === "import") {
        const resolved = await import(yieldItem[AbsolutePath]);
        if (isFirstImport) {
          await gen.next();
          isFirstImport = false;
        }
        const result = await gen.next(resolved);
        if (result.done) {
          isDone = true;
          returnValue = result.value;
          break;
        }
      } else {
        const result = await gen.next();
        if (result.done) {
          isDone = true;
          returnValue = result.value;
          break;
        }
      }
    }
    if (!isDone) {
      const result = await gen.next();
      returnValue = result.value;
    }
    postMessage({
      [EventType]: Return,
      [EventValue]: {
        [InvocationId]: data[InvocationId],
        [Value]: returnValue
      }
    });
  }
  async function handleSynchronization(data) {
    const valueName = data[Name];
    replaceContents(globalThis[valueName], data[Value]);
  }

  // On unhandled promise rejection
  self.addEventListener("unhandledrejection", event => {
    event.preventDefault();
    postMessage({
      [EventType]: Error\$1,
      [EventValue]: event.reason
    });
    close();
  });

  // On uncaught exception
  self.addEventListener("error", event => {
    event.preventDefault();
    postMessage({
      [EventType]: Error\$1,
      [EventValue]: event.error
    });
    close();
  });
  globalThis.onmessage = async e => {
    switch (e.data[EventType]) {
      case Init:
        handleInit(e.data[EventValue]);
        break;
      case Invocation:
        handleInvocation(e.data[EventValue]);
        break;
      case ClaimAcceptance:
        handleClaimAcceptance(e.data[EventValue]);
        break;
      case Synchronization:
        handleSynchronization(e.data[EventValue]);
        break;
    }
  };
  return state;
}();
`;
async function $claim(value) {}
function $unclaim(value) {}
const workerPools = new WeakMap();
const valueOwnershipQueue = new WeakMap();
function threaded(configOrFn, maybeFn) {
  const config = {
    debug: false,
    maxThreads: typeof navigator !== "undefined" ? navigator.hardwareConcurrency : 4
  };
  let fn;
  if (typeof configOrFn === "function") {
    fn = configOrFn;
  } else {
    Object.assign(config, configOrFn);
    fn = maybeFn;
  }
  let context = {};
  const workerPool = [];
  const invocationQueue = new Map();
  workerPools.set(fn, workerPool);
  let invocationCount = 0;
  const init = (async () => {
    const fnStr = fn.toString();
    const yieldList = await parseTopLevelYieldStatements(fnStr);
    // @ts-ignore - Call function without arguments
    const gen = fn();
    for (const yieldItem of yieldList) {
      // @ts-ignore - Pass empty object to prevent TypeError when user has destructured import
      const result = await gen.next({});
      if (yieldItem[Type] !== "variable") continue;
      context[yieldItem[Name]] = result.value;
    }
    for (const key in context) {
      // Initialize the ownership queue
      valueOwnershipQueue.set(context[key], []);
    }
    const workerCode = [inlineWorker, `__internal.${UserFunction} = ${fnStr};`];
    const serializedVariables = serialize(context);
    for (const [key, value] of Object.entries(serializedVariables)) {
      if (value[WasType] !== Function) continue;
      // globalthis. is necessary to prevent duplicate variable names when the function is named
      workerCode.unshift(`globalThis.${key} = ${value.value};`);
      delete serializedVariables[key];
    }
    const workerCodeString = workerCode.join("\r\n");
    for (let i = 0; i < config.maxThreads; i++) {
      const worker = new (await Worker)(encodeURI("data:application/javascript;base64," + btoa(workerCodeString)), {
        type: "module"
      });
      setupWorkerListeners(worker, context, valueOwnershipQueue, invocationQueue, workerPool, workerCodeString, i);
      workerPool.push(worker);
      worker.postMessage({
        [EventType]: Init,
        [EventValue]: {
          [ProcessId]: i,
          [YieldList]: yieldList,
          [Variables]: serializedVariables,
          [Code]: workerCodeString,
          [DebugEnabled]: config.debug
        }
      });
    }
  })();
  const wrapper = async (...args) => {
    await init;
    const worker = workerPool[invocationCount % config.maxThreads];
    const pwr = Promise.withResolvers();
    invocationQueue.set(invocationCount, pwr);
    worker.postMessage({
      [EventType]: Invocation,
      [EventValue]: {
        [InvocationId]: invocationCount++,
        [Args]: args
      }
    });
    return pwr.promise;
  };
  wrapper.dispose = () => {
    for (const worker of workerPool) {
      worker.terminate();
    }
    workerPools.delete(fn);
    invocationQueue.forEach(pwr => pwr.reject("Disposed"));
    invocationQueue.clear();
  };
  return wrapper;
}

export { $claim, $unclaim, threaded };