judge: add span tracing (#1106)

packages/hydrojudge/package.json

@@ -8,6 +8,9 @@
     "dependencies": {
         "@hydrooj/common": "workspace:^",
         "@hydrooj/utils": "workspace:^",
+        "@opentelemetry/api": "^1.9.0",
+        "@opentelemetry/exporter-trace-otlp-http": "^0.208.0",
+        "@opentelemetry/sdk-node": "^0.208.0",
         "@zip.js/zip.js": "^2.8.11",
         "cac": "^6.7.14",
         "p-queue": "^9.0.1",

packages/hydrojudge/src/config.ts

@@ -32,6 +32,10 @@ export const JudgeSettings = Schema.object({
     host: Schema.any(),
     secret: Schema.string().description('Judge Token Secret').default(randomstring(32)),
     disable: Schema.boolean().description('Disable builtin judge').default(false),
+    tracing: Schema.object({
+        endpoint: Schema.string().role('url').description('Tempo endpoint').default('http://localhost:4318'),
+        samplePercentage: Schema.number().description('Sample percentage').default(0).min(0).max(1),
+    }),
     detail: Schema.union([
         Schema.const('full'),
         Schema.const('case'),

packages/hydrojudge/src/daemon.ts

@@ -20,6 +20,7 @@ import HydroHost from './hosts/hydro';
 import Vj4Host from './hosts/vj4';
 import log from './log';
 import { versionCheck } from './sandbox';
+import { initTracing } from './tracing';
 
 const hosts: Record<string, HydroHost | Vj4Host> = {};
 let exit = false;
@@ -46,6 +47,8 @@ process.on('unhandledRejection', (reason, p) => {
 async function daemon() {
     const shouldRun = await versionCheck((msg) => log.error(msg));
     if (!shouldRun) process.exit(1);
+    const tracing = getConfig('tracing');
+    if (tracing?.endpoint && tracing?.samplePercentage) initTracing(tracing.endpoint, tracing.samplePercentage);
     const _hosts = getConfig('hosts');
     const queue = new PQueue({ concurrency: Infinity });
     await fs.ensureDir(getConfig('tmp_dir'));

packages/hydrojudge/src/flow.ts

@@ -131,6 +131,4 @@ export const runFlow = async (ctx: Context, task: Task) => {
             subtasks: infos,
         });
     }
-    ctx.stat.done = new Date();
-    if (process.env.DEV) ctx.next({ message: JSON.stringify(ctx.stat) });
 };

packages/hydrojudge/src/hosts/builtin.ts

@@ -16,6 +16,7 @@ import { Context } from '../judge/interface';
 import logger from '../log';
 import { versionCheck } from '../sandbox';
 import { JudgeTask } from '../task';
+import { initTracing } from '../tracing';
 
 const session: Session = {
     config: { detail: getConfig('detail') },
@@ -56,6 +57,13 @@ export async function apply(ctx: HydroContext) {
             await versionCheck(warn, error);
         });
     });
+    const tracing = getConfig('tracing');
+    if (tracing?.endpoint && tracing?.samplePercentage) {
+        ctx.effect(() => {
+            const sdk = initTracing(tracing.endpoint, tracing.samplePercentage);
+            return () => sdk.shutdown();
+        });
+    }
     await fs.ensureDir(getConfig('tmp_dir'));
     const info = await sysinfo.get();
     const handle = async (t) => {

packages/hydrojudge/src/judge/generate.ts

@@ -11,7 +11,6 @@ import { JudgeTask } from '../task';
 import { parseMemoryMB, parseTimeMS } from '../utils';
 
 export const judge = async (ctx: JudgeTask) => {
-    ctx.stat.judge = new Date();
     ctx.next({ status: STATUS.STATUS_COMPILING });
     if (!('content' in ctx.code)) throw new SystemError('Unsupported input');
     const [generator, std] = ctx.code.content.toString().split('\n').map((i) => i.trim());
@@ -135,8 +134,6 @@ export const judge = async (ctx: JudgeTask) => {
         const result = await runGenerator(i);
         if (result) await runStd(i, { src: result });
     }
-    ctx.stat.done = new Date();
-    if (process.env.DEV) ctx.next({ message: JSON.stringify(ctx.stat) });
     ctx.end({
         status: totalStatus,
         score: totalStatus === STATUS.STATUS_ACCEPTED ? 100 : 0,

packages/hydrojudge/src/judge/run.ts

@@ -5,7 +5,6 @@ import signals from '../signals';
 import { JudgeTask } from '../task';
 
 export const judge = async (ctx: JudgeTask) => {
-    ctx.stat.judge = new Date();
     ctx.next({ status: STATUS.STATUS_COMPILING });
     const execute = await ctx.compile(ctx.lang, ctx.code);
     ctx.next({ status: STATUS.STATUS_JUDGING, progress: 0 });
@@ -57,8 +56,6 @@ export const judge = async (ctx: JudgeTask) => {
     if ([STATUS.STATUS_WRONG_ANSWER, STATUS.STATUS_RUNTIME_ERROR].includes(status)) {
         await ctx.runAnalysis(execute, { content: ctx.input });
     }
-    ctx.stat.done = new Date();
-    if (process.env.DEV) ctx.next({ message: JSON.stringify(ctx.stat) });
     ctx.end({
         status,
         time: Math.floor(time * 1000000) / 1000000,

packages/hydrojudge/src/task.ts

@@ -1,4 +1,5 @@
 import { basename, join } from 'path';
+import { context, SpanStatusCode, trace } from '@opentelemetry/api';
 import {
     CompilableSource, FileInfo, JudgeMeta, JudgeResultBody, STATUS, TestCase,
 } from '@hydrooj/common';
@@ -25,7 +26,6 @@ const testlibFile = {
 };
 
 export class JudgeTask {
-    stat: Record<string, Date> = Object.create(null);
     source: string;
     rid: string;
     lang: string;
@@ -43,16 +43,23 @@ export class JudgeTask {
     env: Record<string, string>;
     callbackCache?: TestCase[];
     compileCache: Record<string, Pick<Execute, 'execute' | 'copyIn' | typeof Symbol.asyncDispose>> = {};
+    tracer = trace.getTracer('hydrojudge');
+    span = this.tracer.startSpan('judge.task');
+    mainContext = trace.setSpan(context.active(), this.span);
 
     constructor(public session: Session, public request: JudgeRequest) {
-        this.stat.receive = new Date();
         logger.debug('%o', request);
     }
 
+    private startChildSpan(name: string, attributes?: Record<string, any>) {
+        const span = this.tracer.startSpan(name, { attributes }, this.mainContext);
+        span[Symbol.dispose] = () => span.end();
+        return span as typeof span & { [Symbol.dispose]: () => void };
+    }
+
     async handle() {
         const host = this.session.config.host || 'local';
         try {
-            this.stat.handle = new Date();
             this.rid = this.request.rid.toString();
             this.lang = this.request.lang;
             this.code = { content: this.request.code };
@@ -70,6 +77,13 @@ export class JudgeTask {
                 HYDRO_USER: (this.request.uid || 0).toString(),
                 HYDRO_CONTEST: tid,
             };
+            this.span.setAttributes({
+                host,
+                rid: this.rid,
+                source: this.source,
+                lang: this.lang,
+                tid,
+            });
             Object.assign(this, this.session.getReporter(this));
             logger.info('Submission: %s/%s/%s', host, this.source, this.rid);
             await this.doSubmission();
@@ -86,6 +100,7 @@ export class JudgeTask {
                     status: STATUS.STATUS_FORMAT_ERROR, score: 0, time: 0, memory: 0,
                 });
             } else {
+                this.span.recordException(e);
                 logger.error(e);
                 this.next({ message: { message: e.message, params: e.params, ...process.env.DEV ? { stack: e.stack } : {} } });
                 this.end({
@@ -96,12 +111,14 @@ export class JudgeTask {
             this.finished = true;
             // eslint-disable-next-line no-await-in-loop
             for (const clean of this.clean) await clean()?.catch(() => null);
+            this.span.end();
         }
     }
 
     async cacheOpen(source: string, files: FileInfo[]) {
         // Backward compatibility for vj4
         if ((this.session as any).cacheOpen) return (this.session as any).cacheOpen(source, files);
+        using span = this.startChildSpan('judge.cacheOpen', { source, fileCount: files?.length || 0 });
         const host = this.session.config?.host || 'local';
         const filePath = join(getConfig('cache_dir'), host, source);
         await Lock.acquire(filePath);
@@ -125,6 +142,7 @@ export class JudgeTask {
             const allFilesToRemove = Object.keys(etags).filter((name) => !allFiles.has(name) && fs.existsSync(join(filePath, name)));
             await Promise.all(allFilesToRemove.map((name) => fs.remove(join(filePath, name))));
             if (filenames.length) {
+                span.setAttribute('syncedFiles', filenames.length);
                 logger.info(`Getting problem data: ${this.session?.config.host || 'local'}/${source}`);
                 this.next({ message: 'Syncing testdata, please wait...' });
                 await this.session.fetchFile(source, Object.fromEntries(
@@ -139,6 +157,7 @@ export class JudgeTask {
             await fs.writeFile(join(filePath, 'lastUsage'), Date.now().toString());
             return filePath;
         } catch (e) {
+            span.recordException(e);
             logger.warn('CacheOpen Fail: %s %o %o', source, files, e);
             throw e;
         } finally {
@@ -147,14 +166,13 @@ export class JudgeTask {
     }
 
     async doSubmission() {
-        this.stat.cache_start = new Date();
         this.folder = await this.cacheOpen(this.source, this.data);
         if (this.files?.code) {
             const target = await this.session.fetchFile(null, { [this.files.code]: '' });
             this.code = { src: target };
             this.clean.push(() => fs.remove(target));
         }
-        this.stat.read_cases = new Date();
+        using readCasesSpan = this.startChildSpan('judge.readCases', { folder: this.folder, lang: this.lang });
         this.config = await readCases(
             this.folder,
             {
@@ -171,12 +189,14 @@ export class JudgeTask {
                     ? null : this.session.getLang(this.lang),
             },
         );
-        this.stat.judge = new Date();
+        readCasesSpan.end();
         const type = this.request.contest?.toString() === '000000000000000000000000' ? 'run'
             : this.request.type === 'generate' ? 'generate'
                 : this.files?.hack ? 'hack'
                     : this.config.type || 'default';
         if (!judge[type]) throw new FormatError('Unrecognized problemType: {0}', [type]);
+        using executeSpan = this.startChildSpan('judge.execute', { judgeType: type });
+        this.mainContext = trace.setSpan(context.active(), executeSpan);
         await judge[type].judge(this);
     }
 
@@ -186,10 +206,12 @@ export class JudgeTask {
     }
 
     async compile(lang: string, code: CopyInFile) {
+        using compileSpan = this.startChildSpan('judge.compile', { lang });
         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);
+        compileSpan.setStatus({ code: SpanStatusCode.OK });
         await this.pushClean(result.clean);
         return result;
     }
@@ -213,6 +235,7 @@ export class JudgeTask {
                 clean: () => Promise.resolve(null),
             };
         }
+        using span = this.startChildSpan('judge.compileLocalFile', { type, checkerType: checkerType || '', cached: false });
         const withTestlib = type !== 'std' && (type !== 'checker' || checkerType === 'testlib');
         const extra = type === 'std' ? this.config.user_extra_files : this.config.judge_extra_files;
         const copyIn = {
@@ -234,12 +257,14 @@ export class JudgeTask {
             }
         } else lang = this.session.getLang(langId, false);
         if (!lang) throw new FormatError(`Unknown ${type} language.`);
+        span.setAttribute('lang', langId);
         // TODO cache compiled binary
         const result = await compile(lang, { src: file }, copyIn);
         if (!result._cacheable) {
             await this.pushClean(result.clean);
             return result;
         }
+        span.setAttribute('cacheable', true);
         await Lock.acquire(this.folder);
         try {
             const loc = join(this.folder, `_${type}.cache`);
@@ -265,6 +290,7 @@ export class JudgeTask {
     async runAnalysis(execute: Execute, input: CopyInFile) {
         const langConfig = this.session.getLang(this.lang);
         if (!langConfig.analysis) return;
+        using span = this.startChildSpan('judge.runAnalysis', { lang: this.lang });
         try {
             const r = await runQueued(langConfig.analysis, {
                 copyIn: {
@@ -282,6 +308,7 @@ export class JudgeTask {
             if (out.length) this.next({ compilerText: out.substring(0, 1024) });
             if (process.env.DEV) console.log(r);
         } catch (e) {
+            span.recordException(e);
             logger.info('Failed to run analysis');
             logger.error(e);
         }

packages/hydrojudge/src/tracing.ts

@@ -0,0 +1,48 @@
+import { OTLPTraceExporter as OTLPTraceExporterHttp } from '@opentelemetry/exporter-trace-otlp-http';
+import { NodeSDK, resources, tracing } from '@opentelemetry/sdk-node';
+import { Logger } from './log';
+
+const logger = new Logger('tracing');
+
+export function initTracing(endpoint: string, samplePercentage = 1.0) {
+    const traceExporter = new OTLPTraceExporterHttp({
+        url: endpoint.includes('/v1/traces') ? endpoint : `${endpoint}/v1/traces`,
+    });
+    const sdk = new NodeSDK({
+        resource: resources.resourceFromAttributes({
+            'service.name': 'hydrojudge',
+            'service.version': require('../package.json').version,
+        }),
+        traceExporter,
+        sampler: new tracing.TraceIdRatioBasedSampler(samplePercentage),
+        instrumentations: [],
+    });
+    sdk.start();
+    const originalExport = traceExporter.export;
+    if (originalExport && process.env.HYDROJUDGE_DEBUG_TRACING) {
+        // eslint-disable-next-line consistent-return
+        traceExporter.export = function wrappedExport(spans: any, resultCallback: any) {
+            const spanCount = spans?.length || 0;
+            logger.info('[Tracing] Exporting %d span(s) to Tempo...', spanCount);
+            const wrappedCallback = (result: any) => {
+                if (result.code === 0) {
+                    logger.info('[Tracing] Successfully exported %d span(s) - result code: %d', spanCount, result.code);
+                } else {
+                    logger.error('[Tracing] Export failed: %o', {
+                        code: result.code,
+                        error: result.error,
+                        details: result,
+                    });
+                }
+                resultCallback?.(result);
+            };
+            try {
+                return originalExport.call(this, spans, wrappedCallback);
+            } catch (error) {
+                logger.error('[Tracing] Export exception: %o', error);
+                resultCallback?.({ code: 1, error });
+            }
+        };
+    }
+    return sdk;
+}