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

2026-06-19 23:20:10 +03:00
parent d0cafe1332
commit b999b7ec4f
3 changed files with 202 additions and 151 deletions
@@ -1,3 +1,3 @@
 module autopriority

-go 1.24
+go 1.25.0
@@ -12,156 +12,194 @@ import (
 	"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
+	CreateToolhelp32SnapshotProcess = 0x00000002
+	ProcessSetInformation           = 0x00000200
+	ProcessQueryInformation         = 0x00000400
+	PriorityClassIdle               = 0x00000040
+	PriorityClassNormal             = 0x00000020
+	PriorityClassHigh               = 0x00000080
+	MessageBoxIconWarning           = 0x00000030
+	ErrorAlreadyExists              = 1832
 )

-type pe32 struct {
+type processEntry32 struct {
 	Size          uint32
 	CntUsage      uint32
 	PID           uint32
-	HeapID    uintptr
+	DefaultHeapID uintptr
 	ModuleID      uint32
-	Threads   uint32
 	ParentPID     uint32
 	PrioClass     int32
 	Flags         uint32
 	ExeFile       [260]uint16
 }

-type pmc struct {
+type processMemoryCounters struct {
 	CBM                      uint32
-	PageFaults  uint32
-	PeakWS      uintptr
-	WorkingSet  uintptr
-	QuotaPeakPG uintptr
-	QuotaPG     uintptr
-	QuotaPeakPM uintptr
-	QuotaPM     uintptr
-	PeakPF      uintptr
-	Pagefile    uintptr
-	Reserved    [2]uintptr
+	PageFaultCount           uint32
+	PeakWorkingSetSize       uintptr
+	WorkingSetSize           uintptr
+	QuotaPeakPagedPoolUsage  uintptr
+	QuotaPagedPoolUsage      uintptr
+	QuotaPeakNonPagedPoolUsage uintptr
+	QuotaNonPagedPoolUsage    uintptr
+	PeakPagefileUsage        uintptr
+	PagefileUsage            uintptr
+	PrivateUsage             uintptr
 }

-type procInfo struct {
+type processInfo struct {
 	PID  uint32
 	Name string
 }

-func wide(s string) *uint16 {
-	return syscall.StringToUTF16Ptr(s)
-}
+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")
+	procGetLastError   = k32.NewProc("GetLastError")
+)

 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 {
+	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 {
-			out = append(out, procInfo{
-				PID:  pe.PID,
-				Name: syscall.UTF16ToString(pe.ExeFile[:]),
-			})
-			r, _, _ := pProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe)))
+		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)
-		}
+	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)
 		}
 	}