backend/omr-service/src/worker.ts

import { config } from './config.js';
import { predictPagesService } from './services/predict-pages.js';
import * as predictorService from './services/predictor.service.js';

interface PageTask {
	taskId: string;
	scoreId: string;
	pageIndex: number;
	imageData: Buffer;
}

interface AllPagesTask {
	taskId: string;
	scoreId: string;
	imageDataList: Buffer[];
	layoutList?: any[];
}

interface PredictPagesTask {
	taskId: string;
	scoreId: string;
	predictPages: {
		images: Array<{ data: Buffer; layout?: any; enableGauge?: boolean }>;
		outputWidth?: number;
		processes?: string[];
	};
}

type TaskItem = PageTask | AllPagesTask | PredictPagesTask;

function isAllPagesTask(task: TaskItem): task is AllPagesTask {
	return 'imageDataList' in task;
}

function isPredictPagesTask(task: TaskItem): task is PredictPagesTask {
	return 'predictPages' in task;
}

class TaskWorker {
	private queue: TaskItem[] = [];
	private running = 0;
	private maxConcurrent: number;

	constructor(maxConcurrent: number = 2) {
		this.maxConcurrent = maxConcurrent;
	}

	queueTask(task: TaskItem) {
		this.queue.push(task);
		this.processQueue();
	}

	private async processQueue() {
		if (this.running >= this.maxConcurrent) {
			return;
		}

		const task = this.queue.shift();
		if (!task) {
			return;
		}

		this.running++;

		try {
			if (isPredictPagesTask(task)) {
				await predictPagesService({
					taskId: task.taskId,
					scoreId: task.scoreId,
					...task.predictPages,
				});
			} else if (isAllPagesTask(task)) {
				if (task.layoutList?.length) {
					await predictorService.predictWithLayout({
						taskId: task.taskId,
						scoreId: task.scoreId,
						imageDataList: task.imageDataList,
						layoutList: task.layoutList,
					});
				} else {
					await predictorService.predictAll({
						taskId: task.taskId,
						scoreId: task.scoreId,
						imageDataList: task.imageDataList,
					});
				}
			} else {
				await predictorService.predictPage({
					taskId: task.taskId,
					scoreId: task.scoreId,
					pageIndex: task.pageIndex,
					imageData: task.imageData,
				});
			}
		} catch (error) {
			console.error(`Task ${task.taskId} failed:`, error);
		} finally {
			this.running--;
			this.processQueue();
		}
	}

	get pendingCount() {
		return this.queue.length;
	}

	get runningCount() {
		return this.running;
	}
}

export const taskWorker = new TaskWorker(config.taskWorker.maxConcurrent);