This commit is contained in:
xkm
2025-04-27 22:04:12 +08:00
parent 326a2a007c
commit 7c02d88f3a
4 changed files with 91 additions and 116 deletions

View File

@@ -4,91 +4,45 @@ import (
"context"
"encoding/json"
"fmt"
"log"
"os"
"path/filepath"
"time"
"runbin/internal/config"
"runbin/internal/model"
"github.com/docker/docker/api/types/container"
"github.com/docker/docker/client"
)
func monitorMemory(ctx context.Context, cli *client.Client, containerID string, maxMem *uint64) {
// 创建独立上下文用于内存监控
memCtx, cancel := context.WithCancel(context.Background())
defer cancel()
ticker := time.NewTicker(64 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-memCtx.Done():
return
case <-ticker.C:
// 检查容器是否仍在运行
_, err := cli.ContainerInspect(memCtx, containerID)
if err != nil {
return
}
statsResp, err := cli.ContainerStats(memCtx, containerID, false)
if err != nil {
continue // 忽略临时错误
}
var statsJSON container.StatsResponse
if err := json.NewDecoder(statsResp.Body).Decode(&statsJSON); err != nil {
statsResp.Body.Close()
continue
}
statsResp.Body.Close()
if statsJSON.MemoryStats.Stats != nil {
if currentMem, ok := statsJSON.MemoryStats.Stats["anon"]; ok {
log.Println(currentMem)
if currentMem > *maxMem {
*maxMem = currentMem
}
}
}
}
}
type Usage struct {
ExitStatus int64 `json:"exit_status"`
MaxMemory int64 `json:"max_memory"`
RealTime float64 `json:"real_time"`
}
func (w *Worker) RunCppTask(ctx context.Context, task *model.Paste, cli *client.Client) error {
// Create temporary workspace
tmpDir, err := os.MkdirTemp("", "cpp_compile_")
if err != nil {
return fmt.Errorf("create temp dir error: %v", err)
func compileCpp(ctx context.Context, task *model.Paste, cli *client.Client, tmpDir string, cfg *config.WorkerConfig) error {
hostConfig := &container.HostConfig{
Binds: []string{tmpDir + ":/app"},
Resources: container.Resources{
Memory: int64(cfg.Limit.Memory * 1024 * 1024),
CPUQuota: int64(cfg.Limit.Cpu * 100000),
},
NetworkMode: "none",
}
defer os.RemoveAll(tmpDir)
// Write user code to main.cpp
// 写入main.cpp
codePath := filepath.Join(tmpDir, "main.cpp")
if err := os.WriteFile(codePath, []byte(task.Code), 0644); err != nil {
return fmt.Errorf("write code file error: %v", err)
}
// Get execution limits from config
timeout := time.Duration(w.cfg.Limit.Time) * time.Second
compliteCtx, cancel := context.WithTimeout(ctx, timeout)
compliteCtx, cancel := context.WithTimeout(ctx, time.Duration(cfg.Limit.Time)*time.Second)
defer cancel()
hostConfig := &container.HostConfig{
Binds: []string{tmpDir + ":/app"},
Resources: container.Resources{
Memory: int64(w.cfg.Limit.Memory * 1024 * 1024),
CPUQuota: int64(w.cfg.Limit.Cpu * 100000),
},
NetworkMode: "none",
}
// Create compile container configuration
// 创建容器
resp, err := cli.ContainerCreate(compliteCtx, &container.Config{
Image: "gcc:14",
Image: cfg.CompilerImage,
Cmd: []string{"sh", "-c", "g++ -std=c++20 /app/main.cpp -o /app/output > /app/compile.txt 2>&1"},
}, hostConfig, nil, nil, filepath.Base(tmpDir)+"_builder")
if err != nil {
@@ -98,21 +52,15 @@ func (w *Worker) RunCppTask(ctx context.Context, task *model.Paste, cli *client.
Force: true,
})
// Start compile container execution
// 启动容器
if err := cli.ContainerStart(compliteCtx, resp.ID, container.StartOptions{}); err != nil {
return fmt.Errorf("failed to start compile container: %v", err)
}
// Wait for compile container completion
// 等待
statusCh, errCh := cli.ContainerWait(compliteCtx, resp.ID, container.WaitConditionNotRunning)
// Handle compile container execution results
select {
case <-statusCh:
// Get compile container detailed state
containerState, err := cli.ContainerInspect(compliteCtx, resp.ID)
if err != nil {
return fmt.Errorf("container state inspection error: %v", err)
}
case status := <-statusCh:
// Read compilation log
if logData, err := os.ReadFile(filepath.Join(tmpDir, "compile.txt")); err == nil {
@@ -120,10 +68,10 @@ func (w *Worker) RunCppTask(ctx context.Context, task *model.Paste, cli *client.
}
// Non-zero exit code indicates compilation failure
if containerState.State.ExitCode != 0 {
if status.StatusCode != 0 {
task.Status = model.StatusCompileError
return nil
}
return nil
case err := <-errCh:
return err
case <-compliteCtx.Done():
@@ -131,20 +79,31 @@ func (w *Worker) RunCppTask(ctx context.Context, task *model.Paste, cli *client.
task.CompileLog = "Compile process exceeded time limit"
return nil
}
}
// Write input to input.txt
func runCpp(ctx context.Context, task *model.Paste, cli *client.Client, tmpDir string, cfg *config.WorkerConfig) error {
hostConfig := &container.HostConfig{
Binds: []string{tmpDir + ":/app"},
Resources: container.Resources{
Memory: int64(cfg.Limit.Memory * 1024 * 1024),
CPUQuota: int64(cfg.Limit.Cpu * 100000),
},
NetworkMode: "none",
}
// 写入 input.txt
inputPath := filepath.Join(tmpDir, "input.txt")
if err := os.WriteFile(inputPath, []byte(task.Stdin), 0644); err != nil {
return fmt.Errorf("write input file error: %v", err)
}
runCtx, cancel := context.WithTimeout(ctx, timeout)
runCtx, cancel := context.WithTimeout(ctx, time.Duration(cfg.Limit.Time)*time.Second)
defer cancel()
// Create runner container configuration
resp, err = cli.ContainerCreate(runCtx, &container.Config{
Image: "gcc:14",
Cmd: []string{"sh", "-c", "sh -c \"/app/output < /app/input.txt > /app/stdout.txt\" > /app/stderr.txt 2>&1"},
resp, err := cli.ContainerCreate(runCtx, &container.Config{
Image: cfg.CompilerImage,
Cmd: []string{"sh", "-c", `/usr/bin/time --format='{"exit_status":%x,"max_memory":%M,"real_time":%e}' -o /app/usage.json /app/output < /app/input.txt > /app/stdout.txt 2> /app/stderr.txt`},
}, hostConfig, nil, nil, filepath.Base(tmpDir)+"_runner")
if err != nil {
return fmt.Errorf("create runner container error: %v", err)
@@ -158,34 +117,17 @@ func (w *Worker) RunCppTask(ctx context.Context, task *model.Paste, cli *client.
return fmt.Errorf("failed to start runner container: %v", err)
}
// 启动内存监控协程(使用独立上下文)
var maxMem uint64
go monitorMemory(context.Background(), cli, resp.ID, &maxMem)
// 等待容器完成
statusCh, errCh = cli.ContainerWait(runCtx, resp.ID, container.WaitConditionNotRunning)
statusCh, errCh := cli.ContainerWait(runCtx, resp.ID, container.WaitConditionNotRunning)
// 处理执行结果
select {
case <-statusCh:
// Get container detailed state
containerState, err := cli.ContainerInspect(runCtx, resp.ID)
if err != nil {
return fmt.Errorf("container state inspection error: %v", err)
}
startTime, startErr := time.Parse(time.RFC3339Nano, containerState.State.StartedAt)
finishTime, finishErr := time.Parse(time.RFC3339Nano, containerState.State.FinishedAt)
if startErr == nil && finishErr == nil && !startTime.IsZero() && !finishTime.IsZero() {
timeUsage := finishTime.Sub(startTime)
task.ExecutionTimeMs = int(timeUsage.Milliseconds())
}
case status := <-statusCh:
// 非零退出码表示运行时错误
if containerState.State.OOMKilled {
if status.StatusCode == 137 {
task.Status = model.StatusMemoryLimitExceed
} else if containerState.State.ExitCode != 0 {
} else if status.StatusCode != 0 {
task.Status = model.StatusRuntimeError
} else {
task.Status = model.StatusCompleted
@@ -199,6 +141,7 @@ func (w *Worker) RunCppTask(ctx context.Context, task *model.Paste, cli *client.
// Process execution results by reading output files
stdoutPath := filepath.Join(tmpDir, "stdout.txt")
stderrPath := filepath.Join(tmpDir, "stderr.txt")
usagePath := filepath.Join(tmpDir, "usage.json")
// Read program output
if outData, err := os.ReadFile(stdoutPath); err == nil {
@@ -207,18 +150,31 @@ func (w *Worker) RunCppTask(ctx context.Context, task *model.Paste, cli *client.
if outData, err := os.ReadFile(stderrPath); err == nil {
task.Stderr = string(outData)
}
// 确保内存统计有效性至少1MB
if maxMem > 0 {
task.MemoryUsageKb = int(maxMem / 1024)
} else {
// 如果统计失败尝试从OOM状态获取
if task.Status == model.StatusMemoryLimitExceed {
task.MemoryUsageKb = int(w.cfg.Limit.Memory * 1024) // 使用配置的内存限制值
} else {
task.MemoryUsageKb = -1 // 表示统计不可用
}
if usageData, err := os.ReadFile(usagePath); err == nil {
var usage Usage
fmt.Println(string(usageData))
fmt.Println(json.Unmarshal(usageData, &usage))
task.MemoryUsageKb = int(usage.MaxMemory)
task.ExecutionTimeMs = int(usage.RealTime * 1000)
}
return nil
}
func (w *Worker) RunCppTask(ctx context.Context, task *model.Paste, cli *client.Client) error {
// 临时文件夹
tmpDir, err := os.MkdirTemp("", "cpp_compile_")
if err != nil {
return fmt.Errorf("create temp dir error: %v", err)
}
defer os.RemoveAll(tmpDir)
if err := compileCpp(ctx, task, cli, tmpDir, w.cfg); err != nil {
return err
}
if task.Status == model.StatusCompileError {
return nil
}
return runCpp(ctx, task, cli, tmpDir, w.cfg)
}