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
	iocp     windows.Handle
	peak     int64
	exceeded bool
}

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

const (
	jobObjectExtendedLimitInformationClass           = 9
	jobObjectAssociateCompletionPortInformationClass = 7

	jobObjectLimitProcessMemory  = 0x00000100
	jobObjectLimitKillOnJobClose = 0x00002000

	jobObjectMsgProcessMemoryLimit = 9
	jobObjectMsgJobMemoryLimit     = 10
)

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
}

type jobObjectAssociateCompletionPort struct {
	CompletionKey  uintptr
	CompletionPort windows.Handle
}

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(extended): %w", err)
	}

	iocp, err := windows.CreateIoCompletionPort(windows.InvalidHandle, 0, 0, 1)
	if err != nil {
		windows.CloseHandle(job)
		return fmt.Errorf("CreateIoCompletionPort: %w", err)
	}
	assoc := jobObjectAssociateCompletionPort{
		CompletionKey:  uintptr(job),
		CompletionPort: iocp,
	}
	if _, err := windows.SetInformationJobObject(
		job,
		jobObjectAssociateCompletionPortInformationClass,
		uintptr(unsafe.Pointer(&assoc)),
		uint32(unsafe.Sizeof(assoc)),
	); err != nil {
		windows.CloseHandle(iocp)
		windows.CloseHandle(job)
		return fmt.Errorf("SetInformationJobObject(iocp): %w", err)
	}

	l.job = job
	l.iocp = iocp
	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{}
	}

	if l.iocp != 0 {
		for {
			var bytes uint32
			var key uintptr
			var overlapped *windows.Overlapped
			err := windows.GetQueuedCompletionStatus(l.iocp, &bytes, &key, &overlapped, 0)
			if err != nil {
				break
			}
			if bytes == jobObjectMsgProcessMemoryLimit || bytes == jobObjectMsgJobMemoryLimit {
				l.exceeded = true
			}
		}
	}

	var info jobObjectExtendedLimitInformation
	var ret uint32
	if err := windows.QueryInformationJobObject(
		l.job,
		jobObjectExtendedLimitInformationClass,
		uintptr(unsafe.Pointer(&info)),
		uint32(unsafe.Sizeof(info)),
		&ret,
	); err == nil {
		l.peak = int64(info.PeakProcessMemoryUsed)
	}

	return limitStats{PeakMemory: l.peak, MemoryExceeded: l.exceeded}
}

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

func cleanupCgroupRoot() {}