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. Two detection paths: an explicit process watch list
// (GAME_PROCS) and a foreign-VRAM threshold via nvidia-smi
// (GPU_FOREIGN_VRAM_MB) that catches anything not on the ignore list.
package game

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os/exec"
	"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 threshold) 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
	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
}

// New builds a Detector from the configured watch list, VRAM threshold in
// MiB (0 disables the nvidia-smi path) and ignore list. Names are matched
// case-insensitively, with or without a trailing ".exe".
func New(procs []string, vramMB int, ignore []string, log *slog.Logger) *Detector {
	if log == nil {
		log = slog.Default()
	}
	return &Detector{
		procs:  nameSet(procs),
		vramMB: vramMB,
		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).
func (d *Detector) Check(ctx context.Context) ([]string, error) {
	ps, psErr := processes()
	if d.vramMB <= 0 || d.noNvidia {
		return d.detect(ps, nil), psErr
	}
	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")
		return d.detect(ps, nil), nil
	}
	if err != nil {
		return d.detect(ps, nil), err
	}
	return d.detect(ps, apps), nil
}

// detect is the pure core of Check: given the process table and (optionally)
// the nvidia-smi compute-apps list, it returns the foreign holders.
func (d *Detector) detect(ps []Process, apps []computeApp) []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 {
		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))
		}
	}
	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))
}

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