core: judge: add concurrency control to task consumer (#719)

Co-authored-by: undefined <i@undefined.moe>

packages/hydrojudge/src/hosts/builtin.ts

@@ -90,7 +90,10 @@ export async function postInit(ctx) {
         await (new JudgeTask(session, JSON.parse(JSON.stringify(Object.assign(rdoc, t))))).handle().catch(logger.error);
     };
     const parallelism = Math.max(getConfig('parallelism'), 2);
-    for (let i = 1; i < parallelism; i++) TaskModel.consume({ type: 'judge' }, handle);
+    const taskConsumer = TaskModel.consume({ type: 'judge' }, handle, true, parallelism);
+    ctx.on('system/setting', () => {
+        taskConsumer.setConcurrency(Math.max(getConfig('parallelism'), 2));
+    });
     TaskModel.consume({ type: 'judge', priority: { $gt: -50 } }, handle);
     TaskModel.consume({ type: 'generate' }, handle);
 }

packages/hydrooj/src/handler/judge.ts

@@ -21,14 +21,13 @@ import record from '../model/record';
 import * as setting from '../model/setting';
 import storage from '../model/storage';
 import * as system from '../model/system';
-import task from '../model/task';
+import task, { Consumer } from '../model/task';
 import user from '../model/user';
 import * as bus from '../service/bus';
 import { updateJudge } from '../service/monitor';
 import {
     ConnectionHandler, Handler, post, subscribe, Types,
 } from '../service/server';
-import { sleep } from '../utils';
 
 const logger = new Logger('judge');
 
