gpu-turnstile/internal/game/game.go

// Package game detects processes outside gpu-turnstile's control that hold
// the GPU — typically a game — so the proxy can block new GPU work and free
// VRAM while they run. Three detection paths: an explicit process watch
// list (GAME_PROCS), a foreign-VRAM threshold via nvidia-smi
// (GPU_FOREIGN_VRAM_MB), and a per-process GPU 3D-engine utilization
// threshold via Windows PDH counters (GPU_FOREIGN_UTIL_PCT). The latter two
// catch anything not on the ignore list without naming individual games.
package game

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os/exec"
	"slices"
	"strconv"
	"strings"
)

// Process is one running OS process.
type Process struct {
	PID  int
	Name string
}

// computeApp is one process holding GPU memory, as reported by nvidia-smi.
type computeApp struct {
	PID    int
	UsedMB int
}

// Detector checks whether a foreign process holds the GPU. The zero value
// (no watch list, no thresholds) never detects anything; main only starts
// the poll loop when at least one path is configured.
type Detector struct {
	procs    map[string]bool // normalized names from GAME_PROCS
	vramMB   int             // foreign VRAM threshold; 0 = disabled
	utilPct  int             // foreign 3D-engine utilization threshold; 0 = disabled
	ignore   map[string]bool // normalized names never counted as foreign
	log      *slog.Logger
	noNvidia bool              // nvidia-smi was not found; VRAM path disabled for good
	sampler  *gpuEngineSampler // open PDH query, opened lazily on first Check
	noPDH    bool              // engine counters unavailable; util path disabled for good
}

// New builds a Detector from the configured watch list, VRAM threshold in
// MiB, 3D-engine utilization threshold in percent (both 0 = disabled) and
// ignore list. Names are matched case-insensitively, with or without a
// trailing ".exe".
func New(procs []string, vramMB, utilPct int, ignore []string, log *slog.Logger) *Detector {
	if log == nil {
		log = slog.Default()
	}
	return &Detector{
		procs:   nameSet(procs),
		vramMB:  vramMB,
		utilPct: utilPct,
		ignore:  nameSet(ignore),
		log:     log,
	}
}

// normName lowercases a process name and strips a trailing ".exe" so the
// watch and ignore lists match on Windows and Linux spellings alike.
func normName(s string) string {
	return strings.TrimSuffix(strings.ToLower(strings.TrimSpace(s)), ".exe")
}

func nameSet(names []string) map[string]bool {
	set := make(map[string]bool, len(names))
	for _, n := range names {
		if n = normName(n); n != "" {
			set[n] = true
		}
	}
	return set
}

// Check looks once for foreign GPU holders and returns a human-readable
// description of each (empty when the GPU is free for gpu-turnstile's
// consumers). A failing nvidia-smi call is returned as an error only when
// the process list found nothing; a missing nvidia-smi binary disables the
// VRAM path permanently (logged once), as do missing engine counters.
func (d *Detector) Check(ctx context.Context) ([]string, error) {
	ps, psErr := processes()
	utils := d.engineUtil()
	var apps []computeApp
	if d.vramMB > 0 && !d.noNvidia {
		var err error
		apps, err = queryComputeApps(ctx)
		if errors.Is(err, exec.ErrNotFound) {
			d.noNvidia = true
			d.log.Warn("GPU_FOREIGN_VRAM_MB is set but nvidia-smi was not found; VRAM detection disabled")
		} else if err != nil {
			return d.detect(ps, nil, utils), err
		}
	}
	return d.detect(ps, apps, utils), psErr
}

// engineUtil samples per-process 3D-engine utilization via PDH. The first
// call only primes the rate counters and returns nil. A failing open or
// sample disables the path permanently (logged once).
func (d *Detector) engineUtil() map[int]float64 {
	if d.utilPct <= 0 || d.noPDH {
		return nil
	}
	if d.sampler == nil {
		s, err := openGPUEngineSampler()
		if err != nil {
			d.noPDH = true
			d.log.Warn("GPU_FOREIGN_UTIL_PCT is set but per-process GPU counters are unavailable; engine detection disabled", "err", err)
			return nil
		}
		d.sampler = s
	}
	utils, err := d.sampler.sample()
	if err != nil {
		if errors.Is(err, errNotPrimed) {
			return nil
		}
		d.noPDH = true
		d.log.Warn("per-process GPU counters failed; engine detection disabled", "err", err)
		return nil
	}
	return utils
}

