judge/runner/limiter_linux.go

//go:build linux

package runner

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var (
	cgroupRootOnce  sync.Once
	cgroupRoot      string
	cgroupRootOwned bool
	cgroupInitErr   error
	cgroupCounter   int64
)

const cgroupFSRoot = "/sys/fs/cgroup"

func ensureCgroupRoot() (string, error) {
	cgroupRootOnce.Do(func() {
		if root, err := createOwnedCgroup(); err == nil {
			cgroupRoot = root
			cgroupRootOwned = true
			return
		}

		root, err := createScopeCgroup()
		if err != nil {
			cgroupInitErr = err
			return
		}
		cgroupRoot = root
	})
	return cgroupRoot, cgroupInitErr
}

func hasController(path, name string) (bool, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return false, err
	}
	needle := " " + name + " "
	return strings.Contains(" "+strings.TrimSpace(string(data))+" ", needle), nil
}

func createOwnedCgroup() (string, error) {
	ok, err := hasController(filepath.Join(cgroupFSRoot, "cgroup.controllers"), "memory")
	if err != nil {
		return "", err
	}
	if !ok {
		return "", fmt.Errorf("memory controller not available in root cgroup")
	}

	enabled, err := hasController(filepath.Join(cgroupFSRoot, "cgroup.subtree_control"), "memory")
	if err != nil {
		return "", err
	}
	if !enabled {
		if err := os.WriteFile(filepath.Join(cgroupFSRoot, "cgroup.subtree_control"), []byte("+memory"), 0644); err != nil {
			return "", fmt.Errorf("enable +memory in root subtree_control: %w", err)
		}
	}

	name := fmt.Sprintf("judge.%d", os.Getpid())
	root := filepath.Join(cgroupFSRoot, name)
	if err := os.Mkdir(root, 0755); err != nil {
		if os.IsExist(err) {
		} else {
			return "", err
		}
	}

	ok, err = hasController(filepath.Join(root, "cgroup.controllers"), "memory")
	if err != nil {
		_ = os.Remove(root)
		return "", err
	}
	if !ok {
		_ = os.Remove(root)
		return "", fmt.Errorf("memory controller not inherited into %s", root)
	}
	if err := os.WriteFile(filepath.Join(root, "cgroup.subtree_control"), []byte("+memory"), 0644); err != nil {
		_ = os.Remove(root)
		return "", fmt.Errorf("enable +memory in %s: %w", root, err)
	}

	return root, nil
}

func createScopeCgroup() (string, error) {
	data, err := os.ReadFile("/proc/self/cgroup")
	if err != nil {
		return "", fmt.Errorf("read /proc/self/cgroup: %w", err)
	}
	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 == "" {
		return "", fmt.Errorf("cgroup v2 not found in /proc/self/cgroup (unified hierarchy required)")
	}
	ownCg := filepath.Join(cgroupFSRoot, rel)

	ok, err := hasController(filepath.Join(ownCg, "cgroup.controllers"), "memory")
	if err != nil {
		return "", fmt.Errorf("cgroup %s not accessible: %w", ownCg, err)
	}
	if !ok {
		return "", fmt.Errorf("memory controller not delegated to %s", ownCg)
	}

	initCg := filepath.Join(ownCg, "judge.init")
	if err := os.MkdirAll(initCg, 0755); err != nil {
		return "", fmt.Errorf("mkdir %s: %w", initCg, err)
	}
	if err := os.WriteFile(filepath.Join(initCg, "cgroup.procs"), []byte(strconv.Itoa(os.Getpid())), 0644); err != nil {
		return "", fmt.Errorf("move judge into %s: %w", initCg, err)
	}

	if err := os.WriteFile(filepath.Join(ownCg, "cgroup.subtree_control"), []byte("+memory"), 0644); err != nil {
		enabled, _ := hasController(filepath.Join(ownCg, "cgroup.subtree_control"), "memory")
		if !enabled {
			return "", fmt.Errorf("enable +memory in %s/cgroup.subtree_control: %w", ownCg, err)
		}
	}
	return ownCg, nil
}

func cleanupCgroupRoot() {
	if cgroupRoot == "" || !cgroupRootOwned {
		return
	}
	_ = os.Remove(cgroupRoot)
	cgroupRoot = ""
	cgroupRootOwned = false
}

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)
	}
	swapMax := filepath.Join(l.cgPath, "memory.swap.max")
	if _, err := os.Stat(swapMax); err == nil {
		_ = os.WriteFile(swapMax, []byte("0"), 0644)
	}
	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)
	}
}