runbin/internal/worker/RunCppTask.go

package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"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
					}
				}
			}
		}
	}
}

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)
	}
	defer os.RemoveAll(tmpDir)

	// Write user code to 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)
	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",
		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 {
		return fmt.Errorf("create compile container error: %v", err)
	}
	defer cli.ContainerRemove(ctx, resp.ID, container.RemoveOptions{
		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)
		}

		// Read compilation log
		if logData, err := os.ReadFile(filepath.Join(tmpDir, "compile.txt")); err == nil {
			task.CompileLog = string(logData)
		}

		// Non-zero exit code indicates compilation failure
		if containerState.State.ExitCode != 0 {
			task.Status = model.StatusCompileError
			return nil
		}
	case err := <-errCh:
		return err
	case <-compliteCtx.Done():
		task.Status = model.StatusCompileError
		task.CompileLog = "Compile process exceeded time limit"
		return nil
	}

	// Write input to 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)
	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"},
	}, hostConfig, nil, nil, filepath.Base(tmpDir)+"_runner")
	if err != nil {
		return fmt.Errorf("create runner container error: %v", err)
	}
	defer cli.ContainerRemove(ctx, resp.ID, container.RemoveOptions{
		Force: true,
	})

	// Start runner container execution
	if err := cli.ContainerStart(runCtx, resp.ID, container.StartOptions{}); err != nil {
		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)

	// 处理执行结果
	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())
		}

		// 非零退出码表示运行时错误
		if containerState.State.OOMKilled {
			task.Status = model.StatusMemoryLimitExceed
		} else if containerState.State.ExitCode != 0 {
			task.Status = model.StatusRuntimeError
		} else {
			task.Status = model.StatusCompleted
		}
	case err := <-errCh:
		return err
	case <-runCtx.Done():
		task.Status = model.StatusTimeLimitExceed
	}

	// Process execution results by reading output files
	stdoutPath := filepath.Join(tmpDir, "stdout.txt")
	stderrPath := filepath.Join(tmpDir, "stderr.txt")

	// Read program output
	if outData, err := os.ReadFile(stdoutPath); err == nil {
		task.Stdout = string(outData)
	}
	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 // 表示统计不可用
		}
	}

	return nil
}