1

refactor: error handling, single rss call, stale PID cleanup, utf16 trim, O_APPEND log

This commit is contained in:
2026-06-19 23:20:10 +03:00
parent d0cafe1332
commit b999b7ec4f
3 changed files with 202 additions and 151 deletions
+1 -1
View File
@@ -1,3 +1,3 @@
module autopriority module autopriority
go 1.24 go 1.25.0
View File
+201 -150
View File
@@ -12,156 +12,194 @@ import (
"unsafe" "unsafe"
) )
const (
CreateToolhelp32SnapshotProcess = 0x00000002
ProcessSetInformation = 0x00000200
ProcessQueryInformation = 0x00000400
PriorityClassIdle = 0x00000040
PriorityClassNormal = 0x00000020
PriorityClassHigh = 0x00000080
MessageBoxIconWarning = 0x00000030
ErrorAlreadyExists = 1832
)
type processEntry32 struct {
Size uint32
CntUsage uint32
PID uint32
DefaultHeapID uintptr
ModuleID 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 ( var (
k32 = syscall.NewLazyDLL("kernel32.dll") k32 = syscall.NewLazyDLL("kernel32.dll")
ps = syscall.NewLazyDLL("psapi.dll") ps = syscall.NewLazyDLL("psapi.dll")
pCreateSnap = k32.NewProc("CreateToolhelp32Snapshot") procCreateSnap = k32.NewProc("CreateToolhelp32Snapshot")
pProcess32First = k32.NewProc("Process32FirstW") procProcess32First = k32.NewProc("Process32FirstW")
pProcess32Next = k32.NewProc("Process32NextW") procProcess32Next = k32.NewProc("Process32NextW")
pGetProcMemInfo = ps.NewProc("GetProcessMemoryInfo") procGetMemInfo = ps.NewProc("GetProcessMemoryInfo")
pSetPriority = k32.NewProc("SetPriorityClass") procSetPriority = k32.NewProc("SetPriorityClass")
pGetPriority = k32.NewProc("GetPriorityClass") procGetPriority = k32.NewProc("GetPriorityClass")
pOpenProcess = k32.NewProc("OpenProcess") procOpenProcess = k32.NewProc("OpenProcess")
pCloseHandle = k32.NewProc("CloseHandle") procCloseHandle = k32.NewProc("CloseHandle")
pCreateMutex = k32.NewProc("CreateMutexW") procCreateMutex = k32.NewProc("CreateMutexW")
pMessageBox = k32.NewProc("MessageBoxW") procMessageBox = k32.NewProc("MessageBoxW")
pGetLastError = k32.NewProc("GetLastError") procGetLastError = k32.NewProc("GetLastError")
) )
const (
SNAPPROCESS = 0x00000002
SET_INFO = 0x0200
QUERY_INFO = 0x0400
IDLE_PRIO = 0x00000040
NORMAL_PRIO = 0x00000020
HIGH_PRIO = 0x00000080
EXISTING = 1832
)
type pe32 struct {
Size uint32
CntUsage uint32
PID uint32
HeapID uintptr
ModuleID uint32
Threads uint32
ParentPID uint32
PrioClass int32
Flags uint32
ExeFile [260]uint16
}
type pmc struct {
CBM uint32
PageFaults uint32
PeakWS uintptr
WorkingSet uintptr
QuotaPeakPG uintptr
QuotaPG uintptr
QuotaPeakPM uintptr
QuotaPM uintptr
PeakPF uintptr
Pagefile uintptr
Reserved [2]uintptr
}
type procInfo struct {
PID uint32
Name string
}
func wide(s string) *uint16 {
return syscall.StringToUTF16Ptr(s)
}
func closeH(h syscall.Handle) { func closeH(h syscall.Handle) {
_, _, _ = pCloseHandle.Call(uintptr(h)) _, _, _ = procCloseHandle.Call(uintptr(h))
} }
func openProc(pid uint32, acc uint32) syscall.Handle { func openProc(pid uint32, acc uint32) (syscall.Handle, error) {
r, _, _ := pOpenProcess.Call(uintptr(acc), 0, uintptr(pid)) r, _, e := procOpenProcess.Call(uintptr(acc), 0, uintptr(pid))
return syscall.Handle(r)
}
func setPrio(pid uint32, cls uint32) bool {
h := openProc(pid, SET_INFO)
if h == 0 {
return false
}
defer closeH(h)
r, _, _ := pSetPriority.Call(uintptr(h), uintptr(cls))
return r != 0
}
func curPrio(pid uint32) uint32 {
h := openProc(pid, QUERY_INFO)
if h == 0 {
return 0
}
defer closeH(h)
r, _, _ := pGetPriority.Call(uintptr(h))
return uint32(r)
}
func rss(pid uint32) uint64 {
h := openProc(pid, QUERY_INFO)
if h == 0 {
return 0
}
defer closeH(h)
var m pmc
r, _, _ := pGetProcMemInfo.Call(uintptr(h), uintptr(unsafe.Pointer(&m)), uintptr(unsafe.Sizeof(m)))
if r == 0 { if r == 0 {
return 0 return 0, e
} }
return uint64(m.WorkingSet) return syscall.Handle(r), nil
} }
func allProcs() []procInfo { func setPrio(pid uint32, cls uint32) error {
var out []procInfo h, err := openProc(pid, ProcessSetInformation)
pe := pe32{Size: uint32(unsafe.Sizeof(pe32{}))} if err != nil {
snap, _, _ := pCreateSnap.Call(SNAPPROCESS, 0) return err
}
defer closeH(h)
r, _, e := procSetPriority.Call(uintptr(h), uintptr(cls))
if r == 0 {
return e
}
return nil
}
func curPrio(pid uint32) (uint32, error) {
h, err := openProc(pid, ProcessQueryInformation)
if err != nil {
return 0, err
}
defer closeH(h)
r, _, e := procGetPriority.Call(uintptr(h))
if r == 0 {
return 0, e
}
return uint32(r), nil
}
func rss(pid uint32) (uint64, error) {
h, err := openProc(pid, ProcessQueryInformation)
if err != nil {
return 0, err
}
defer closeH(h)
var m processMemoryCounters
r, _, e := procGetMemInfo.Call(uintptr(h), uintptr(unsafe.Pointer(&m)), uintptr(unsafe.Sizeof(m)))
if r == 0 {
return 0, e
}
return uint64(m.WorkingSetSize), 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 { if snap == 0 {
return out return nil, e
} }
defer closeH(syscall.Handle(snap)) defer closeH(syscall.Handle(snap))
r, _, _ := pProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe))) pe := processEntry32{Size: uint32(unsafe.Sizeof(processEntry32{}))}
if r != 0 { r, _, e := procProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe)))
for { if r == 0 {
out = append(out, procInfo{ return nil, e
PID: pe.PID, }
Name: syscall.UTF16ToString(pe.ExeFile[:]),
}) var out []processInfo
r, _, _ := pProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe))) for {
if r == 0 { name := utf16Trim(pe.ExeFile[:])
break if name != "" {
} out = append(out, processInfo{PID: pe.PID, Name: name})
}
r, _, e = procProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe)))
if r == 0 {
break
} }
} }
return out return out, nil
} }
func prioName(c uint32) string { func prioName(c uint32) string {
switch c { switch c {
case IDLE_PRIO: case PriorityClassIdle:
return "IDLE" return "IDLE"
case NORMAL_PRIO: case PriorityClassNormal:
return "NORMAL" return "NORMAL"
case HIGH_PRIO: case PriorityClassHigh:
return "HIGH" return "HIGH"
default: default:
return fmt.Sprintf("0x%X", c) return fmt.Sprintf("0x%X", c)
} }
} }
func isAlive(pid uint32) bool {
h, _ := openProc(pid, ProcessQueryInformation)
if h == 0 {
return false
}
closeH(h)
return true
}
func ensureSingleInstance() {
h, _, _ := procCreateMutex.Call(0, 0, uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("Global\\autopriority_v1"))))
le, _, _ := procGetLastError.Call()
if uint32(le) == ErrorAlreadyExists {
closeH(syscall.Handle(h))
procMessageBox.Call(0,
uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("autoPriority is already running."))),
uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("autoPriority"))),
MessageBoxIconWarning,
)
os.Exit(1)
}
}
func main() { func main() {
mem := flag.Uint64("mem", 512*1024*1024, "memory threshold in bytes") mem := flag.Uint64("mem", 512*1024*1024, "memory threshold in bytes")
interval := flag.Duration("interval", time.Minute, "scan interval") interval := flag.Duration("interval", time.Minute, "scan interval")
dryRun := flag.Bool("dry-run", false, "log only, do not change priorities") dryRun := flag.Bool("dry-run", false, "log only, do not change priorities")
doLog := flag.Bool("log", true, "write log file to temp directory")
flag.Parse() flag.Parse()
if *interval < 10*time.Second { if *interval < 10*time.Second {
@@ -169,17 +207,12 @@ func main() {
} }
var logFile *os.File var logFile *os.File
if *doLog { logDir := os.TempDir()
logDir := os.TempDir() if logDir != "" {
if logDir == "" {
exePath, _ := os.Executable()
logDir = filepath.Dir(exePath)
}
f, err := os.OpenFile(filepath.Join(logDir, "autopriority.log"), f, err := os.OpenFile(filepath.Join(logDir, "autopriority.log"),
os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
if err == nil { if err == nil {
logFile = f logFile = f
defer logFile.Close()
} }
} }
@@ -197,31 +230,23 @@ func main() {
logf("autoPriority started (mem=%d bytes, interval=%s, dry-run=%v)", *mem, *interval, *dryRun) logf("autoPriority started (mem=%d bytes, interval=%s, dry-run=%v)", *mem, *interval, *dryRun)
// Single instance ensureSingleInstance()
mh, _, _ := pCreateMutex.Call(0, 0, uintptr(unsafe.Pointer(wide("Global\\autopriority_v1"))))
lastErr, _, _ := pGetLastError.Call() if err := setPrio(uint32(os.Getpid()), PriorityClassIdle); err != nil {
if lastErr == uintptr(EXISTING) { logf("warning: could not set own priority to IDLE: %v", err)
closeH(syscall.Handle(mh)) } else {
pMessageBox.Call(0, logf("own priority set to IDLE")
uintptr(unsafe.Pointer(wide("autoPriority is already running."))),
uintptr(unsafe.Pointer(wide("autoPriority"))),
0x30,
)
os.Exit(1)
} }
// Lower own priority
pSetPriority.Call(uintptr(syscall.Handle(^uintptr(0))), IDLE_PRIO)
logf("own priority set to IDLE")
// promoted: PID → name (for restore logging)
promoted := make(map[uint32]string) promoted := make(map[uint32]string)
// Restore on exit if logFile != nil {
defer logFile.Close()
}
defer func() { defer func() {
for pid, name := range promoted { for pid, name := range promoted {
if setPrio(pid, NORMAL_PRIO) { if err := setPrio(pid, PriorityClassNormal); err == nil {
logf("RESTORE %s (PID %d) NORMAL", name, pid) logf("RESTORE %s (PID %d) -> NORMAL", name, pid)
} }
} }
logf("autoPriority stopped") logf("autoPriority stopped")
@@ -232,30 +257,56 @@ func main() {
defer ticker.Stop() defer ticker.Stop()
scan := func() { scan := func() {
for _, p := range allProcs() { procs, err := allProcs()
if err != nil {
logf("allProcs error: %v", err)
return
}
for pid := range promoted {
if !isAlive(pid) {
delete(promoted, pid)
}
}
for _, p := range procs {
if p.PID == myPID || p.PID == 0 { if p.PID == myPID || p.PID == 0 {
continue continue
} }
if _, ok := promoted[p.PID]; ok { if _, ok := promoted[p.PID]; ok {
continue continue
} }
if rss(p.PID) < *mem {
rssBytes, err := rss(p.PID)
if err != nil {
logf("rss(%d) error: %v", p.PID, err)
continue
}
if rssBytes < *mem {
continue continue
} }
orig := curPrio(p.PID) rssMB := rssBytes / 1024 / 1024
origPrio, prioErr := curPrio(p.PID)
origStr := prioName(origPrio)
if prioErr != nil {
origStr = "unknown"
}
if *dryRun { if *dryRun {
logf("[DRY-RUN] %s (PID %d) RSS=%dMB, priority=%s — would set HIGH", logf("[DRY-RUN] %s (PID %d) RSS=%dMB, priority=%s — would set HIGH",
p.Name, p.PID, rss(p.PID)/1024/1024, prioName(orig)) p.Name, p.PID, rssMB, origStr)
continue continue
} }
if setPrio(p.PID, HIGH_PRIO) { if err := setPrio(p.PID, PriorityClassHigh); err != nil {
promoted[p.PID] = p.Name logf("setPrio(%d, HIGH) error: %v", p.PID, err)
logf("PROMOTE %s (PID %d) RSS=%dMB, %s → HIGH", continue
p.Name, p.PID, rss(p.PID)/1024/1024, prioName(orig))
} }
promoted[p.PID] = p.Name
logf("PROMOTE %s (PID %d) RSS=%dMB, %s -> HIGH",
p.Name, p.PID, rssMB, origStr)
} }
} }