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

// QueryVRAMMB returns used and total GPU VRAM in MiB via nvidia-smi.
// Unlike the per-process list this works under WDDM too.
func QueryVRAMMB(ctx context.Context) (used, total int, err error) {
	out, err := exec.CommandContext(ctx, "nvidia-smi",
		"--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits").Output()
	if err != nil {
		return 0, 0, err
	}
	usedStr, totalStr, ok := strings.Cut(strings.TrimSpace(string(out)), ",")
	if !ok {
		return 0, 0, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(string(out)))
	}
	used, err = strconv.Atoi(strings.TrimSpace(usedStr))
	if err != nil {
		return 0, 0, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(string(out)))
	}
	total, err = strconv.Atoi(strings.TrimSpace(totalStr))
	if err != nil {
		return 0, 0, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(string(out)))
	}
	return used, total, 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
}