rewrite: single file, zero deps, Windows API, logging, CLI flags
This commit is contained in:
@@ -0,0 +1,262 @@
|
||||
//go:build windows
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"syscall"
|
||||
"time"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
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")
|
||||
)
|
||||
|
||||
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))
|
||||
}
|
||||
|
||||
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)))
|
||||
if r == 0 {
|
||||
return 0
|
||||
}
|
||||
return uint64(m.WorkingSet)
|
||||
}
|
||||
|
||||
func allProcs() []procInfo {
|
||||
var out []procInfo
|
||||
pe := pe32{Size: uint32(unsafe.Sizeof(pe32{}))}
|
||||
snap, _, _ := pCreateSnap.Call(SNAPPROCESS, 0)
|
||||
if snap == 0 {
|
||||
return out
|
||||
}
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func prioName(c uint32) string {
|
||||
switch c {
|
||||
case IDLE_PRIO:
|
||||
return "IDLE"
|
||||
case NORMAL_PRIO:
|
||||
return "NORMAL"
|
||||
case HIGH_PRIO:
|
||||
return "HIGH"
|
||||
default:
|
||||
return fmt.Sprintf("0x%X", c)
|
||||
}
|
||||
}
|
||||
|
||||
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 next to executable")
|
||||
flag.Parse()
|
||||
|
||||
if *interval < 10*time.Second {
|
||||
*interval = 10 * time.Second
|
||||
}
|
||||
|
||||
var logFile *os.File
|
||||
if *doLog {
|
||||
exePath, _ := os.Executable()
|
||||
f, err := os.OpenFile(filepath.Join(filepath.Dir(exePath), "autopriority.log"),
|
||||
os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 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...)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
// 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
|
||||
defer func() {
|
||||
for pid, name := range promoted {
|
||||
if setPrio(pid, NORMAL_PRIO) {
|
||||
logf("RESTORE %s (PID %d) → NORMAL", name, pid)
|
||||
}
|
||||
}
|
||||
logf("autoPriority stopped")
|
||||
}()
|
||||
|
||||
myPID := uint32(os.Getpid())
|
||||
ticker := time.NewTicker(*interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
scan := func() {
|
||||
for _, p := range allProcs() {
|
||||
if p.PID == myPID || p.PID == 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := promoted[p.PID]; ok {
|
||||
continue
|
||||
}
|
||||
if rss(p.PID) < *mem {
|
||||
continue
|
||||
}
|
||||
|
||||
orig := curPrio(p.PID)
|
||||
|
||||
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))
|
||||
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))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
scan()
|
||||
for range ticker.C {
|
||||
scan()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user