nstdlib-nightly/lib/internal/worker/messaging.mjs

Node.js standard library converted to runtime-agnostic ES modules.

// Source: https://github.com/nodejs/node/blob/65eff1eb/lib/internal/worker/messaging.js

import {
  isMainThread,
  threadId as currentThreadId,
} from "nstdlib/stub/binding/worker";
import { codes as __codes__ } from "nstdlib/lib/internal/errors";
import { MessageChannel } from "nstdlib/lib/internal/worker/io";
import { validateNumber } from "nstdlib/lib/internal/validators";

const { SharedArrayBuffer } = globalThis;

const {
  ERR_WORKER_MESSAGING_ERRORED,
  ERR_WORKER_MESSAGING_FAILED,
  ERR_WORKER_MESSAGING_SAME_THREAD,
  ERR_WORKER_MESSAGING_TIMEOUT,
} = __codes__;

const messageTypes = {
  REGISTER_MAIN_THREAD_PORT: "registerMainThreadPort",
  UNREGISTER_MAIN_THREAD_PORT: "unregisterMainThreadPort",
  SEND_MESSAGE_TO_WORKER: "sendMessageToWorker",
  RECEIVE_MESSAGE_FROM_WORKER: "receiveMessageFromWorker",
};

// This is only populated by main thread and always empty in other threads
const threadsPorts = new Map();

// This is only populated in child threads and always undefined in main thread
let mainThreadPort;

// SharedArrayBuffer must always be Int32, so it's * 4.
// We need one for the operation status (performing / performed) and one for the result (success / failure).
const WORKER_MESSAGING_SHARED_DATA = 2 * 4;
const WORKER_MESSAGING_STATUS_INDEX = 0;
const WORKER_MESSAGING_RESULT_INDEX = 1;

// Response codes
const WORKER_MESSAGING_RESULT_DELIVERED = 0;
const WORKER_MESSAGING_RESULT_NO_LISTENERS = 1;
const WORKER_MESSAGING_RESULT_LISTENER_ERROR = 2;

// This event handler is always executed on the main thread only
function handleMessageFromThread(message) {
  switch (message.type) {
    case messageTypes.REGISTER_MAIN_THREAD_PORT:
      {
        const { threadId, port } = message;

        // Register the port
        threadsPorts.set(threadId, port);

        // Handle messages on this port
        // When a new thread wants to register a children
        // this take care of doing that.
        // This way any thread can be linked to the main one.
        port.on("message", handleMessageFromThread);

        // Never block the thread on this port
        port.unref();
      }

      break;
    case messageTypes.UNREGISTER_MAIN_THREAD_PORT:
      threadsPorts.get(message.threadId).close();
      threadsPorts.delete(message.threadId);
      break;
    case messageTypes.SEND_MESSAGE_TO_WORKER:
      {
        // Send the message to the target thread
        const { source, destination, value, transferList, memory } = message;
        sendMessageToWorker(source, destination, value, transferList, memory);
      }
      break;
  }
}

function handleMessageFromMainThread(message) {
  switch (message.type) {
    case messageTypes.RECEIVE_MESSAGE_FROM_WORKER:
      receiveMessageFromWorker(message.source, message.value, message.memory);
      break;
  }
}

function sendMessageToWorker(source, destination, value, transferList, memory) {
  // We are on the main thread, we can directly process the message
  if (destination === 0) {
    receiveMessageFromWorker(source, value, memory);
    return;
  }

  // Search the port to the target thread
  const port = threadsPorts.get(destination);

  if (!port) {
    const status = new Int32Array(memory);
    AtomicsStore(
      status,
      WORKER_MESSAGING_RESULT_INDEX,
      WORKER_MESSAGING_RESULT_NO_LISTENERS,
    );
    AtomicsStore(status, WORKER_MESSAGING_STATUS_INDEX, 1);
    AtomicsNotify(status, WORKER_MESSAGING_STATUS_INDEX, 1);
    return;
  }

  port.postMessage(
    {
      type: messageTypes.RECEIVE_MESSAGE_FROM_WORKER,
      source,
      destination,
      value,
      memory,
    },
    transferList,
  );
}

