Overlay 2.0 (#12)

src/core/services/tokenizer/parser-enrichment-worker-runtime.ts

@@ -0,0 +1,149 @@
+import type { MergedToken } from '../../../types';
+import { createLogger } from '../../../logger';
+import { enrichTokensWithMecabPos1 } from './parser-enrichment-stage';
+
+const logger = createLogger('main:tokenizer');
+const DISABLE_WORKER_ENV = 'SUBMINER_DISABLE_MECAB_ENRICHMENT_WORKER';
+
+interface WorkerRequest {
+  id: number;
+  tokens: MergedToken[];
+  mecabTokens: MergedToken[] | null;
+}
+
+interface WorkerResponse {
+  id?: unknown;
+  result?: unknown;
+  error?: unknown;
+}
+
+type PendingRequest = {
+  resolve: (value: MergedToken[]) => void;
+  reject: (reason?: unknown) => void;
+};
+
+class ParserEnrichmentWorkerRuntime {
+  private worker: import('node:worker_threads').Worker | null = null;
+  private nextRequestId = 1;
+  private pending = new Map<number, PendingRequest>();
+  private initAttempted = false;
+
+  async enrichTokens(
+    tokens: MergedToken[],
+    mecabTokens: MergedToken[] | null,
+  ): Promise<MergedToken[]> {
+    const worker = await this.getWorker();
+    if (!worker) {
+      return enrichTokensWithMecabPos1(tokens, mecabTokens);
+    }
+
+    return new Promise<MergedToken[]>((resolve, reject) => {
+      const id = this.nextRequestId++;
+      this.pending.set(id, { resolve, reject });
+      const request: WorkerRequest = { id, tokens, mecabTokens };
+      worker.postMessage(request);
+    });
+  }
+
+  private async getWorker(): Promise<import('node:worker_threads').Worker | null> {
+    if (process.env[DISABLE_WORKER_ENV] === '1') {
+      return null;
+    }
+    if (this.worker) {
+      return this.worker;
+    }
+    if (this.initAttempted) {
+      return null;
+    }
+
+    this.initAttempted = true;
+
+    let workerThreads: typeof import('node:worker_threads');
+    try {
+      workerThreads = await import('node:worker_threads');
+    } catch {
+      return null;
+    }
+
+    let workerPath = '';
+    try {
+      workerPath = require.resolve('./parser-enrichment-worker-thread.js');
+    } catch {
+      return null;
+    }
+
+    try {
+      const worker = new workerThreads.Worker(workerPath);
+      worker.on('message', (message: WorkerResponse) => this.handleWorkerMessage(message));
+      worker.on('error', (error: Error) => this.handleWorkerFailure(error));
+      worker.on('exit', (code: number) => {
+        if (code !== 0) {
+          this.handleWorkerFailure(new Error(`parser enrichment worker exited with code ${code}`));
+        } else {
+          this.worker = null;
+        }
+      });
+      this.worker = worker;
+      return worker;
+    } catch (error) {
+      logger.debug(`Failed to start parser enrichment worker: ${(error as Error).message}`);
+      return null;
+    }
+  }
+
+  private handleWorkerMessage(message: WorkerResponse): void {
+    if (typeof message.id !== 'number') {
+      return;
+    }
+
+    const request = this.pending.get(message.id);
+    if (!request) {
+      return;
+    }
+    this.pending.delete(message.id);
+
+    if (typeof message.error === 'string' && message.error.length > 0) {
+      request.reject(new Error(message.error));
+      return;
+    }
+
+    if (!Array.isArray(message.result)) {
+      request.reject(new Error('Parser enrichment worker returned invalid payload'));
+      return;
+    }
+
+    request.resolve(message.result as MergedToken[]);
+  }
+
+  private handleWorkerFailure(error: Error): void {
+    logger.debug(
+      `Parser enrichment worker unavailable, falling back to main thread: ${error.message}`,
+    );
+    for (const pending of this.pending.values()) {
+      pending.reject(error);
+    }
+    this.pending.clear();
+
+    if (this.worker) {
+      this.worker.removeAllListeners();
+      this.worker = null;
+    }
+  }
+}
+
+let runtime: ParserEnrichmentWorkerRuntime | null = null;
+
+export async function enrichTokensWithMecabPos1Async(
+  tokens: MergedToken[],
+  mecabTokens: MergedToken[] | null,
+): Promise<MergedToken[]> {
+  if (!runtime) {
+    runtime = new ParserEnrichmentWorkerRuntime();
+  }
+
+  try {
+    return await runtime.enrichTokens(tokens, mecabTokens);
+  } catch {
+    return enrichTokensWithMecabPos1(tokens, mecabTokens);
+  }
+}