Shared worker shim + promised based insanity

README.md

Promised based typesafe shared worker messaging

I built this thing, it's kind of insane. In the worker, you can call it almost like any other async function... including type safety.

I don't really have time to explain more, but maybe i will later.

usage.ts

import BfSharedWorker from "./main.ts"

const worker = new BfSharedWorker();
const result = await worker.adder(1,2);
const typeErrorLol = await worker.adder("lol", 2);

worker.ts

class SharedWorkerShim extends EventTarget implements SharedWorker {
  port: MessagePort;
  private backingWorker: Worker;

  constructor(url: string | URL, options: string | WorkerOptions = {}) {
    super();
    const { port1, port2 } = new MessageChannel();
    this.port = new Proxy(port1, {
      get: (target, property) => {
        if (property === "start") {
          return () => {
            target["start"](); // Call the original start method
            // Dispatch a fake 'onconnect' event when start is called
            setTimeout(() => {
              this.backingWorker.postMessage({
                type: "onconnect",
                ports: [port2],
              }, [port2]);
            }, 0);
          };
        }
        // @ts-expect-error we're doing pretty crazy stuff.
        return target[property];
      },
    });

    if (typeof options === "string") {
      options = { name: options };
    }
    options.type = "module";
    this.backingWorker = new Worker(url, options);
  }

  onerror(ev: ErrorEvent) {
    console.error(ev);
  }

  close() {
    this.port.close();
    this.backingWorker.terminate();
    this.backingWorker.onmessage = null;
    this.backingWorker.onerror = null;
  }
}

if (typeof SharedWorker === "undefined") {
  globalThis.SharedWorker = SharedWorkerShim;
}

type PublicMethodKeys<T> = {
  [K in keyof T]: T[K] extends (...args: infer P) => infer R
    ? (P extends unknown[] ? (R extends unknown ? K : never) : never)
    : never;
}[keyof T];

// Maps each method of T to its argument types as a tuple.
type MethodArguments<T> = {
  [K in keyof T]: T[K] extends (...args: infer P) => unknown ? P : never;
};

// Maps each method of T to its return type.
type MethodReturnTypes<T> = {
  [K in keyof T]: T[K] extends (...args: unknown[]) => infer R ? R : never;
};

type WorkerMessage<T extends BfSharedWorker<T>> = {
  type: PublicMethodKeys<T>;
  args: MethodArguments<T>[PublicMethodKeys<T>];
};

export enum BfSharedWorkerTypes {
  "MAIN_WORKER" = "MAIN_WORKER",
  "MAIN_THREAD" = "MAIN_THREAD",
  "IMPORTED_WORKER" = "IMPORTED_WORKER",
  "IMPORTED_MAIN_THREAD" = "IMPORTED_MAIN_THREAD",
}
export default class BfSharedWorker<T extends BfSharedWorker<T>> {
  [key: string | symbol]: unknown;
  protected mainUrl = "/resources/workers/bf-shared-worker.ts"; // my server and my import map both resolve this correctly.
  protected worker?: SharedWorker;
  protected supportsRunningAsMainThread = false;
  protected supportsRunningAsWorker = true;
  protected workerName = this.constructor.name;

  private get workerType() {
    const isWorker = typeof globalThis.window === "undefined";
    const isMainRunningInMainThread = typeof globalThis.window !== "undefined";
    const isMain = import.meta.main;
    const isImported = !import.meta.main;

    if (isMain && isMainRunningInMainThread) {
      return BfSharedWorkerTypes.MAIN_THREAD;
    }
    if (isMain && !isMainRunningInMainThread) {
      return BfSharedWorkerTypes.MAIN_WORKER;
    }
    if (isImported && isWorker) {
      return BfSharedWorkerTypes.IMPORTED_WORKER;
    }
    return BfSharedWorkerTypes.IMPORTED_MAIN_THREAD;
  }

