135 lines
3.1 KiB
Go
135 lines
3.1 KiB
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
|
|
}
|
|
}
|