add memory limit

runner/limiter_windows.go

@@ -0,0 +1,134 @@
+//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
+	}
+}