File size: 2,406 Bytes
6f1c297
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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);