diff --git a/go.mod b/go.mod index 7b79382..e938cb5 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,3 @@ module autopriority -go 1.24 +go 1.25.0 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..e69de29 diff --git a/main.go b/main.go index 9837842..c7a7871 100644 --- a/main.go +++ b/main.go @@ -12,156 +12,194 @@ import ( "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 ( k32 = syscall.NewLazyDLL("kernel32.dll") ps = syscall.NewLazyDLL("psapi.dll") - pCreateSnap = k32.NewProc("CreateToolhelp32Snapshot") - pProcess32First = k32.NewProc("Process32FirstW") - pProcess32Next = k32.NewProc("Process32NextW") - pGetProcMemInfo = ps.NewProc("GetProcessMemoryInfo") - pSetPriority = k32.NewProc("SetPriorityClass") - pGetPriority = k32.NewProc("GetPriorityClass") - pOpenProcess = k32.NewProc("OpenProcess") - pCloseHandle = k32.NewProc("CloseHandle") - pCreateMutex = k32.NewProc("CreateMutexW") - pMessageBox = k32.NewProc("MessageBoxW") - pGetLastError = k32.NewProc("GetLastError") + 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") + 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) { - _, _, _ = pCloseHandle.Call(uintptr(h)) + _, _, _ = procCloseHandle.Call(uintptr(h)) } -func openProc(pid uint32, acc uint32) syscall.Handle { - r, _, _ := pOpenProcess.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))) +func openProc(pid uint32, acc uint32) (syscall.Handle, error) { + r, _, e := procOpenProcess.Call(uintptr(acc), 0, uintptr(pid)) if r == 0 { - return 0 + return 0, e } - return uint64(m.WorkingSet) + return syscall.Handle(r), nil } -func allProcs() []procInfo { - var out []procInfo - pe := pe32{Size: uint32(unsafe.Sizeof(pe32{}))} - snap, _, _ := pCreateSnap.Call(SNAPPROCESS, 0) +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 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 { - return out + return nil, e } defer closeH(syscall.Handle(snap)) - r, _, _ := pProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe))) - if r != 0 { - for { - out = append(out, procInfo{ - PID: pe.PID, - Name: syscall.UTF16ToString(pe.ExeFile[:]), - }) - r, _, _ := pProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe))) - if r == 0 { - break - } + pe := processEntry32{Size: uint32(unsafe.Sizeof(processEntry32{}))} + r, _, e := procProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe))) + if r == 0 { + return nil, e + } + + var out []processInfo + for { + name := utf16Trim(pe.ExeFile[:]) + 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 { switch c { - case IDLE_PRIO: + case PriorityClassIdle: return "IDLE" - case NORMAL_PRIO: + case PriorityClassNormal: return "NORMAL" - case HIGH_PRIO: + case PriorityClassHigh: return "HIGH" default: 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() { 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") - doLog := flag.Bool("log", true, "write log file to temp directory") flag.Parse() if *interval < 10*time.Second { @@ -169,17 +207,12 @@ func main() { } var logFile *os.File - if *doLog { - logDir := os.TempDir() - if logDir == "" { - exePath, _ := os.Executable() - logDir = filepath.Dir(exePath) - } + logDir := os.TempDir() + if logDir != "" { 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 { 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) - // Single instance - mh, _, _ := pCreateMutex.Call(0, 0, uintptr(unsafe.Pointer(wide("Global\\autopriority_v1")))) - lastErr, _, _ := pGetLastError.Call() - if lastErr == uintptr(EXISTING) { - closeH(syscall.Handle(mh)) - pMessageBox.Call(0, - uintptr(unsafe.Pointer(wide("autoPriority is already running."))), - uintptr(unsafe.Pointer(wide("autoPriority"))), - 0x30, - ) - os.Exit(1) + ensureSingleInstance() + + 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") } - // 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) - // Restore on exit + if logFile != nil { + defer logFile.Close() + } defer func() { for pid, name := range promoted { - if setPrio(pid, NORMAL_PRIO) { - logf("RESTORE %s (PID %d) → NORMAL", name, pid) + if err := setPrio(pid, PriorityClassNormal); err == nil { + logf("RESTORE %s (PID %d) -> NORMAL", name, pid) } } logf("autoPriority stopped") @@ -232,30 +257,56 @@ func main() { defer ticker.Stop() 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 { continue } if _, ok := promoted[p.PID]; ok { 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 } - orig := curPrio(p.PID) + rssMB := rssBytes / 1024 / 1024 + origPrio, prioErr := curPrio(p.PID) + origStr := prioName(origPrio) + if prioErr != nil { + origStr = "unknown" + } if *dryRun { 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 } - if setPrio(p.PID, HIGH_PRIO) { - promoted[p.PID] = p.Name - logf("PROMOTE %s (PID %d) RSS=%dMB, %s → HIGH", - p.Name, p.PID, rss(p.PID)/1024/1024, prioName(orig)) + if err := setPrio(p.PID, PriorityClassHigh); err != nil { + logf("setPrio(%d, HIGH) error: %v", p.PID, err) + continue } + + promoted[p.PID] = p.Name + logf("PROMOTE %s (PID %d) RSS=%dMB, %s -> HIGH", + p.Name, p.PID, rssMB, origStr) } }