regex-tools/src/regex/execution/EngineSupervisor.ts

import type {
  RegexExecutionRequest,
  RegexExecutionResult,
  RegexReplacementRequest,
  RegexReplacementResult,
} from "../model/match";
import { WorkerRequestError, WorkerSupervisor } from "./WorkerSupervisor";
import type {
  EngineWorkerOperation,
  EngineWorkerResult,
} from "./worker-protocol";
import {
  ENGINE_WORKER_REGISTRY,
  type EngineWorkerRegistry,
} from "./engine-registry";
import type { RegexEngineInfo, RegexFlavourId } from "../model/flavour";

export type EngineRuntimeState =
  | { readonly status: "unloaded" }
  | { readonly status: "running"; readonly operation: "execute" | "replace" }
  | { readonly status: "ready" }
  | { readonly status: "recovering"; readonly reason: string }
  | { readonly status: "failed"; readonly error: Error };

export class EngineSupervisor {
  private readonly registry: EngineWorkerRegistry;
  private readonly supervisors = new Map<
    RegexFlavourId,
    WorkerSupervisor<EngineWorkerOperation, EngineWorkerResult>
  >();
  private readonly loaded = new Map<
    RegexFlavourId,
    {
      readonly generation: number;
      readonly info: RegexEngineInfo;
    }
  >();
  private readonly loading = new Map<
    RegexFlavourId,
    Promise<RegexEngineInfo>
  >();
  private readonly operationRevisions = new Map<RegexFlavourId, number>();
  private disposed = false;

  constructor(registry: EngineWorkerRegistry = ENGINE_WORKER_REGISTRY) {
    this.registry = registry;
  }

  private supervisorFor(
    flavour: RegexFlavourId,
  ): WorkerSupervisor<EngineWorkerOperation, EngineWorkerResult> {
    if (this.disposed) {
      throw new Error("Engine supervisor has been disposed.");
    }
    const current = this.supervisors.get(flavour);
    if (current) return current;
    const registration = this.registry.require(flavour);
    const supervisor = new WorkerSupervisor<
      EngineWorkerOperation,
      EngineWorkerResult
    >(registration.label, registration.createWorker);
    this.supervisors.set(flavour, supervisor);
    return supervisor;
  }

  async load(flavour: RegexFlavourId): Promise<RegexEngineInfo> {
    const registration = this.registry.require(flavour);
    const supervisor = this.supervisorFor(flavour);
    const loaded = this.loaded.get(flavour);
    if (
      loaded &&
      supervisor.hasWorker &&
      loaded.generation === supervisor.currentGeneration
    ) {
      return loaded.info;
    }
    const current = this.loading.get(flavour);
    if (current && supervisor.hasWorker) return current;
    const loading: Promise<RegexEngineInfo> = supervisor
      .run({ kind: "load" }, registration.startupTimeoutMs)
      .then((response) => {
        if (response.kind !== "load") {
          throw new Error(
            "Engine worker returned the wrong load response kind",
          );
        }
        if (response.info.flavour !== flavour) {
          throw new Error(
            `Engine worker for ${flavour} loaded a ${response.info.flavour} runtime.`,
          );
        }
        this.loaded.set(flavour, {
          generation: supervisor.currentGeneration,
          info: response.info,
        });
        return response.info;
      })
      .finally(() => {
        if (this.loading.get(flavour) === loading) {
          this.loading.delete(flavour);
        }
      });
    this.loading.set(flavour, loading);
    return loading;
  }

  private async prepareOperation(flavour: RegexFlavourId): Promise<{
    readonly revision: number;
    readonly supervisor: WorkerSupervisor<
      EngineWorkerOperation,
      EngineWorkerResult
    >;
  }> {
    const revision = (this.operationRevisions.get(flavour) ?? 0) + 1;
    this.operationRevisions.set(flavour, revision);
    const supervisor = this.supervisorFor(flavour);
    if (supervisor.isRunning) supervisor.cancel();
    await this.load(flavour);
    if (this.operationRevisions.get(flavour) !== revision) {
      throw new WorkerRequestError(
        "cancelled",
        `${this.registry.require(flavour).label} request was superseded`,
      );
    }
    return { revision, supervisor };
  }

  async execute(
    request: RegexExecutionRequest,
    timeoutMs: number,
  ): Promise<RegexExecutionResult> {
    const { supervisor } = await this.prepareOperation(request.flavour);
    const response = await supervisor.run(
      { kind: "execute", request },
      timeoutMs,
    );
    if (response.kind !== "execute") {
      throw new Error("Engine worker returned the wrong response kind");
    }
    if (response.result.engine.flavour !== request.flavour) {
      throw new Error(
        `Engine worker for ${request.flavour} returned a ${response.result.engine.flavour} result.`,
      );
    }
    return response.result;
  }

  async replace(
    request: RegexReplacementRequest,
    timeoutMs: number,
  ): Promise<RegexReplacementResult> {
    const { supervisor } = await this.prepareOperation(request.flavour);
    const response = await supervisor.run(
      { kind: "replace", request },
      timeoutMs,
    );
    if (response.kind !== "replace") {
      throw new Error("Engine worker returned the wrong response kind");
    }
    if (response.result.execution.engine.flavour !== request.flavour) {
      throw new Error(
        `Engine worker for ${request.flavour} returned a ${response.result.execution.engine.flavour} result.`,
      );
    }
    return response.result;
  }

  cancel(): void {
    for (const [flavour, supervisor] of this.supervisors) {
      this.operationRevisions.set(
        flavour,
        (this.operationRevisions.get(flavour) ?? 0) + 1,
      );
      supervisor.cancel();
    }
  }

  dispose(): void {
    if (this.disposed) return;
    this.disposed = true;
    for (const supervisor of this.supervisors.values()) supervisor.dispose();
    this.supervisors.clear();
    this.loaded.clear();
    this.loading.clear();
    this.operationRevisions.clear();
  }
}