add memory limit

runner/limiter_linux.go

@@ -0,0 +1,153 @@
+//go:build linux
+
+package runner
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	cgroupRootOnce sync.Once
+	cgroupRoot     string
+	cgroupInitErr  error
+	cgroupCounter  int64
+)
+
+const cgroupFSRoot = "/sys/fs/cgroup"
+
+func ensureCgroupRoot() (string, error) {
+	cgroupRootOnce.Do(func() {
+		data, err := os.ReadFile("/proc/self/cgroup")
+		if err != nil {
+			cgroupInitErr = fmt.Errorf("read /proc/self/cgroup: %w", err)
+			return
+		}
+		var rel string
+		for _, line := range strings.Split(strings.TrimSpace(string(data)), "\n") {
+			if strings.HasPrefix(line, "0::") {
+				rel = strings.TrimPrefix(line, "0::")
+				break
+			}
+		}
+		if rel == "" {
+			cgroupInitErr = fmt.Errorf("cgroup v2 not found in /proc/self/cgroup (unified hierarchy required)")
+			return
+		}
+		ownCg := filepath.Join(cgroupFSRoot, rel)
+
+		controllers, err := os.ReadFile(filepath.Join(ownCg, "cgroup.controllers"))
+		if err != nil {
+			cgroupInitErr = fmt.Errorf("cgroup %s not accessible: %w", ownCg, err)
+			return
+		}
+		if !strings.Contains(" "+string(controllers)+" ", " memory ") {
+			cgroupInitErr = fmt.Errorf("memory controller not delegated to %s (controllers: %s)", ownCg, strings.TrimSpace(string(controllers)))
+			return
+		}
+
+		initCg := filepath.Join(ownCg, "judge.init")
+		if err := os.MkdirAll(initCg, 0755); err != nil {
+			cgroupInitErr = fmt.Errorf("mkdir %s: %w", initCg, err)
+			return
+		}
+		if err := os.WriteFile(filepath.Join(initCg, "cgroup.procs"), []byte(strconv.Itoa(os.Getpid())), 0644); err != nil {
+			cgroupInitErr = fmt.Errorf("move judge into %s: %w", initCg, err)
+			return
+		}
+
+		if err := os.WriteFile(filepath.Join(ownCg, "cgroup.subtree_control"), []byte("+memory"), 0644); err != nil {
+			current, _ := os.ReadFile(filepath.Join(ownCg, "cgroup.subtree_control"))
+			if !strings.Contains(" "+string(current)+" ", " memory ") {
+				cgroupInitErr = fmt.Errorf("enable +memory in %s/cgroup.subtree_control: %w", ownCg, err)
+				return
+			}
+		}
+
+		cgroupRoot = ownCg
+	})
+	return cgroupRoot, cgroupInitErr
+}
+
+type linuxLimiter struct {
+	memLimit int64
+	cgPath   string
+}
+
+func newLimiter(memLimit int64) limiter {
+	return &linuxLimiter{memLimit: memLimit}
+}
+
+func (l *linuxLimiter) prepare(cmd *exec.Cmd) error {
+	if l.memLimit <= 0 {
+		return nil
+	}
+	root, err := ensureCgroupRoot()
+	if err != nil {
+		return err
+	}
+	name := fmt.Sprintf("judge.test.%d.%d", os.Getpid(), atomic.AddInt64(&cgroupCounter, 1))
+	l.cgPath = filepath.Join(root, name)
+	if err := os.Mkdir(l.cgPath, 0755); err != nil {
+		l.cgPath = ""
+		return fmt.Errorf("mkdir %s: %w", name, err)
+	}
+	if err := os.WriteFile(filepath.Join(l.cgPath, "memory.max"), []byte(strconv.FormatInt(l.memLimit, 10)), 0644); err != nil {
+		_ = os.Remove(l.cgPath)
+		l.cgPath = ""
+		return fmt.Errorf("write memory.max: %w", err)
+	}
+	return nil
+}
+
+func (l *linuxLimiter) afterStart(cmd *exec.Cmd) error {
+	if l.cgPath == "" || cmd.Process == nil {
+		return nil
+	}
+	return os.WriteFile(filepath.Join(l.cgPath, "cgroup.procs"), []byte(strconv.Itoa(cmd.Process.Pid)), 0644)
+}
+
+func (l *linuxLimiter) collect() limitStats {
+	if l.cgPath == "" {
+		return limitStats{}
+	}
+	var s limitStats
+	if data, err := os.ReadFile(filepath.Join(l.cgPath, "memory.peak")); err == nil {
+		if n, err := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64); err == nil {
+			s.PeakMemory = n
+		}
+	}
+	if data, err := os.ReadFile(filepath.Join(l.cgPath, "memory.events")); err == nil {
+		for _, line := range strings.Split(string(data), "\n") {
+			fields := strings.Fields(line)
+			if len(fields) != 2 {
+				continue
+			}
+			if (fields[0] == "oom_kill" || fields[0] == "oom_group_kill") && fields[1] != "0" {
+				s.MemoryExceeded = true
+			}
+		}
+	}
+	return s
+}
+
+func (l *linuxLimiter) cleanup() {
+	if l.cgPath == "" {
+		return
+	}
+	for i := 0; i < 10; i++ {
+		err := os.Remove(l.cgPath)
+		if err == nil || os.IsNotExist(err) {
+			l.cgPath = ""
+			return
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+}