From 599a177f462f65a2a444d5e21c9b7e6ebd954abb Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Mon, 27 Nov 2023 23:59:02 +0800
Subject: [PATCH] core&judge: allow concurrency setting within single
 connection

---
 packages/hydrojudge/src/hosts/hydro.ts |  7 ++++--
 packages/hydrooj/src/handler/judge.ts  | 34 ++++++++++++++++++--------
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/packages/hydrojudge/src/hosts/hydro.ts b/packages/hydrojudge/src/hosts/hydro.ts
index d1f93e3a4..0f6651a04 100644
--- a/packages/hydrojudge/src/hosts/hydro.ts
+++ b/packages/hydrojudge/src/hosts/hydro.ts
@@ -178,8 +178,11 @@ export default class Hydro {
                 Authorization: `Bearer ${this.config.cookie.split('sid=')[1].split(';')[0]}`,
             },
         });
-        const content = this.config.minPriority !== undefined
-            ? `{"key":"prio","prio":${this.config.minPriority}}`
+        const config: { prio?: number, concurrency?: number } = {};
+        if (this.config.minPriority !== undefined) config.prio = this.config.minPriority;
+        if (this.config.concurrency !== undefined) config.concurrency = this.config.concurrency;
+        const content = Object.keys(config).length
+            ? JSON.stringify({ key: 'config', ...config })
             : '{"key":"ping"}';
         setInterval(() => this.ws?.send?.(content), 30000);
         this.ws.on('message', (data) => {
diff --git a/packages/hydrooj/src/handler/judge.ts b/packages/hydrooj/src/handler/judge.ts
index afb5f1450..02523cb2b 100644
--- a/packages/hydrooj/src/handler/judge.ts
+++ b/packages/hydrooj/src/handler/judge.ts
@@ -3,7 +3,7 @@ import yaml from 'js-yaml';
 import { omit } from 'lodash';
 import { ObjectId } from 'mongodb';
 import {
-    JudgeResultBody, ProblemConfigFile, RecordDoc, TestCase,
+    JudgeResultBody, ProblemConfigFile, RecordDoc, Task, TestCase,
 } from '../interface';
 import { Logger } from '../logger';
 import * as builtin from '../model/builtin';
@@ -191,11 +191,12 @@ export class SubmissionDataDownloadHandler extends Handler {
 
 class JudgeConnectionHandler extends ConnectionHandler {
     category = '#judge';
-    processing: any = null;
+    processing: Task[] = [];
     closed = false;
     query: any = { type: 'judge' };
     rdocs: Record<string, RecordDoc> = {};
     ip: string;
+    concurrency = 1;
 
     async prepare() {
         logger.info('Judge daemon connected from ', this.request.ip);
@@ -211,7 +212,7 @@ class JudgeConnectionHandler extends ConnectionHandler {
     }
 
     async newTask() {
-        if (this.processing) return;
+        if (this.processing.length >= this.concurrency) return;
         let t;
         let rdoc: RecordDoc;
         while (!t) {
@@ -225,7 +226,7 @@ class JudgeConnectionHandler extends ConnectionHandler {
         }
         this.send({ task: { ...rdoc, ...t } });
         this.rdocs[rdoc._id.toHexString()] = rdoc;
-        this.processing = t;
+        this.processing.push(t);
         const $set = { status: builtin.STATUS.STATUS_FETCHED };
         rdoc = await record.update(t.domainId, t.rid, $set, {});
         bus.broadcast('record/change', rdoc, $set, {});
@@ -243,7 +244,7 @@ class JudgeConnectionHandler extends ConnectionHandler {
             if (msg.key === 'next') await next({ ...msg, rdoc });
             if (msg.key === 'end') {
                 if (!msg.nop) await end({ judger: this.user._id, ...msg, rdoc }).catch((e) => logger.error(e));
-                this.processing = null;
+                this.processing = this.processing.filter((t) => t.rid.toHexString() !== msg.rid);
                 delete this.rdocs[msg.rid];
                 await this.newTask();
             }
@@ -253,16 +254,29 @@ class JudgeConnectionHandler extends ConnectionHandler {
             this.query.priority = { $gt: msg.prio };
         } else if (msg.key === 'lang' && msg.lang instanceof Array && msg.lang.every((i) => typeof i === 'string')) {
             this.query.lang = { $in: msg.lang };
+        } else if (msg.key === 'config') {
+            if (Number.isSafeInteger(msg.prio)) {
+                this.query.priority = { $gt: msg.prio };
+            }
+            if (Number.isSafeInteger(msg.concurrency) && msg.concurrency > 0) {
+                const old = this.concurrency;
+                this.concurrency = msg.concurrency;
+                if (old <= this.concurrency) {
+                    for (let i = old; i < this.concurrency; i++) {
+                        await this.newTask(); // eslint-disable-line no-await-in-loop
+                    }
+                }
+            }
         }
     }
 
     async cleanup() {
-        logger.info('Judge daemon disconnected from ', this.request.ip);
-        if (this.processing) {
-            await record.reset(this.processing.domainId, this.processing.rid, false);
-            await task.add(this.processing);
-        }
         this.closed = true;
+        logger.info('Judge daemon disconnected from ', this.request.ip);
+        await Promise.all(this.processing.map(async (t) => {
+            await record.reset(t.domainId, t.rid, false);
+            return await task.add(t);
+        }));
     }
 }