feat(judge): add memory usage monitoring and handling

This commit is contained in:
cfngc4594 2025-03-02 01:30:38 +08:00
parent 9dede32073
commit efdf5089fa

View File

@ -81,7 +81,35 @@ async function compileCode(container: Docker.Container, filePath: string, fileNa
}); });
} }
async function runCode(container: Docker.Container, fileName: string, timeout: number) { async function monitorMemoryUsage(container: Docker.Container, memoryLimit: number, timeoutHandle: NodeJS.Timeout) {
return new Promise<void>((resolve, reject) => {
const interval = setInterval(async () => {
try {
const stats = await container.stats({ stream: false });
const memoryUsage = stats.memory_stats.usage / (1024 * 1024); // Convert to MB
console.log(`Memory usage: ${memoryUsage.toFixed(2)} MB`);
if (memoryUsage > memoryLimit) {
console.warn(`Memory limit exceeded: ${memoryUsage.toFixed(2)} MB > ${memoryLimit} MB`);
clearInterval(interval);
clearTimeout(timeoutHandle); // Clear the timeout timer
await container.stop({ t: 0 });
await container.remove();
reject(new Error("Memory limit exceeded, container stopped and removed"));
}
} catch (error) {
console.error("Error monitoring memory:", error);
clearInterval(interval);
reject(error);
}
}, 500); // Check every 500ms
resolve();
});
}
async function runCode(container: Docker.Container, fileName: string, timeout: number, memoryLimit: number) {
const runExec = await container.exec({ const runExec = await container.exec({
Cmd: [`./${fileName}`], Cmd: [`./${fileName}`],
AttachStdout: true, AttachStdout: true,
@ -90,6 +118,7 @@ async function runCode(container: Docker.Container, fileName: string, timeout: n
return new Promise<string>((resolve, reject) => { return new Promise<string>((resolve, reject) => {
let timeoutHandle: NodeJS.Timeout; let timeoutHandle: NodeJS.Timeout;
let memoryMonitor: Promise<void>;
runExec.start({}, async (error, stream) => { runExec.start({}, async (error, stream) => {
if (error) { if (error) {
@ -118,23 +147,29 @@ async function runCode(container: Docker.Container, fileName: string, timeout: n
docker.modem.demuxStream(stream, stdoutStream, stderrStream); docker.modem.demuxStream(stream, stdoutStream, stderrStream);
// Timeout monitoring
timeoutHandle = setTimeout(async () => { timeoutHandle = setTimeout(async () => {
try {
console.warn("Execution timed out, stopping container..."); console.warn("Execution timed out, stopping container...");
try {
await container.stop({ t: 0 }); await container.stop({ t: 0 });
console.warn("Container stopped, removing...");
await container.remove(); await container.remove();
reject(new Error("Execution timed out and container was removed")); reject(new Error("Execution timed out and container was removed"));
} catch (stopError) { } catch (stopError) {
console.error('Error stopping/removing container:', stopError);
reject(new Error("Execution timed out, but failed to stop/remove container")); reject(new Error("Execution timed out, but failed to stop/remove container"));
} }
}, timeout); }, timeout); // Use the configured timeout
// Memory monitoring
memoryMonitor = monitorMemoryUsage(container, memoryLimit, timeoutHandle);
stream.on("end", () => { stream.on("end", () => {
clearTimeout(timeoutHandle); clearTimeout(timeoutHandle);
memoryMonitor.then(() => {
const stdout = stdoutChunks.join(""); const stdout = stdoutChunks.join("");
const stderr = stderrChunks.join(""); const stderr = stderrChunks.join("");
resolve(stderr ? stdout + stderr : stdout); resolve(stderr ? stdout + stderr : stdout);
}).catch(reject);
}); });
stream.on("error", (error) => { stream.on("error", (error) => {
@ -157,7 +192,7 @@ async function cleanupContainer(container: Docker.Container) {
} }
export async function judge(language: string, value: string) { export async function judge(language: string, value: string) {
const { image, tag, fileName, extension, workingDir, timeout } = LanguageConfigs[language]; const { image, tag, fileName, extension, workingDir, timeout, memoryLimit } = LanguageConfigs[language];
const filePath = `${fileName}.${extension}`; const filePath = `${fileName}.${extension}`;
let container: Docker.Container | undefined; let container: Docker.Container | undefined;
@ -173,7 +208,7 @@ export async function judge(language: string, value: string) {
return compileOutput; return compileOutput;
} }
return await runCode(container, fileName, timeout); return await runCode(container, fileName, timeout, memoryLimit);
} catch (error) { } catch (error) {
console.error("Error during judging:", error); console.error("Error during judging:", error);
throw error; throw error;