426 lines
10 KiB
Go
426 lines
10 KiB
Go
//go:build windows
|
|
|
|
package main
|
|
|
|
import (
|
|
"errors"
|
|
"flag"
|
|
"fmt"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"syscall"
|
|
"time"
|
|
"unsafe"
|
|
)
|
|
|
|
const (
|
|
CreateToolhelp32SnapshotProcess = 0x00000002
|
|
ProcessSetInformation = 0x00000200
|
|
ProcessQueryInformation = 0x00000400
|
|
PriorityClassIdle = 0x00000040
|
|
PriorityClassNormal = 0x00000020
|
|
PriorityClassHigh = 0x00000080
|
|
PriorityClassAboveNormal = 0x00008000
|
|
PriorityClassRealtime = 0x00000100
|
|
PriorityClassBelowNormal = 0x00004000
|
|
MessageBoxIconWarning = 0x00000030
|
|
ErrorAlreadyExists = 183
|
|
ErrorNoMoreFiles = 18
|
|
ErrorAccessDenied = 5
|
|
ErrorInvalidParameter = 87
|
|
)
|
|
|
|
type processEntry32 struct {
|
|
Size uint32
|
|
CntUsage uint32
|
|
PID uint32
|
|
DefaultHeapID uintptr
|
|
ModuleID uint32
|
|
CntThreads uint32
|
|
ParentPID uint32
|
|
PrioClass int32
|
|
Flags uint32
|
|
ExeFile [260]uint16
|
|
}
|
|
|
|
type processMemoryCounters struct {
|
|
CBM uint32
|
|
PageFaultCount uint32
|
|
PeakWorkingSetSize uintptr
|
|
WorkingSetSize uintptr
|
|
QuotaPeakPagedPoolUsage uintptr
|
|
QuotaPagedPoolUsage uintptr
|
|
QuotaPeakNonPagedPoolUsage uintptr
|
|
QuotaNonPagedPoolUsage uintptr
|
|
PeakPagefileUsage uintptr
|
|
PagefileUsage uintptr
|
|
PrivateUsage uintptr
|
|
}
|
|
|
|
type processInfo struct {
|
|
PID uint32
|
|
Name string
|
|
}
|
|
|
|
var (
|
|
k32 = syscall.NewLazyDLL("kernel32.dll")
|
|
ps = syscall.NewLazyDLL("psapi.dll")
|
|
|
|
procCreateSnap = k32.NewProc("CreateToolhelp32Snapshot")
|
|
procProcess32First = k32.NewProc("Process32FirstW")
|
|
procProcess32Next = k32.NewProc("Process32NextW")
|
|
procGetMemInfo = ps.NewProc("GetProcessMemoryInfo")
|
|
procSetPriority = k32.NewProc("SetPriorityClass")
|
|
procGetPriority = k32.NewProc("GetPriorityClass")
|
|
procOpenProcess = k32.NewProc("OpenProcess")
|
|
procCloseHandle = k32.NewProc("CloseHandle")
|
|
procCreateMutex = k32.NewProc("CreateMutexW")
|
|
procMessageBox = k32.NewProc("MessageBoxW")
|
|
)
|
|
|
|
func closeH(h syscall.Handle) {
|
|
_, _, _ = procCloseHandle.Call(uintptr(h))
|
|
}
|
|
|
|
func openProc(pid uint32, acc uint32) (syscall.Handle, error) {
|
|
r, _, e := procOpenProcess.Call(uintptr(acc), 0, uintptr(pid))
|
|
if r == 0 {
|
|
return 0, fmt.Errorf("OpenProcess(%d) failed: %w", pid, e)
|
|
}
|
|
return syscall.Handle(r), nil
|
|
}
|
|
|
|
func setPrio(pid uint32, cls uint32) error {
|
|
h, err := openProc(pid, ProcessSetInformation)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer closeH(h)
|
|
r, _, e := procSetPriority.Call(uintptr(h), uintptr(cls))
|
|
if r == 0 {
|
|
return fmt.Errorf("SetPriorityClass(%d) failed: %w", pid, e)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ignoreOpenError(err error) bool {
|
|
return errors.Is(err, syscall.Errno(ErrorAccessDenied)) ||
|
|
errors.Is(err, syscall.Errno(ErrorInvalidParameter))
|
|
}
|
|
|
|
func curPrioWithHandle(h syscall.Handle) (uint32, error) {
|
|
r, _, e := procGetPriority.Call(uintptr(h))
|
|
if r == 0 {
|
|
return 0, fmt.Errorf("GetPriorityClass failed: %w", e)
|
|
}
|
|
return uint32(r), nil
|
|
}
|
|
|
|
func utf16Trim(p []uint16) string {
|
|
for i, v := range p {
|
|
if v == 0 {
|
|
return syscall.UTF16ToString(p[:i])
|
|
}
|
|
}
|
|
return syscall.UTF16ToString(p)
|
|
}
|
|
|
|
func allProcs() ([]processInfo, error) {
|
|
snap, _, e := procCreateSnap.Call(CreateToolhelp32SnapshotProcess, 0)
|
|
if snap == 0 {
|
|
return nil, fmt.Errorf("CreateToolhelp32Snapshot failed: %w", e)
|
|
}
|
|
defer closeH(syscall.Handle(snap))
|
|
|
|
pe := processEntry32{Size: uint32(unsafe.Sizeof(processEntry32{}))}
|
|
r, _, e := procProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe)))
|
|
if r == 0 {
|
|
return nil, fmt.Errorf("Process32First failed: %w", e)
|
|
}
|
|
|
|
var out []processInfo
|
|
for {
|
|
name := utf16Trim(pe.ExeFile[:])
|
|
if name != "" {
|
|
out = append(out, processInfo{PID: pe.PID, Name: name})
|
|
}
|
|
pe.Size = uint32(unsafe.Sizeof(processEntry32{}))
|
|
r, _, e = procProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe)))
|
|
if r == 0 {
|
|
if e == syscall.Errno(ErrorNoMoreFiles) {
|
|
break
|
|
}
|
|
return nil, fmt.Errorf("Process32Next failed: %w", e)
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func isAboveNormal(c uint32) bool {
|
|
return c == PriorityClassAboveNormal ||
|
|
c == PriorityClassHigh ||
|
|
c == PriorityClassRealtime
|
|
}
|
|
|
|
func prioName(c uint32) string {
|
|
switch c {
|
|
case PriorityClassIdle:
|
|
return "IDLE"
|
|
case PriorityClassBelowNormal:
|
|
return "BELOW_NORMAL"
|
|
case PriorityClassNormal:
|
|
return "NORMAL"
|
|
case PriorityClassAboveNormal:
|
|
return "ABOVE_NORMAL"
|
|
case PriorityClassHigh:
|
|
return "HIGH"
|
|
case PriorityClassRealtime:
|
|
return "REALTIME"
|
|
default:
|
|
return fmt.Sprintf("0x%X", c)
|
|
}
|
|
}
|
|
|
|
func procName(pid uint32) (string, error) {
|
|
snap, _, e := procCreateSnap.Call(CreateToolhelp32SnapshotProcess, 0)
|
|
if snap == 0 {
|
|
return "", fmt.Errorf("CreateToolhelp32Snapshot failed: %w", e)
|
|
}
|
|
defer closeH(syscall.Handle(snap))
|
|
|
|
pe := processEntry32{Size: uint32(unsafe.Sizeof(processEntry32{}))}
|
|
r, _, e := procProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe)))
|
|
if r == 0 {
|
|
return "", fmt.Errorf("Process32First failed: %w", e)
|
|
}
|
|
|
|
for {
|
|
if pe.PID == pid {
|
|
return utf16Trim(pe.ExeFile[:]), nil
|
|
}
|
|
pe.Size = uint32(unsafe.Sizeof(processEntry32{}))
|
|
r, _, e = procProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe)))
|
|
if r == 0 {
|
|
if e == syscall.Errno(ErrorNoMoreFiles) {
|
|
break
|
|
}
|
|
return "", fmt.Errorf("Process32Next failed: %w", e)
|
|
}
|
|
}
|
|
return "", nil
|
|
}
|
|
|
|
func isAlive(pid uint32, expectedName string) bool {
|
|
h, err := openProc(pid, ProcessQueryInformation)
|
|
if err != nil || h == 0 {
|
|
return false
|
|
}
|
|
closeH(h)
|
|
name, err := procName(pid)
|
|
if err != nil || name == "" {
|
|
return false
|
|
}
|
|
return name == expectedName
|
|
}
|
|
|
|
func ensureSingleInstance() syscall.Handle {
|
|
h, _, err := procCreateMutex.Call(0, 0, uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("Global\\autopriority_v1"))))
|
|
if h == 0 {
|
|
fmt.Fprintf(os.Stderr, "autoPriority: failed to create mutex: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
if err == syscall.Errno(ErrorAlreadyExists) {
|
|
procMessageBox.Call(0,
|
|
uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("autoPriority is already running."))),
|
|
uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("autoPriority"))),
|
|
MessageBoxIconWarning,
|
|
)
|
|
os.Exit(1)
|
|
}
|
|
return syscall.Handle(h)
|
|
}
|
|
|
|
func main() {
|
|
mem := flag.Uint64("mem", 512*1024*1024, "memory threshold in bytes")
|
|
interval := flag.Duration("interval", time.Minute, "scan interval")
|
|
dryRun := flag.Bool("dry-run", false, "log only, do not change priorities")
|
|
flag.Parse()
|
|
|
|
if *interval < 10*time.Second {
|
|
*interval = 10 * time.Second
|
|
}
|
|
|
|
var logFile *os.File
|
|
logPath := filepath.Join(os.TempDir(), "autopriority.log")
|
|
os.Remove(logPath)
|
|
f, err := os.OpenFile(logPath,
|
|
os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0644)
|
|
if err == nil {
|
|
logFile = f
|
|
defer logFile.Close()
|
|
}
|
|
|
|
ts := func() string {
|
|
return time.Now().Format("02.01.2006 15:04:05")
|
|
}
|
|
|
|
logf := func(format string, a ...any) {
|
|
if logFile == nil {
|
|
return
|
|
}
|
|
fmt.Fprintf(logFile, "[%s] ", ts())
|
|
fmt.Fprintf(logFile, format+"\n", a...)
|
|
logFile.Sync()
|
|
}
|
|
|
|
logf("autoPriority started (mem=%d bytes, interval=%s, dry-run=%v)", *mem, *interval, *dryRun)
|
|
|
|
mutex := ensureSingleInstance()
|
|
defer closeH(mutex)
|
|
|
|
if err := setPrio(uint32(os.Getpid()), PriorityClassIdle); err != nil {
|
|
logf("warning: could not set own priority to IDLE: %v", err)
|
|
} else {
|
|
logf("own priority set to IDLE")
|
|
}
|
|
|
|
promoted := make(map[uint32]string)
|
|
blocked := make(map[uint32]string)
|
|
|
|
defer func() {
|
|
for pid, name := range promoted {
|
|
cur, err := procName(pid)
|
|
if err != nil || cur != name {
|
|
logf("SKIP RESTORE PID %d: process name mismatch (expected %s, got %s)", pid, name, cur)
|
|
continue
|
|
}
|
|
if err := setPrio(pid, PriorityClassNormal); err != nil {
|
|
logf("RESTORE %s (PID %d) -> NORMAL error: %v", name, pid, err)
|
|
} else {
|
|
logf("RESTORE %s (PID %d) -> NORMAL", name, pid)
|
|
}
|
|
}
|
|
logf("autoPriority stopped")
|
|
}()
|
|
|
|
myPID := uint32(os.Getpid())
|
|
ticker := time.NewTicker(*interval)
|
|
defer ticker.Stop()
|
|
|
|
stop := make(chan os.Signal, 1)
|
|
signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
|
|
|
|
scan := func() {
|
|
procs, err := allProcs()
|
|
if err != nil {
|
|
logf("allProcs error: %v", err)
|
|
return
|
|
}
|
|
|
|
for pid, name := range promoted {
|
|
if !isAlive(pid, name) {
|
|
delete(promoted, pid)
|
|
}
|
|
}
|
|
|
|
for pid, name := range blocked {
|
|
if !isAlive(pid, name) {
|
|
delete(blocked, pid)
|
|
}
|
|
}
|
|
|
|
for _, p := range procs {
|
|
if p.PID == myPID || p.PID == 0 {
|
|
continue
|
|
}
|
|
if _, ok := promoted[p.PID]; ok {
|
|
continue
|
|
}
|
|
if _, ok := blocked[p.PID]; ok {
|
|
continue
|
|
}
|
|
|
|
h, err := openProc(p.PID, ProcessQueryInformation|ProcessSetInformation)
|
|
if err != nil {
|
|
if !ignoreOpenError(err) {
|
|
logf("openProc(%s, PID %d) error: %v", p.Name, p.PID, err)
|
|
}
|
|
continue
|
|
}
|
|
|
|
var m processMemoryCounters
|
|
m.CBM = uint32(unsafe.Sizeof(m))
|
|
r, _, e := procGetMemInfo.Call(uintptr(h), uintptr(unsafe.Pointer(&m)), uintptr(unsafe.Sizeof(m)))
|
|
if r == 0 {
|
|
closeH(h)
|
|
logf("GetProcessMemoryInfo(%d) error: %v", p.PID, e)
|
|
continue
|
|
}
|
|
rssBytes := uint64(m.WorkingSetSize)
|
|
if rssBytes < *mem {
|
|
origPrio, prioErr := curPrioWithHandle(h)
|
|
if prioErr == nil && isAboveNormal(origPrio) {
|
|
if *dryRun {
|
|
logf("[DRY-RUN] %s (PID %d) RSS=%dMB, priority=%s — would set NORMAL",
|
|
p.Name, p.PID, rssBytes/1024/1024, prioName(origPrio))
|
|
} else {
|
|
sr, _, se := procSetPriority.Call(uintptr(h), uintptr(PriorityClassNormal))
|
|
if sr != 0 {
|
|
logf("DEMOTE %s (PID %d) RSS=%dMB, %s -> NORMAL",
|
|
p.Name, p.PID, rssBytes/1024/1024, prioName(origPrio))
|
|
} else {
|
|
blocked[p.PID] = p.Name
|
|
logf("BLOCK %s (PID %d): %v (added to exclusion list)", p.Name, p.PID, se)
|
|
}
|
|
}
|
|
}
|
|
closeH(h)
|
|
continue
|
|
}
|
|
|
|
rssMB := rssBytes / 1024 / 1024
|
|
origPrio, prioErr := curPrioWithHandle(h)
|
|
if prioErr != nil {
|
|
closeH(h)
|
|
continue
|
|
}
|
|
if origPrio == PriorityClassHigh {
|
|
closeH(h)
|
|
continue
|
|
}
|
|
origStr := prioName(origPrio)
|
|
|
|
if *dryRun {
|
|
closeH(h)
|
|
logf("[DRY-RUN] %s (PID %d) RSS=%dMB, priority=%s — would set HIGH",
|
|
p.Name, p.PID, rssMB, origStr)
|
|
continue
|
|
}
|
|
|
|
sr, _, se := procSetPriority.Call(uintptr(h), uintptr(PriorityClassHigh))
|
|
closeH(h)
|
|
if sr == 0 {
|
|
blocked[p.PID] = p.Name
|
|
logf("BLOCK %s (PID %d): %v (added to exclusion list)", p.Name, p.PID, se)
|
|
continue
|
|
}
|
|
|
|
promoted[p.PID] = p.Name
|
|
logf("PROMOTE %s (PID %d) RSS=%dMB, %s -> HIGH",
|
|
p.Name, p.PID, rssMB, origStr)
|
|
}
|
|
}
|
|
|
|
scan()
|
|
for {
|
|
select {
|
|
case <-stop:
|
|
logf("received shutdown signal")
|
|
return
|
|
case <-ticker.C:
|
|
scan()
|
|
}
|
|
}
|
|
}
|