From 6b759cbe637969083aa185fc5037fdc6959eccd9 Mon Sep 17 00:00:00 2001 From: Yang Gao <6821729+criyle@users.noreply.github.com> Date: Sun, 28 Jan 2024 01:23:05 +0800 Subject: [PATCH] core: judge: add concurrency control to task consumer (#719) Co-authored-by: undefined --- packages/hydrojudge/src/hosts/builtin.ts | 5 +- packages/hydrooj/src/handler/judge.ts | 114 +++++++++--------- packages/hydrooj/src/handler/record.ts | 19 ++- packages/hydrooj/src/model/schedule.ts | 8 +- packages/hydrooj/src/model/task.ts | 59 +++++++-- packages/hydrooj/src/service/server.ts | 2 +- .../ui-default/components/editor/index.ts | 2 +- .../components/preview/preview.page.ts | 2 +- .../problemconfig/ProblemConfigEditor.tsx | 2 +- 9 files changed, 130 insertions(+), 83 deletions(-) diff --git a/packages/hydrojudge/src/hosts/builtin.ts b/packages/hydrojudge/src/hosts/builtin.ts index dcc32e849..968e14cd0 100644 --- a/packages/hydrojudge/src/hosts/builtin.ts +++ b/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); } diff --git a/packages/hydrooj/src/handler/judge.ts b/packages/hydrooj/src/handler/judge.ts index 5607debc5..07f98623c 100644 --- a/packages/hydrooj/src/handler/judge.ts +++ b/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 { function processPayload(body: Partial) { const $set: Partial = {}; 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) { 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) { 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 = {}; concurrency = 1; - queue: PQueue = new PQueue({ concurrency: 1 }); - callback: (res?: any) => void = null; + consumer: Consumer = null; + startTimeout: NodeJS.Timeout = null; + tasks: Record 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; - } - const rdoc = await record.get(t.domainId, t.rid); - if (!rdoc) continue; + async newTask(t: Task) { + const rdoc = await record.get(t.domainId, t.rid); + if (!rdoc) return; - 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 }); - } + 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.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); })); } } diff --git a/packages/hydrooj/src/handler/record.ts b/packages/hydrooj/src/handler/record.ts index 78bf82314..3f8ceaebc 100644 --- a/packages/hydrooj/src/handler/record.ts +++ b/packages/hydrooj/src/handler/record.ts @@ -253,6 +253,8 @@ class RecordMainConnectionHandler extends ConnectionHandler { pretest = false; tdoc: Tdoc; applyProjection = false; + queue: Map Promise> = 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) { + 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 { diff --git a/packages/hydrooj/src/model/schedule.ts b/packages/hydrooj/src/model/schedule.ts index ef71f779a..92e3f5878 100644 --- a/packages/hydrooj/src/model/schedule.ts +++ b/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; } } diff --git a/packages/hydrooj/src/model/task.ts b/packages/hydrooj/src/model/task.ts index 5adb9f959..008daad73 100644 --- a/packages/hydrooj/src/model/task.ts +++ b/packages/hydrooj/src/model/task.ts @@ -25,37 +25,70 @@ async function getFirst(query: Filter) { return null; } -class Consumer { +export class Consumer { consuming: boolean; + processing: Set = 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, 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, destroyOnError = true, concurrency = 1) { + return new Consumer(query, cb, destroyOnError, concurrency); } } diff --git a/packages/hydrooj/src/service/server.ts b/packages/hydrooj/src/service/server.ts index 10d69c76e..991646487 100644 --- a/packages/hydrooj/src/service/server.ts +++ b/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) { diff --git a/packages/ui-default/components/editor/index.ts b/packages/ui-default/components/editor/index.ts index 8c026d452..ef6bbbba3 100644 --- a/packages/ui-default/components/editor/index.ts +++ b/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(); diff --git a/packages/ui-default/components/preview/preview.page.ts b/packages/ui-default/components/preview/preview.page.ts index da7e195c7..252b0bd0a 100644 --- a/packages/ui-default/components/preview/preview.page.ts +++ b/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; } diff --git a/packages/ui-default/components/problemconfig/ProblemConfigEditor.tsx b/packages/ui-default/components/problemconfig/ProblemConfigEditor.tsx index f3c15721a..d381d605b 100644 --- a/packages/ui-default/components/problemconfig/ProblemConfigEditor.tsx +++ b/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());