  constructor(...args: Array<unknown>) {
    console.log("Constructor", this.workerType);
    switch (this.workerType) {
      case BfSharedWorkerTypes.IMPORTED_MAIN_THREAD:
      case BfSharedWorkerTypes.IMPORTED_WORKER:
        return this.initAsImported(...args);
      case BfSharedWorkerTypes.MAIN_THREAD:
        return this.initAsMainThread(...args);

      case BfSharedWorkerTypes.MAIN_WORKER:
        return this.initAsMainWorker(...args);
    }
  }

  private initAsMainThread(...args: Array<unknown>) {
    if (!this.supportsRunningAsMainThread) {
      throw new Error(
        "This worker does not support running as main thread",
      );
    }
    return this;
  }

  private initAsImported(...args: Array<unknown>) {
    const url = new URL(import.meta.resolve(this.mainUrl));
    this.worker = new SharedWorker(url, this.workerName);
    // Start the port for communication with the SharedWorker
    console.log("going to start");
    this.worker.port.start();
    const thisSelf: T = this as unknown as T; // Capture the actual type, including subclasses

    return new Proxy(thisSelf, {
      get(target, prop, receiver) {
        if (typeof prop !== "string") {
          return Reflect.get(target, prop, receiver);
        }

        const origMethod = target[prop];
        if (typeof origMethod === "function" && prop in target) {
          // Intercept method calls
          return (...methodArgs: unknown[]) => {
            // Redirect the call through sendMessage
            // Type safety for arguments is limited
            return thisSelf.sendMessage(
              prop as PublicMethodKeys<T>,
              methodArgs as MethodArguments<T>[PublicMethodKeys<T>],
            );
          };
        }
        return Reflect.get(target, prop, receiver);
      },
    });
  }

  private initAsMainWorker(...args: Array<unknown>) {
    console.log("initAsMainWorker", this.workerType);
    globalThis.addEventListener("connect", (e) => {
      this.onConnect((e as MessageEvent).ports as MessagePort[]);
    });
    globalThis.addEventListener("message", (e) => {
      if (e.data.type === "onconnect") {
        this.onConnect(e.data.ports as MessagePort[]);
      } else {
        console.log("worker message", e);
      }
    });
    return this;
  }

  private onConnect(ports: MessagePort[]) {
    console.log("onConnect");
    const port = ports[0];
    port.onmessage = this.receiveMessage.bind(this);
  }

  private isFunctionProperty<K extends keyof BfSharedWorker<T>>(
    prop: K
  ): this is BfSharedWorker<T> & Record<K, (...args: unknown[]) => unknown> {
    return typeof this[prop] === 'function';
  }
  

  private async receiveMessage(
    e: MessageEvent<{type: PublicMethodKeys<T>, args: MethodArguments<T>[PublicMethodKeys<T>], port: MessagePort}>,
  ) {
    const { type, args, port } = e.data;
    
    try {
      if (type in this && this.isFunctionProperty(type)) {
        // TypeScript should now recognize 'this[type]' as a function
        const returnable = await this[type](...args);
        port.postMessage(returnable);
      }
      throw new Error(`Method ${String(type)} not found`);
    } catch (error) {
      port.postMessage({ error });
    }

  }
  
  

  private sendMessage<K extends PublicMethodKeys<this>>(
    type: K,
    args: MethodArguments<this>[K],
    transferrables: Transferable[] = [],
  ): Promise<MethodReturnTypes<this>[K]> {
    // Implementation of sendMessage to handle the proxied method calls
    return new Promise((resolve, reject) => {
      const { port1, port2 } = new MessageChannel();
      const txfr = [
        port2,
        ...transferrables,
      ];
      port1.onmessage = (e) => {
        if (e.data.error) {
          reject(e.data.error);
        } else {
          resolve(e.data);
        }
      };
      if (this.worker) {
        console.log("sending message", type, args);
        return this.worker.port.postMessage({ port: port2, type, args }, txfr);
      }
      reject(new Error("Worker not initialized"));
    });
  }
  
  // these would normally be extended in a new worker

  ping() {
    console.log("ping");
    return "pong";
  }

  adder(a: number, b: number) {
    return a + b;
  }
}

if (import.meta.main) {
  const worker = new BfSharedWorker();
}