autoPriority/main.go

//go:build windows

package main

import (
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"syscall"
	"time"
	"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")

	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) {
	_, _, _ = 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, 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 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 nil, 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, 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, nil
}

func prioName(c uint32) string {
	switch c {
	case PriorityClassIdle:
		return "IDLE"
	case PriorityClassNormal:
		return "NORMAL"
	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")
	flag.Parse()

	if *interval < 10*time.Second {
		*interval = 10 * time.Second
	}

	var logFile *os.File
	logDir := os.TempDir()
	if logDir != "" {
		f, err := os.OpenFile(filepath.Join(logDir, "autopriority.log"),
			os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
		if err == nil {
			logFile = f
		}
	}

	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)

	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")
	}

	promoted := make(map[uint32]string)

	if logFile != nil {
		defer logFile.Close()
	}
	defer func() {
		for pid, name := range promoted {
			if err := setPrio(pid, PriorityClassNormal); err == nil {
				logf("RESTORE %s (PID %d) -> NORMAL", name, pid)
			}
		}
		logf("autoPriority stopped")
	}()

	myPID := uint32(os.Getpid())
	ticker := time.NewTicker(*interval)
	defer ticker.Stop()

	scan := func() {
		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
			}

			rssBytes, err := rss(p.PID)
			if err != nil {
				logf("rss(%d) error: %v", p.PID, err)
				continue
			}
			if rssBytes < *mem {
				continue
			}

			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, rssMB, origStr)
				continue
			}

			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)
		}
	}

	scan()
	for range ticker.C {
		scan()
	}
}