function receiveMessageFromWorker(source, value, memory) {
  let response = WORKER_MESSAGING_RESULT_NO_LISTENERS;

  try {
    if (process.emit("workerMessage", value, source)) {
      response = WORKER_MESSAGING_RESULT_DELIVERED;
    }
  } catch {
    response = WORKER_MESSAGING_RESULT_LISTENER_ERROR;
  }

  // Populate the result
  const status = new Int32Array(memory);
  AtomicsStore(status, WORKER_MESSAGING_RESULT_INDEX, response);
  AtomicsStore(status, WORKER_MESSAGING_STATUS_INDEX, 1);
  AtomicsNotify(status, WORKER_MESSAGING_STATUS_INDEX, 1);
}

function createMainThreadPort(threadId) {
  // Create a channel that links the new thread to the main thread
  const { port1: mainThreadPortToMain, port2: mainThreadPortToThread } =
    new MessageChannel();

  const registrationMessage = {
    type: messageTypes.REGISTER_MAIN_THREAD_PORT,
    threadId,
    port: mainThreadPortToMain,
  };

  if (isMainThread) {
    handleMessageFromThread(registrationMessage);
  } else {
    mainThreadPort.postMessage(registrationMessage, [mainThreadPortToMain]);
  }

  return mainThreadPortToThread;
}

function destroyMainThreadPort(threadId) {
  const unregistrationMessage = {
    type: messageTypes.UNREGISTER_MAIN_THREAD_PORT,
    threadId,
  };

  if (isMainThread) {
    handleMessageFromThread(unregistrationMessage);
  } else {
    mainThreadPort.postMessage(unregistrationMessage);
  }
}

function setupMainThreadPort(port) {
  mainThreadPort = port;
  mainThreadPort.on("message", handleMessageFromMainThread);

  // Never block the process on this port
  mainThreadPort.unref();
}

async function postMessageToThread(threadId, value, transferList, timeout) {
  if (typeof transferList === "number" && typeof timeout === "undefined") {
    timeout = transferList;
    transferList = [];
  }

  if (typeof timeout !== "undefined") {
    validateNumber(timeout, "timeout", 0);
  }

  if (threadId === currentThreadId) {
    throw new ERR_WORKER_MESSAGING_SAME_THREAD();
  }

  const memory = new SharedArrayBuffer(WORKER_MESSAGING_SHARED_DATA);
  const status = new Int32Array(memory);
  const promise = AtomicsWaitAsync(
    status,
    WORKER_MESSAGING_STATUS_INDEX,
    0,
    timeout,
  ).value;

  const message = {
    type: messageTypes.SEND_MESSAGE_TO_WORKER,
    source: currentThreadId,
    destination: threadId,
    value,
    memory,
    transferList,
  };

  if (isMainThread) {
    handleMessageFromThread(message);
  } else {
    mainThreadPort.postMessage(message, transferList);
  }

  // Wait for the response
  const response = await promise;

  if (response === "timed-out") {
    throw new ERR_WORKER_MESSAGING_TIMEOUT();
  } else if (
    status[WORKER_MESSAGING_RESULT_INDEX] ===
    WORKER_MESSAGING_RESULT_NO_LISTENERS
  ) {
    throw new ERR_WORKER_MESSAGING_FAILED();
  } else if (
    status[WORKER_MESSAGING_RESULT_INDEX] ===
    WORKER_MESSAGING_RESULT_LISTENER_ERROR
  ) {
    throw new ERR_WORKER_MESSAGING_ERRORED();
  }
}

export { createMainThreadPort };
export { destroyMainThreadPort };
export { setupMainThreadPort };
export { postMessageToThread };