spark/src/splatWorker.ts

import { getArrayBuffers } from "./utils.js";
import BundledWorker from "./worker?worker&inline";

// SplatWorker is an internal class that manages a WebWorker for executing
// longer running CPU tasks such as Gsplat file decoding and sorting.
// Although a SplatWorker can be created and used directly, the utility
// function withWorker() is recommended to allocate from a managed
// pool of SplatWorkers.

export class SplatWorker {
  worker: Worker;
  messages: Record<
    number,
    { resolve: (value: unknown) => void; reject: (reason?: unknown) => void }
  > = {};
  messageIdNext = 0;

  constructor() {
    // this.worker = new Worker(new URL("./worker", import.meta.url), { type: "module" });
    this.worker = new BundledWorker();
    this.worker.onmessage = (event) => this.onMessage(event);
  }

  makeMessageId(): number {
    return ++this.messageIdNext;
  }

  makeMessagePromiseId(): { id: number; promise: Promise<unknown> } {
    const id = this.makeMessageId();
    const promise = new Promise((resolve, reject) => {
      this.messages[id] = { resolve, reject };
    });
    return { id, promise };
  }

  onMessage(event: MessageEvent) {
    // console.log("SplatWorker.onMessage:", event);
    const { id, result, error } = event.data;
    // console.log(`SplatWorker.onMessage(${id}):`, result, error);
    const handler = this.messages[id];
    if (handler) {
      delete this.messages[id];
      if (error) {
        handler.reject(error);
      } else {
        handler.resolve(result);
      }
    }
  }

  // Invoke an RPC on the worker with the given name and arguments.
  // The normal usage of a worker is to run one activity at a time,
  // but this function allows for concurrent calls, tagging each request
  // with a unique message Id and awaiting a response to that same Id.
  // The method will automatically transfer any ArrayBuffers in the
  // arguments to the worker. If you'd like to transfer a copy of a
  // buffer then you must clone it before passing to this function.
  async call(name: string, args: unknown): Promise<unknown> {
    const { id, promise } = this.makeMessagePromiseId();
    // console.log(`SplatWorker.call(${name}):`, args);
    this.worker.postMessage(
      { name, args, id },
      { transfer: getArrayBuffers(args) },
    );
    return promise;
  }
}

let maxWorkers = 4;

let numWorkers = 0;
const freeWorkers: SplatWorker[] = [];
const workerQueue: ((worker: SplatWorker) => void)[] = [];

// Set the maximum number of workers to allocate for the pool. (default: 4)
export function setWorkerPool(count = 4) {
  maxWorkers = count;
}

// Allocate a worker from the pool. If none are available and we are below the
// maximum, create a new one. Otherwise, add the request to a queue and wait
// for it to be fulfilled.
export async function allocWorker(): Promise<SplatWorker> {
  const worker = freeWorkers.shift();
  if (worker) {
    return worker;
  }

  if (numWorkers < maxWorkers) {
    const worker = new SplatWorker();
    numWorkers += 1;
    return worker;
  }

  return new Promise((resolve) => {
    workerQueue.push(resolve);
  });
}

// Return a worker to the pool. Pass the worker to any pending waiter.
export function freeWorker(worker: SplatWorker) {
  if (numWorkers > maxWorkers) {
    // Worker no longer needed
    numWorkers -= 1;
    return;
  }

  const waiter = workerQueue.shift();
  if (waiter) {
    waiter(worker);
    return;
  }

  freeWorkers.push(worker);
}

// Allocate a worker from the pool and invoke the callback with the worker.
// When the callback completes, the worker will be returned to the pool.
export async function withWorker<T>(
  callback: (worker: SplatWorker) => Promise<T>,
): Promise<T> {
  const worker = await allocWorker();
  try {
    return await callback(worker);
  } finally {
    freeWorker(worker);
  }
}