From 19c3d9e95c3cb0caf54ec32ad9f8616b31f37911 Mon Sep 17 00:00:00 2001
From: Yang Gao <6821729+criyle@users.noreply.github.com>
Date: Wed, 24 Jan 2024 10:28:25 +0800
Subject: [PATCH] core: judge: optimize task assign (#718)

---
 packages/hydrojudge/src/hosts/hydro.ts |  1 +
 packages/hydrooj/src/handler/judge.ts  | 12 ++++++++++--
 packages/hydrooj/src/service/server.ts |  2 ++
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/packages/hydrojudge/src/hosts/hydro.ts b/packages/hydrojudge/src/hosts/hydro.ts
index 94a20ebb9..9275991ac 100644
--- a/packages/hydrojudge/src/hosts/hydro.ts
+++ b/packages/hydrojudge/src/hosts/hydro.ts
@@ -224,6 +224,7 @@ export default class Hydro implements Session {
                         this.ws.send(JSON.stringify({ key: 'status', info: { mid, ...inf } }));
                     }, 1200000);
                 }
+                this.ws.send(content);
                 resolve(null);
             });
         });
diff --git a/packages/hydrooj/src/handler/judge.ts b/packages/hydrooj/src/handler/judge.ts
index fdeb30f7e..7c449510a 100644
--- a/packages/hydrooj/src/handler/judge.ts
+++ b/packages/hydrooj/src/handler/judge.ts
@@ -217,6 +217,7 @@ class JudgeConnectionHandler extends ConnectionHandler {
     query: any = { type: { $in: ['judge', 'generate'] } };
     rdocs: Record<string, RecordDoc> = {};
     concurrency = 1;
+    callback: (res?: any) => void = null;
 
     async prepare() {
         logger.info('Judge daemon connected from ', this.request.ip);
@@ -235,7 +236,12 @@ class JudgeConnectionHandler extends ConnectionHandler {
         while (!this.closed) {
             /* eslint-disable no-await-in-loop */
             if (this.processing.length >= this.concurrency) {
-                await sleep(500);
+                await Promise.race([
+                    sleep(500),
+                    new Promise((resolve) => {
+                        this.callback = resolve;
+                    }),
+                ]);
                 continue;
             }
             const t = await task.getFirst(this.query);
@@ -264,9 +270,10 @@ class JudgeConnectionHandler extends ConnectionHandler {
             if (!rdoc) return;
             if (msg.key === 'next') await next({ ...msg, domainId: rdoc.domainId });
             if (msg.key === 'end') {
-                if (!msg.nop) await end({ judger: this.user._id, ...msg, domainId: rdoc.domainId }).catch((e) => logger.error(e));
                 this.processing = this.processing.filter((t) => t.rid.toHexString() !== msg.rid);
                 delete this.rdocs[msg.rid];
+                this.callback?.();
+                if (!msg.nop) await end({ judger: this.user._id, ...msg, domainId: rdoc.domainId }).catch((e) => logger.error(e));
             }
         } else if (msg.key === 'status') {
             await updateJudge(msg.info);
@@ -279,6 +286,7 @@ class JudgeConnectionHandler extends ConnectionHandler {
             }
             if (Number.isSafeInteger(msg.concurrency) && msg.concurrency > 0) {
                 this.concurrency = msg.concurrency;
+                this.callback?.();
             }
             if (msg.lang instanceof Array && msg.lang.every((i) => typeof i === 'string')) {
                 this.query.lang = { $in: msg.lang };
diff --git a/packages/hydrooj/src/service/server.ts b/packages/hydrooj/src/service/server.ts
index 940150c09..4e628d23d 100644
--- a/packages/hydrooj/src/service/server.ts
+++ b/packages/hydrooj/src/service/server.ts
@@ -417,6 +417,7 @@ export function Connection(
                 h.compression = new Shorty();
                 conn.send('shorty');
             }
+            conn.pause();
             if (h._prepare) await h._prepare(args);
             if (h.prepare) await h.prepare(args);
             // eslint-disable-next-line @typescript-eslint/no-shadow
@@ -449,6 +450,7 @@ export function Connection(
                 h.message?.(JSON.parse(e.data.toString()));
             };
             conn.onclose = clean;
+            conn.resume();
             await bus.parallel('connection/active', h);
         } catch (e) {
             await h.onerror(e);