judge: let sandbox client manages cache clean up & fix cache clean up for parallel compile tasks (#955)

packages/hydrojudge/src/compile.ts

@@ -2,7 +2,7 @@ import { LangConfig, STATUS } from '@hydrooj/common';
 import { CompileError } from './error';
 import { Execute } from './interface';
 import {
-    CopyIn, CopyInFile, del, runQueued,
+    CopyIn, CopyInFile, runQueued,
 } from './sandbox';
 import { compilerText } from './utils';
 
@@ -14,7 +14,7 @@ export default async function compile(
     const command = copyIn['compile.sh'] ? '/bin/bash compile.sh' : lang.compile;
     if (command) {
         const {
-            status, stdout, stderr, fileIds,
+            status, stdout, stderr, fileIds, [Symbol.asyncDispose]: cleanup,
         } = await runQueued(
             command,
             {
@@ -36,7 +36,7 @@ export default async function compile(
         return {
             execute,
             copyIn: { ...copyIn, [target]: { fileId: fileIds[target] } },
-            clean: () => del(fileIds[target]),
+            clean: async () => await cleanup(),
             _cacheable: target,
         };
     }

packages/hydrojudge/src/judge/default.ts

@@ -1,14 +1,14 @@
 import { NormalizedCase, STATUS } from '@hydrooj/common';
 import checkers from '../checkers';
 import { runFlow } from '../flow';
-import { del, runQueued } from '../sandbox';
+import { runQueued } from '../sandbox';
 import signals from '../signals';
 import { Context, ContextSubTask } from './interface';
 
 function judgeCase(c: NormalizedCase) {
     return async (ctx: Context, ctxSubtask: ContextSubTask, runner?: Function) => {
         const { address_space_limit, process_limit } = ctx.session.getLang(ctx.lang);
-        const res = await runQueued(
+        await using res = await runQueued(
             ctx.execute.execute,
             {
                 stdin: { src: c.input },
@@ -57,7 +57,6 @@ function judgeCase(c: NormalizedCase) {
             if (code < 32 && signalled) message = signals[code];
             else message = { message: 'Your program returned {0}.', params: [code] };
         }
-        await Promise.allSettled(Object.values(res.fileIds).map((id) => del(id)));
         if (runner && ctx.rerun && c.time <= 5000 && status === STATUS.STATUS_TIME_LIMIT_EXCEEDED) {
             ctx.rerun--;
             return await runner(ctx, ctxSubtask);

packages/hydrojudge/src/judge/generate.ts

@@ -27,7 +27,7 @@ export const judge = async (ctx: JudgeTask) => {
     let totalStatus = 0;
 
     async function runGenerator(i: number) {
-        const res = await runQueued(
+        await using res = await runQueued(
             `${executeGenerator.execute} ${i}`,
             {
                 stdin: { content: ctx.input || '' },
@@ -41,9 +41,8 @@ export const judge = async (ctx: JudgeTask) => {
             1,
         );
         const tmp = path.join(tmpdir(), `${ctx.request.rid}.${i}.in`);
-        ctx.clean.push(() => {
+        ctx.pushClean(() => {
             if (fs.existsSync(tmp)) fs.removeSync(tmp);
-            return res.fileIds['stdout'] ? client.deleteFile(res.fileIds['stdout']) : Promise.resolve();
         });
         const {
             code, signalled, time, memory, fileIds, stderr,
@@ -80,7 +79,7 @@ export const judge = async (ctx: JudgeTask) => {
         return status === STATUS.STATUS_ACCEPTED ? tmp : null;
     }
     async function runStd(i: number, stdin: CopyInFile) {
-        const res = await runQueued(
+        await using res = await runQueued(
             `${executeStd.execute} ${i}`,
             {
                 stdin,
@@ -94,9 +93,8 @@ export const judge = async (ctx: JudgeTask) => {
             1,
         );
         const tmp = path.join(tmpdir(), `${ctx.request.rid}.${i}.out`);
-        ctx.clean.push(() => {
+        ctx.pushClean(() => {
             if (fs.existsSync(tmp)) fs.removeSync(tmp);
-            return res.fileIds['stdout'] ? client.deleteFile(res.fileIds['stdout']) : Promise.resolve();
         });
         const {
             code, signalled, time, memory, fileIds, stderr,

packages/hydrojudge/src/judge/hack.ts

@@ -2,7 +2,7 @@
 import { STATUS } from '@hydrooj/common';
 import { fs, parseMemoryMB, parseTimeMS } from '@hydrooj/utils';
 import checkers from '../checkers';
-import { del, runQueued } from '../sandbox';
+import { runQueued } from '../sandbox';
 import signals from '../signals';
 import { Context } from './interface';
 
@@ -12,9 +12,12 @@ export async function judge(ctx: Context) {
         ctx.compile(ctx.lang, ctx.code),
         ctx.compileLocalFile('checker', ctx.config.checker, ctx.config.checker_type),
         ctx.compileLocalFile('validator', ctx.config.validator),
-        ctx.session.fetchFile(null, { [ctx.files.hack]: '' }),
+        (async () => {
+            const f = await ctx.session.fetchFile(null, { [ctx.files.hack]: '' });
+            ctx.pushClean(() => fs.unlink(f));
+            return f;
+        })(),
     ]);
-    ctx.clean.push(() => fs.unlink(input));
     ctx.next({ status: STATUS.STATUS_JUDGING, progress: 0 });
     const validateResult = await runQueued(
         validator.execute,
@@ -31,7 +34,7 @@ export async function judge(ctx: Context) {
         return ctx.end({ status: STATUS.STATUS_FORMAT_ERROR, message });
     }
     const { address_space_limit, process_limit } = ctx.session.getLang(ctx.lang);
-    const res = await runQueued(
+    await using res = await runQueued(
         execute.execute,
         {
             stdin: { src: input },
@@ -72,8 +75,6 @@ export async function judge(ctx: Context) {
         if (code < 32 && signalled) message = signals[code];
         else message = { message: 'Your program returned {0}.', params: [code] };
     }
-    await Promise.allSettled(Object.values(res.fileIds).map((id) => del(id)));
-
     if (message) ctx.next({ message });
 
     return ctx.end({

packages/hydrojudge/src/judge/submit_answer.ts

@@ -2,7 +2,7 @@ import { NormalizedCase, STATUS } from '@hydrooj/common';
 import { fs } from '@hydrooj/utils';
 import checkers from '../checkers';
 import { runFlow } from '../flow';
-import { del, runQueued } from '../sandbox';
+import { runQueued } from '../sandbox';
 import { Context } from './interface';
 
 function judgeCase(c: NormalizedCase) {
@@ -15,9 +15,8 @@ function judgeCase(c: NormalizedCase) {
         let status = STATUS.STATUS_ACCEPTED;
         let message: any = '';
         let score = 0;
-        const fileIds = [];
         if (ctx.config.subType === 'multi') {
-            const res = await runQueued(
+            await using res = await runQueued(
                 '/usr/bin/unzip foo.zip',
                 {
                     stdin: null,
@@ -36,7 +35,6 @@ function judgeCase(c: NormalizedCase) {
                 status = STATUS.STATUS_WRONG_ANSWER;
             }
             file = { fileId: res.fileIds[name] };
-            fileIds.push(...Object.values(res.fileIds));
         }
         if (status === STATUS.STATUS_ACCEPTED) {
             ({ status, score, message } = await checkers[ctx.config.checker_type]({
@@ -52,7 +50,6 @@ function judgeCase(c: NormalizedCase) {
                 env: { ...ctx.env, HYDRO_TESTCASE: c.id.toString() },
             }));
         }
-        await Promise.allSettled(fileIds.map(del));
         return {
             id: c.id,
             status,

packages/hydrojudge/src/sandbox.ts

@@ -195,8 +195,9 @@ const queue = new PQueue({ concurrency: getConfig('concurrency') || getConfig('p
 export function runQueued(
     execute: Parameter[], pipeMapping: Pick<PipeMap, 'in' | 'out' | 'name'>[],
     params: Parameter, trace?: string, priority?: number,
-): Promise<SandboxAdaptedResult[]>;
-export function runQueued(execute: string, params: Parameter, trace?: string, priority?: number): Promise<SandboxAdaptedResult>;
+): Promise<SandboxAdaptedResult[] & AsyncDisposable>;
+export function runQueued(execute: string, params: Parameter, trace?: string, priority?: number
+): Promise<SandboxAdaptedResult & AsyncDisposable>;
 export function runQueued(
     arg0: string | Parameter[], arg1: Pick<PipeMap, 'in' | 'out' | 'name'>[] | Parameter,
     arg2?: string | Parameter, arg3?: string | number, arg4?: number,
@@ -207,7 +208,9 @@ export function runQueued(
         : [arg0, arg1, arg2 || {}, arg3 || '', arg4 || 0];
     return queue.add(async () => {
         const res = await runPiped(execute, pipeMapping, params, trace);
-        return single ? res[0] : res;
+        const ret = single ? res[0] : res;
+        (ret as any)[Symbol.asyncDispose] = () => Promise.allSettled(res.flatMap((t) => Object.values(t.fileIds || {}).map(del)));
+        return ret;
     }, { priority });
 }
 
@@ -228,7 +231,7 @@ export async function versionCheck(reportWarn: (str: string) => void, reportErro
     }
     const { osinfo } = await sysinfo.get();
     if (sandboxCgroup === 2) {
-        const kernelVersion = osinfo.kernel.split('-')[0];
+        const kernelVersion = osinfo.kernel.match(/^\d+\.\d+\.\d+/)[0];
         if (!gte(kernelVersion, '5.19.0') || !gte(sandboxVersion, '1.6.10')) {
             reportWarn('You are using cgroup v2 without kernel 5.19+. This could result in inaccurate memory usage measurements.');
         }

packages/hydrojudge/src/sandbox/interface.ts

@@ -7,6 +7,7 @@ export interface SandboxVersion {
     pipeProxy?: boolean;
     addressSpaceLimit?: boolean;
     stream?: boolean;
+    procPeak?: boolean;
 }
 
 export interface LocalFile {
@@ -157,6 +158,8 @@ export interface SandboxResult {
     memory: number;
     /** wall clock time in ns */
     runTime: number;
+    /** peak count of processes (thread) during execution */
+    procPeak?: number;
     /** copyOut file name to content (UTF-8 encoded and invalid character replaced) */
     files?: Record<string, string>;
     /** copyOutCached file name to fileId */

packages/hydrojudge/src/task.ts

@@ -31,6 +31,7 @@ export class JudgeTask {
     lang: string;
     code: CopyInFile;
     input?: string;
+    finished: boolean = false;
     clean: (() => Promise<any>)[] = [];
     data: FileInfo[];
     folder: string;
@@ -93,6 +94,7 @@ export class JudgeTask {
                 });
             }
         } finally {
+            this.finished = true;
             // eslint-disable-next-line no-await-in-loop
             for (const clean of this.clean) await clean()?.catch(() => null);
         }
@@ -174,12 +176,17 @@ export class JudgeTask {
         await judge[type].judge(this);
     }
 
+    async pushClean(f: () => any | Promise<any>) {
+        if (this.finished) await f().catch(() => null);
+        else this.clean.push(f);
+    }
+
     async compile(lang: string, code: CopyInFile) {
         const copyIn = Object.fromEntries(
             (this.config.user_extra_files || []).map((i) => [basename(i), { src: i }]),
         ) as CopyIn;
         const result = await compile(this.session.getLang(lang), code, copyIn, this.next);
-        this.clean.push(result.clean);
+        await this.pushClean(result.clean);
         return result;
     }
 
@@ -218,8 +225,10 @@ export class JudgeTask {
         if (!lang) throw new FormatError(`Unknown ${type} language.`);
         // TODO cache compiled binary
         const result = await compile(lang, { src: file }, copyIn);
-        this.clean.push(result.clean);
-        if (!result._cacheable) return result;
+        if (!result._cacheable) {
+            await this.pushClean(result.clean);
+            return result;
+        }
         await Lock.acquire(this.folder);
         try {
             const loc = join(this.folder, `_${type}.cache`);
@@ -235,6 +244,7 @@ export class JudgeTask {
             const currEtag = await fs.readFile(join(this.folder, 'etags'), 'utf-8');
             await fs.writeFile(join(this.folder, 'etags'), JSON.stringify({ ...JSON.parse(currEtag), '*cache': this.compileCache }));
         } finally {
+            await this.pushClean(result.clean);
             Lock.release(this.folder);
         }
         return result;