judge/runner/limiter_windows.go

//go:build windows

package runner

import (
	"fmt"
	"os/exec"
	"unsafe"

	"golang.org/x/sys/windows"
)

type windowsLimiter struct {
	memLimit int64
	job      windows.Handle
	peak     int64
	exceeded bool
}

func newLimiter(memLimit int64) limiter {
	return &windowsLimiter{memLimit: memLimit}
}

const (
	jobObjectExtendedLimitInformationClass = 9
	jobObjectLimitProcessMemory            = 0x00000100
	jobObjectLimitKillOnJobClose           = 0x00002000
)

type ioCounters struct {
	ReadOperationCount  uint64
	WriteOperationCount uint64
	OtherOperationCount uint64
	ReadTransferCount   uint64
	WriteTransferCount  uint64
	OtherTransferCount  uint64
}

type jobObjectBasicLimitInformation struct {
	PerProcessUserTimeLimit int64
	PerJobUserTimeLimit     int64
	LimitFlags              uint32
	MinimumWorkingSetSize   uintptr
	MaximumWorkingSetSize   uintptr
	ActiveProcessLimit      uint32
	Affinity                uintptr
	PriorityClass           uint32
	SchedulingClass         uint32
}

type jobObjectExtendedLimitInformation struct {
	BasicLimitInformation jobObjectBasicLimitInformation
	IoInfo                ioCounters
	ProcessMemoryLimit    uintptr
	JobMemoryLimit        uintptr
	PeakProcessMemoryUsed uintptr
	PeakJobMemoryUsed     uintptr
}

func (l *windowsLimiter) prepare(cmd *exec.Cmd) error {
	if l.memLimit <= 0 {
		return nil
	}
	job, err := windows.CreateJobObject(nil, nil)
	if err != nil {
		return fmt.Errorf("CreateJobObject: %w", err)
	}
	info := jobObjectExtendedLimitInformation{
		BasicLimitInformation: jobObjectBasicLimitInformation{
			LimitFlags: jobObjectLimitProcessMemory | jobObjectLimitKillOnJobClose,
		},
		ProcessMemoryLimit: uintptr(l.memLimit),
	}
	if _, err := windows.SetInformationJobObject(
		job,
		jobObjectExtendedLimitInformationClass,
		uintptr(unsafe.Pointer(&info)),
		uint32(unsafe.Sizeof(info)),
	); err != nil {
		windows.CloseHandle(job)
		return fmt.Errorf("SetInformationJobObject: %w", err)
	}
	l.job = job
	return nil
}

func (l *windowsLimiter) afterStart(cmd *exec.Cmd) error {
	if l.job == 0 || cmd.Process == nil {
		return nil
	}
	procHandle, err := windows.OpenProcess(
		windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE|windows.PROCESS_QUERY_INFORMATION,
		false,
		uint32(cmd.Process.Pid),
	)
	if err != nil {
		return fmt.Errorf("OpenProcess: %w", err)
	}
	defer windows.CloseHandle(procHandle)
	if err := windows.AssignProcessToJobObject(l.job, procHandle); err != nil {
		return fmt.Errorf("AssignProcessToJobObject: %w", err)
	}
	return nil
}

func (l *windowsLimiter) collect() limitStats {
	if l.job == 0 {
		return limitStats{}
	}
	var info jobObjectExtendedLimitInformation
	var ret uint32
	err := windows.QueryInformationJobObject(
		l.job,
		jobObjectExtendedLimitInformationClass,
		uintptr(unsafe.Pointer(&info)),
		uint32(unsafe.Sizeof(info)),
		&ret,
	)
	if err != nil {
		return limitStats{}
	}
	l.peak = int64(info.PeakProcessMemoryUsed)
	if l.memLimit > 0 && l.peak >= l.memLimit {
		l.exceeded = true
	}
	return limitStats{PeakMemory: l.peak, MemoryExceeded: l.exceeded}
}

func (l *windowsLimiter) cleanup() {
	if l.job != 0 {
		windows.CloseHandle(l.job)
		l.job = 0
	}
}

func cleanupCgroupRoot() {}