Przeglądaj źródła

# feat: Docker沙箱改用stats API采集真实内存用量(bytes→KB,超时降级为0)

yangyi 2 tygodni temu
rodzic
commit
87cd32b931

+ 63 - 7
src/main/java/space/anyi/codesandbox/DockerCodeSandBox.java

@@ -7,6 +7,7 @@ import com.github.dockerjava.api.command.ExecCreateCmdResponse;
 import com.github.dockerjava.api.model.Capability;
 import com.github.dockerjava.api.model.Frame;
 import com.github.dockerjava.api.model.HostConfig;
+import com.github.dockerjava.api.model.Statistics;
 import com.github.dockerjava.api.model.StreamType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -18,6 +19,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.*;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * @fileName: DockerCodeSandBox
@@ -44,6 +46,9 @@ public class DockerCodeSandBox implements CodeSandBox {
     /** 获取并发许可的超时(毫秒),超时返回SE */
     private static final int CONCURRENCY_ACQUIRE_TIMEOUT_MS = 5_000;
 
+    /** 采集容器内存统计的超时(毫秒),超时/异常时降级为 0 */
+    private static final long MEMORY_STATS_TIMEOUT_MS = 3_000;
+
     /** 并发控制信号量 */
     private static final Semaphore CONCURRENCY_LIMIT = new Semaphore(MAX_CONCURRENT);
 
@@ -110,7 +115,7 @@ public class DockerCodeSandBox implements CodeSandBox {
                 TestCase tc = codeContext.testCases().get(i);
                 RunResult run = runCase(containerId, codeContext, tc);
                 String verdict = finalVerdict(run, tc);
-                results.add(new TestCaseResponse(tc.id(), verdict, run.durationMs(), 0L));
+                results.add(new TestCaseResponse(tc.id(), verdict, run.durationMs(), run.memoryUsedKb()));
                 log.info("用例 {}/{} 判定: {} (耗时 {}ms, 退出码 {})",
                         i + 1, codeContext.testCases().size(), verdict, run.durationMs(), run.exitCode());
                 if (!Submission.STATUS_AC.equals(verdict)) {
@@ -279,8 +284,10 @@ public class DockerCodeSandBox implements CodeSandBox {
      * @param exitCode 进程退出码,超时/异常为 -1
      * @param durationMs 运行耗时(毫秒)
      * @param verdict 用例判定,为 null 表示进程正常退出,待输出比对后定 AC/WA
+     * @param memoryUsedKb 用例运行期间容器内存占用(KB),采集失败为 0
      */
-    public record RunResult(String stdout, String stderr, int exitCode, long durationMs, String verdict) {
+    public record RunResult(String stdout, String stderr, int exitCode, long durationMs, String verdict,
+                            long memoryUsedKb) {
     }
 
     private RunResult runCase(String containerId, CodeContext codeContext, TestCase tc) {
@@ -328,7 +335,10 @@ public class DockerCodeSandBox implements CodeSandBox {
                     });
 
             boolean completed = callback.awaitCompletion(codeContext.maxTime(), TimeUnit.MILLISECONDS);
+            long memoryUsedKb;
             if (!completed) {
+                // 超时:先采样内存,再强杀容器
+                memoryUsedKb = getMemoryUsageKb(containerId);
                 try {
                     dockerClient.killContainerCmd(containerId).withSignal("KILL").exec();
                 } catch (Exception e) {
@@ -336,6 +346,9 @@ public class DockerCodeSandBox implements CodeSandBox {
                 }
                 callback.close();
                 timedOut = true;
+            } else {
+                // 正常退出后采样内存
+                memoryUsedKb = getMemoryUsageKb(containerId);
             }
 
             long durationMs = (System.nanoTime() - startNanos) / 1_000_000;
@@ -343,20 +356,61 @@ public class DockerCodeSandBox implements CodeSandBox {
             String err = stderr.toString(StandardCharsets.UTF_8);
 
             if (timedOut) {
-                return new RunResult(out, err, -1, durationMs, Submission.STATUS_TLE);
+                return new RunResult(out, err, -1, durationMs, Submission.STATUS_TLE, memoryUsedKb);
             }
             Long exitCodeLong = dockerClient.inspectExecCmd(exec.getId()).exec().getExitCodeLong();
             int exitCode = exitCodeLong == null ? -1 : exitCodeLong.intValue();
             if (truncated[0] || exitCode != 0) {
-                return new RunResult(out, err, exitCode, durationMs, Submission.STATUS_RE);
+                return new RunResult(out, err, exitCode, durationMs, Submission.STATUS_RE, memoryUsedKb);
             }
-            return new RunResult(out, err, exitCode, durationMs, null);
+            return new RunResult(out, err, exitCode, durationMs, null, memoryUsedKb);
         } catch (Exception e) {
             log.error("用例运行系统异常", e);
+            long memoryUsedKb = getMemoryUsageKb(containerId);
             long durationMs = (System.nanoTime() - startNanos) / 1_000_000;
             return new RunResult(stdout.toString(StandardCharsets.UTF_8),
-                    stderr.toString(StandardCharsets.UTF_8), -1, durationMs, Submission.STATUS_SE);
+                    stderr.toString(StandardCharsets.UTF_8), -1, durationMs, Submission.STATUS_SE, memoryUsedKb);
+        }
+    }
+
+    /**
+     * 采集容器当前内存占用(no-stream 模式,仅取一帧)
+     * @param containerId 容器 ID
+     * @return 内存占用(KB),超时或异常时降级为 0
+     */
+    private long getMemoryUsageKb(String containerId) {
+        AtomicReference<Long> usageBytes = new AtomicReference<>();
+        ResultCallback.Adapter<Statistics> callback = null;
+        try {
+            callback = dockerClient.statsCmd(containerId)
+                    .withNoStream(true)
+                    .exec(new ResultCallback.Adapter<Statistics>() {
+                        @Override
+                        public void onNext(Statistics stats) {
+                            if (stats != null && stats.getMemoryStats() != null
+                                    && stats.getMemoryStats().getUsage() != null) {
+                                usageBytes.set(stats.getMemoryStats().getUsage());
+                            }
+                        }
+                    });
+            callback.awaitCompletion(MEMORY_STATS_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            return 0L;
+        } catch (Exception e) {
+            log.warn("采集容器内存失败: {}", e.getMessage());
+            return 0L;
+        } finally {
+            if (callback != null) {
+                try {
+                    callback.close();
+                } catch (Exception e) {
+                    log.warn("关闭内存采集回调失败: {}", e.getMessage());
+                }
+            }
         }
+        Long usage = usageBytes.get();
+        return usage == null ? 0L : usage / 1024L;
     }
 
     private String finalVerdict(RunResult run, TestCase tc) {
@@ -417,6 +471,8 @@ public class DockerCodeSandBox implements CodeSandBox {
         if (results == null || results.isEmpty()) {
             return 0L;
         }
-        return results.stream().map(TestCaseResponse::usedMemory).min(Long::compareTo).orElse(0L);
+        return results.stream().map(TestCaseResponse::usedMemory)
+                .filter(memory -> memory != null && memory > 0)
+                .min(Long::compareTo).orElse(0L);
     }
 }