// detect is the pure core of Check: given the process table and
// (optionally) the nvidia-smi compute-apps list and the PDH engine
// utilization, it returns the foreign holders.
func (d *Detector) detect(ps []Process, apps []computeApp, utils map[int]float64) []string {
	var holders []string
	for _, p := range ps {
		if d.procs[normName(p.Name)] {
			holders = append(holders, fmt.Sprintf("%s (pid %d)", p.Name, p.PID))
		}
	}
	if (d.vramMB > 0 && apps != nil) || (d.utilPct > 0 && utils != nil) {
		names := make(map[int]string, len(ps))
		for _, p := range ps {
			names[p.PID] = p.Name
		}
		for _, a := range apps {
			name := names[a.PID]
			if d.ignore[normName(name)] || a.UsedMB < d.vramMB {
				continue
			}
			if name == "" {
				name = "unknown process"
			}
			holders = append(holders, fmt.Sprintf("%s (pid %d) using %d MiB VRAM", name, a.PID, a.UsedMB))
		}
		// Sorted for stable output (map iteration order is random).
		pids := make([]int, 0, len(utils))
		for pid := range utils {
			pids = append(pids, pid)
		}
		slices.Sort(pids)
		for _, pid := range pids {
			util := utils[pid]
			name := names[pid]
			if d.ignore[normName(name)] || util < float64(d.utilPct) {
				continue
			}
			if name == "" {
				name = "unknown process"
			}
			holders = append(holders, fmt.Sprintf("%s (pid %d) using %.0f%% GPU", name, pid, util))
		}
	}
	return holders
}

// queryComputeApps runs nvidia-smi and parses the per-process VRAM list.
// Note: under Windows' WDDM driver, nvidia-smi only sees compute
// allocations, so graphics-only games may not appear there — GAME_PROCS is
// the reliable path on Windows; on Linux both work.
func queryComputeApps(ctx context.Context) ([]computeApp, error) {
	out, err := exec.CommandContext(ctx, "nvidia-smi",
		"--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits").Output()
	if err != nil {
		return nil, err
	}
	return parseComputeApps(string(out))
}

// GPUStats is one nvidia-smi reading of the whole card.
type GPUStats struct {
	UsedMB  int
	TotalMB int
	TempC   int // -1 when nvidia-smi reports N/A
	FanPct  int // -1 when N/A (some cards don't expose the fan)
}

// QueryGPUStats returns VRAM usage, temperature and fan speed via
// nvidia-smi. Unlike the per-process list this works under WDDM too.
func QueryGPUStats(ctx context.Context) (GPUStats, error) {
	out, err := exec.CommandContext(ctx, "nvidia-smi",
		"--query-gpu=memory.used,memory.total,temperature.gpu,fan.speed", "--format=csv,noheader,nounits").Output()
	if err != nil {
		return GPUStats{}, err
	}
	return parseGPUStats(string(out))
}

// parseGPUStats parses one "used, total, temp, fan" CSV line (MiB, °C,
// percent). The memory fields must be numeric; temperature and fan fall
// back to -1 on "N/A" and friends.
func parseGPUStats(out string) (GPUStats, error) {
	fields := strings.Split(strings.TrimSpace(out), ",")
	if len(fields) != 4 {
		return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(out))
	}
	num := func(s string) (int, error) {
		return strconv.Atoi(strings.TrimSpace(s))
	}
	optional := func(s string) int {
		n, err := num(s)
		if err != nil {
			return -1
		}
		return n
	}
	var st GPUStats
	var err error
	if st.UsedMB, err = num(fields[0]); err != nil {
		return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(out))
	}
	if st.TotalMB, err = num(fields[1]); err != nil {
		return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(out))
	}
	st.TempC = optional(fields[2])
	st.FanPct = optional(fields[3])
	return st, nil
}

// parseComputeApps parses "pid, used_memory" CSV lines (no header, MiB
// units). Unsupported rows ("N/A" on WDDM) are skipped.
func parseComputeApps(out string) ([]computeApp, error) {
	var apps []computeApp
	for _, line := range strings.Split(out, "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		pidStr, memStr, ok := strings.Cut(line, ",")
		if !ok {
			return nil, fmt.Errorf("nvidia-smi: unexpected line %q", line)
		}
		pid, err := strconv.Atoi(strings.TrimSpace(pidStr))
		if err != nil {
			return nil, fmt.Errorf("nvidia-smi: unexpected pid in %q", line)
		}
		mem, err := strconv.Atoi(strings.TrimSpace(memStr))
		if err != nil {
			continue // "N/A" and friends: unsupported under WDDM
		}
		apps = append(apps, computeApp{PID: pid, UsedMB: mem})
	}
	return apps, nil
}