@@ -45,6 +44,8 @@ function parseCaseResult(body: TestCase): Required<TestCase> {
 function processPayload(body: Partial<JudgeResultBody>) {
     const $set: Partial<RecordDoc> = {};
     const $push: any = {};
+    const $unset: any = {};
+    const $inc: any = {};
     if (body.cases?.length) {
         const c = body.cases.map(parseCaseResult);
         $push.testCases = { $each: c };
@@ -64,13 +65,18 @@ function processPayload(body: Partial<JudgeResultBody>) {
     if (Number.isFinite(body.memory)) $set.memory = body.memory;
     if (body.progress !== undefined) $set.progress = body.progress;
     if (body.subtasks) $set.subtasks = body.subtasks;
-    return { $set, $push };
+    if (body.addProgress) $inc.progress = body.addProgress;
+    return {
+        $set, $push, $unset, $inc,
+    };
 }
 
 export async function next(body: Partial<JudgeResultBody>) {
     body.rid = new ObjectId(body.rid);
-    const { $set, $push } = processPayload(body);
-    const rdoc = await record.update(body.domainId, body.rid, $set, $push, {}, body.addProgress ? { progress: body.addProgress } : {});
+    const {
+        $set, $push, $unset, $inc,
+    } = processPayload(body);
+    const rdoc = await record.update(body.domainId, body.rid, $set, $push, $unset, $inc);
     bus.broadcast('record/change', rdoc, $set, $push, body);
     return rdoc;
 }
@@ -210,25 +216,25 @@ export class JudgeFileUpdateHandler extends Handler {
     }
 }
 
-class JudgeConnectionHandler extends ConnectionHandler {
+export class JudgeConnectionHandler extends ConnectionHandler {
     category = '#judge';
-    processing: Task[] = [];
-    started = false;
-    closed = false;
     query: any = { type: { $in: ['judge', 'generate'] } };
-    rdocs: Record<string, RecordDoc> = {};
     concurrency = 1;
-    queue: PQueue = new PQueue({ concurrency: 1 });
-    callback: (res?: any) => void = null;
+    consumer: Consumer = null;
+    startTimeout: NodeJS.Timeout = null;
+    tasks: Record<string, {
+        resolve: (_: any) => void,
+        queue: PQueue,
+        domainId: string,
+        t: Task,
+    }> = {};
 
     async prepare() {
         logger.info('Judge daemon connected from ', this.request.ip);
         this.sendLanguageConfig();
         // TODO deprecated, just for compatibility
-        setTimeout(() => {
-            if (this.started || this.closed) return;
-            this.newTask().catch((e) => logger.error(e));
-            this.started = true;
+        this.startTimeout = setTimeout(() => {
+            this.consumer ||= task.consume(this.query, this.newTask.bind(this), true, this.concurrency);
         }, 15000);
     }
 
@@ -237,31 +243,23 @@ class JudgeConnectionHandler extends ConnectionHandler {
         this.send({ language: setting.langs });
     }
 
-    async newTask() {
-        while (!this.closed) {
-            /* eslint-disable no-await-in-loop */
-            if (this.processing.length >= this.concurrency) {
-                await Promise.race([
-                    sleep(500),
-                    new Promise((resolve) => {
-                        this.callback = resolve;
-                    }),
-                ]);
-                continue;
-            }
-            const t = await task.getFirst(this.query);
-            if (!t) {
-                await sleep(500);
-                continue;
-            }
+    async newTask(t: Task) {
         const rdoc = await record.get(t.domainId, t.rid);
-            if (!rdoc) continue;
+        if (!rdoc) return;
 
+        const rid = rdoc._id.toHexString();
+        let resolve: (_: any) => void;
+        const p = new Promise((r) => { resolve = r; });
+        this.tasks[rid] = {
+            queue: new PQueue({ concurrency: 1 }),
+            domainId: rdoc.domainId,
+            resolve,
+            t,
+        };
         this.send({ task: { ...rdoc, ...t } });
-            this.rdocs[rdoc._id.toHexString()] = rdoc;
-            this.processing.push(t);
-            await next({ status: STATUS.STATUS_FETCHED, domainId: rdoc.domainId, rid: rdoc._id });
-        }
+        this.tasks[rid].queue.add(() => next({ status: STATUS.STATUS_FETCHED, domainId: rdoc.domainId, rid: rdoc._id }));
+        await p;
+        delete this.tasks[rid];
     }
 
     async message(msg) {
@@ -271,47 +269,47 @@ class JudgeConnectionHandler extends ConnectionHandler {
             logger[method]('%o', omit(msg, keys));
         }
         if (['next', 'end'].includes(msg.key)) {
-            const rdoc = this.rdocs[msg.rid];
-            if (!rdoc) return;
-            msg.domainId = rdoc.domainId;
-            if (msg.key === 'next') this.queue.add(() => next(msg));
+            const t = this.tasks[msg.key];
+            if (!t) return;
+            msg.domainId = t.domainId;
+            if (msg.key === 'next') t.queue.add(() => next(msg));
             if (msg.key === 'end') {
-                this.processing = this.processing.filter((t) => t.rid.toHexString() !== msg.rid);
-                delete this.rdocs[msg.rid];
-                this.callback?.();
-                if (!msg.nop) {
-                    this.queue.add(() => end({ judger: this.user._id, ...msg }).catch((e) => logger.error(e)));
-                }
+                if (!msg.nop) t.queue.add(() => end({ judger: this.user._id, ...msg }));
+                t.resolve(null);
             }
         } else if (msg.key === 'status') {
             await updateJudge(msg.info);
         } else if (msg.key === 'prio' && typeof msg.prio === 'number') {
             // TODO deprecated, use config instead
             this.query.priority = { $gt: msg.prio };
+            this.consumer?.setQuery(this.query);
         } else if (msg.key === 'config') {
             if (Number.isSafeInteger(msg.prio)) {
                 this.query.priority = { $gt: msg.prio };
+                this.consumer?.setQuery(this.query);
             }
             if (Number.isSafeInteger(msg.concurrency) && msg.concurrency > 0) {
                 this.concurrency = msg.concurrency;
-                this.callback?.();
+                this.consumer?.setConcurrency(msg.concurrency);
             }
             if (msg.lang instanceof Array && msg.lang.every((i) => typeof i === 'string')) {
                 this.query.lang = { $in: msg.lang };
+                this.consumer?.setQuery(this.query);
             }
         } else if (msg.key === 'start') {
-            if (this.started) return;
-            this.newTask().catch((e) => logger.error(e));
-            this.started = true;
+            clearTimeout(this.startTimeout);
+            this.consumer ||= task.consume(this.query, this.newTask.bind(this), true, this.concurrency);
         }
     }
 
     async cleanup() {
-        this.closed = true;
+        clearTimeout(this.startTimeout);
+        this.consumer?.destroy();
         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);
+        await Promise.all(Object.values(this.tasks).map(async ({ t }) => {
+            const rdoc = await record.reset(t.domainId, t.rid, false);
+            bus.broadcast('record/change', rdoc);
+            return task.add(t);
         }));
     }
 }

packages/hydrooj/src/handler/record.ts

@@ -253,6 +253,8 @@ class RecordMainConnectionHandler extends ConnectionHandler {
     pretest = false;
     tdoc: Tdoc;
     applyProjection = false;
+    queue: Map<string, () => Promise<any>> = new Map();
+    throttleQueueClear: () => void;
 
     @param('tid', Types.ObjectId, true)
     @param('pid', Types.ProblemId, true)
@@ -301,6 +303,7 @@ class RecordMainConnectionHandler extends ConnectionHandler {
             this.checkPriv(PRIV.PRIV_MANAGE_ALL_DOMAIN);
             this.allDomain = true;
         }
+        this.throttleQueueClear = throttle(this.queueClear, 100, { trailing: true });
     }
 
     async message(msg: { rids: string[] }) {
@@ -338,15 +341,25 @@ class RecordMainConnectionHandler extends ConnectionHandler {
             if (!this.user.hasPerm(PERM.PERM_VIEW_PROBLEM)) pdoc = null;
         }
         if (this.applyProjection && typeof rdoc.input !== 'string') rdoc = contest.applyProjection(tdoc, rdoc, this.user);
-        if (this.pretest) this.send({ rdoc: omit(rdoc, ['code', 'input']) });
+        if (this.pretest) this.queueSend(rdoc._id.toHexString(), async () => ({ rdoc: omit(rdoc, ['code', 'input']) }));
         else {
-            this.send({
+            this.queueSend(rdoc._id.toHexString(), async () => ({
                 html: await this.renderHTML('record_main_tr.html', {
                     rdoc, udoc, pdoc, tdoc, allDomain: this.allDomain,
                 }),
-            });
+            }));
         }
     }
+
+    queueSend(rid: string, fn: () => Promise<any>) {
+        this.queue.set(rid, fn);
+        this.throttleQueueClear();
+    }
+
+    async queueClear() {
+        await Promise.all([...this.queue.values()].map(async (fn) => this.send(await fn())));
+        this.queue.clear();
+    }
 }
 
 class RecordDetailConnectionHandler extends ConnectionHandler {

packages/hydrooj/src/model/schedule.ts

@@ -31,10 +31,10 @@ class Consumer {
     consuming: boolean;
     running?: any;
 
-    constructor(public filter: any, public func: Function, public destoryOnError = true) {
+    constructor(public filter: any, public func: Function, public destroyOnError = true) {
         this.consuming = true;
         this.consume();
-        bus.on('app/exit', this.destory);
+        bus.on('app/exit', this.destroy);
     }
 
     async consume() {
@@ -51,12 +51,12 @@ class Consumer {
                 } else await sleep(1000);
             } catch (err) {
                 logger.error(err);
-                if (this.destoryOnError) this.destory();
+                if (this.destroyOnError) this.destroy();
             }
         }
     }
 
-    async destory() {
+    async destroy() {
         this.consuming = false;
     }
 }

packages/hydrooj/src/model/task.ts

@@ -25,37 +25,70 @@ async function getFirst(query: Filter<Task>) {
     return null;
 }
 
-class Consumer {
+export class Consumer {
     consuming: boolean;
+    processing: Set<Task> = new Set();
     running?: any;
+    notify: (res?: any) => void;
 
-    constructor(public filter: any, public func: Function, public destoryOnError = true) {
+    constructor(public filter: any, public func: (t: Task) => Promise<void>, public destroyOnError = true, private concurrency = 1) {
         this.consuming = true;
         this.consume();
-        bus.on('app/exit', this.destory);
+        bus.on('app/exit', this.destroy);
     }
 
     async consume() {
         while (this.consuming) {
             try {
+                if (this.processing.size >= this.concurrency) {
+                    // eslint-disable-next-line no-await-in-loop
+                    await new Promise((resolve) => {
+                        this.notify = resolve;
+                    });
+                    continue;
+                }
                 // eslint-disable-next-line no-await-in-loop
                 const res = await getFirst(this.filter);
-                if (res) {
-                    this.running = res;
+                if (!res) {
+                    let timeout: NodeJS.Timeout = null;
                     // eslint-disable-next-line no-await-in-loop
-                    await this.func(res);
-                    this.running = null;
-                    // eslint-disable-next-line no-await-in-loop
-                } else await sleep(1000);
+                    await new Promise((resolve) => {
+                        timeout = setTimeout(resolve, 1000 / (this.concurrency - this.processing.size));
+                        this.notify = resolve;
+                    });
+                    clearTimeout(timeout);
+                    continue;
+                }
+                this.processing.add(res);
+                this.func(res)
+                    .catch((err) => {
+                        logger.error(err);
+                        if (this.destroyOnError) this.destroy();
+                    })
+                    .finally(() => {
+                        this.processing.delete(res);
+                        this.notify?.();
+                    });
             } catch (err) {
                 logger.error(err);
-                if (this.destoryOnError) this.destory();
+                if (this.destroyOnError) this.destroy();
             }
         }
     }
 
-    async destory() {
+    async destroy() {
         this.consuming = false;
+        this.notify?.();
+    }
+
+    setConcurrency(concurrency: number) {
+        this.concurrency = concurrency;
+        this.notify?.();
+    }
+
+    setQuery(query: string) {
+        this.filter = query;
+        this.notify?.();
     }
 }
 
@@ -95,8 +128,8 @@ class TaskModel {
 
     static getFirst = getFirst;
 
-    static consume(query: any, cb: Function, destoryOnError = true) {
-        return new Consumer(query, cb, destoryOnError);
+    static consume(query: any, cb: (t: Task) => Promise<void>, destroyOnError = true, concurrency = 1) {
+        return new Consumer(query, cb, destroyOnError, concurrency);
     }
 }
 

packages/hydrooj/src/service/server.ts

@@ -417,7 +417,6 @@ 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
@@ -606,6 +605,7 @@ ${ctx.response.status} ${endTime - startTime}ms ${ctx.response.length}`);
                 socket.close(1003, 'Websocket Error');
             } catch (e) { }
         });
+        socket.pause();
         const ctx: any = app.createContext(request, {} as any);
         await domainLayer(ctx, () => baseLayer(ctx, () => layers[1](ctx, () => userLayer(ctx, () => { }))));
         for (const manager of router.wsStack) {

packages/ui-default/components/editor/index.ts

@@ -256,7 +256,7 @@ export default class Editor extends DOMAttachedObject {
     if (hasFocus) this.focus();
   }
 
-  destory() {
+  destroy() {
     this.detach();
     if (this.vditor?.destroy) this.vditor.destroy();
     else if (this.editor?.dispose) this.editor.dispose();

packages/ui-default/components/preview/preview.page.ts

@@ -34,7 +34,7 @@ async function startEdit(filename, value, fileCategory = 'file') {
   });
   const action = await promise;
   value = (editor.value() as string).replace(/\r\n/g, '\n');
-  editor.destory();
+  editor.destroy();
   if (action !== 'ok') return null;
   return value;
 }

packages/ui-default/components/problemconfig/ProblemConfigEditor.tsx

@@ -143,7 +143,7 @@ export default connect(mapStateToProps, mapDispatchToProps)(class MonacoEditor e
   }
 
   componentWillUnmount() {
-    if (this.vjEditor) this.vjEditor.destory();
+    if (this.vjEditor) this.vjEditor.destroy();
     if (this.model) this.model.dispose();
     if (this.editor) this.editor.dispose();
     this.disposable.map((i) => i.dispose());