autoPriority/main.go

//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 to temp directory")
	flag.Parse()

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

	var logFile *os.File
	if *doLog {
		logDir := os.TempDir()
		if logDir == "" {
			exePath, _ := os.Executable()
			logDir = filepath.Dir(exePath)
		}
		f, err := os.OpenFile(filepath.Join(logDir, "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()
	}
}