autoPriority/main.go

//go:build windows

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"
)

type options struct {
	mem, gameMem uint64
	interval     time.Duration
	dryRun       bool
}

func readOptions(args []string, output io.Writer) (options, error) {
	var o options
	fs := flag.NewFlagSet("autoPriority", flag.ContinueOnError)
	fs.SetOutput(output)
	mem := fs.String("mem", "512M", "memory threshold (e.g. 512M, 1G)")
	gameMem := fs.String("game-mem", "2G", "heavy-process threshold; greater than -mem, or 0 to disable")
	fs.DurationVar(&o.interval, "interval", time.Minute, "scan interval (minimum 10s)")
	fs.BoolVar(&o.dryRun, "dry-run", false, "log the plan without changing priorities")
	if err := fs.Parse(args); err != nil {
		return o, err
	}
	if fs.NArg() != 0 {
		return o, fmt.Errorf("unexpected positional arguments")
	}
	var err error
	o.mem, err = parseMemSize(*mem)
	if err != nil || o.mem == 0 {
		return o, fmt.Errorf("invalid -mem: use a positive integer with optional K/M/G suffix")
	}
	o.gameMem, err = parseMemSize(*gameMem)
	if err != nil {
		return o, fmt.Errorf("invalid -game-mem: %w", err)
	}
	if o.gameMem != 0 && o.gameMem <= o.mem {
		return o, fmt.Errorf("-game-mem must be greater than -mem, or 0")
	}
	if o.interval < 10*time.Second {
		o.interval = 10 * time.Second
	}
	return o, nil
}
func runController(ctx context.Context, o options, api processBackend, emit func(scanReport) error) (result error) {
	c := newController(api, o.mem, o.gameMem, o.dryRun, o.interval)
	defer func() {
		report := c.Shutdown()
		result = errors.Join(result, emit(report))
		if report.Unrestored > 0 {
			result = errors.Join(result, fmt.Errorf("не восстановлены приоритеты %d процессов; см. журнал", report.Unrestored))
		}
	}()
	ticker := time.NewTicker(o.interval)
	defer ticker.Stop()
	for {
		if ctx.Err() != nil {
			return nil
		}
		if err := emit(c.Step()); err != nil {
			return err
		}
		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
		}
	}
}
func runApplication(o options) (result error) {
	mutex, exists, err := singleInstance()
	if err != nil {
		return err
	}
	defer closeH(mutex)
	if exists {
		fmt.Fprintln(os.Stderr, "autoPriority: already running")
		return nil
	}
	logFile, err := openRotatingLog(filepath.Join(os.TempDir(), "autopriority.log"), 4<<20)
	if err != nil {
		return err
	}
	defer func() { result = errors.Join(result, logFile.Close()) }()
	emit := func(r scanReport) error { return logFile.Write(formatReport(r, time.Now())) }
	if err = emit(scanReport{Title: fmt.Sprintf("autoPriority: mem=%s, game-mem=%s, interval=%s", formatMemSize(o.mem), formatMemSize(o.gameMem), o.interval), DryRun: o.dryRun}); err != nil {
		return err
	}
	if !o.dryRun {
		pid := uint32(os.Getpid())
		h, e := openProc(pid, ProcessSetInformation)
		if e == nil {
			e = setPrio(h, pid, PriorityClassIdle)
			closeH(h)
		}
		if e != nil {
			if err = emit(scanReport{Errors: []processError{{Name: "autoPriority.exe", PID: pid, Description: fmt.Sprintf("Собственный приоритет IDLE не установлен: %v", e)}}}); err != nil {
				return err
			}
		}
	}
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()
	restored := make(chan struct{})
	tray, err := newTray(cancel, restored)
	if err != nil {
		close(restored)
		return err
	}
	defer func() { result = errors.Join(result, tray.Close()) }()
	// runController performs restoration before this channel closes and before the tray disappears.
	defer close(restored)
	return runController(ctx, o, newWindowsBackend(), emit)
}
func main() {
	o, err := readOptions(os.Args[1:], os.Stderr)
	if errors.Is(err, flag.ErrHelp) {
		return
	}
	if err == nil {
		err = runApplication(o)
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "autoPriority: %v\n", err)
		showError(err)
		os.Exit(1)
	}
}