Env names in reload diff; monitor shows last VRAM check result; harden control channel against floods
- diffConfig reports the env var names (from new struct tags) instead of Go field names, so the reload-env reply names what the user can change - the monitor's GPU line now shows the game detector's last finding (external holders or none) and how long ago the check ran - control channel: 10s per-connection watchdog (abortive force-close), cap of 32 concurrent connections, reload-env rate-limited; command read was already capped at 4 KiB
This commit is contained in:
+64
-22
@@ -728,9 +728,10 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
|
|
||||||
// Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS
|
// Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS
|
||||||
// and/or GPU_FOREIGN_VRAM_MB — hold the lock externally while they run.
|
// and/or GPU_FOREIGN_VRAM_MB — hold the lock externally while they run.
|
||||||
// gw collects the VRAM reading for the status channel.
|
// gw collects the VRAM reading and the last check result for the status
|
||||||
gw := &gpuWatch{}
|
// channel.
|
||||||
if len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0 {
|
gw := &gpuWatch{enabled: len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0}
|
||||||
|
if gw.enabled {
|
||||||
det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUIgnoreProcs, log)
|
det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUIgnoreProcs, log)
|
||||||
go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup, gw)
|
go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup, gw)
|
||||||
}
|
}
|
||||||
@@ -836,26 +837,40 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
|||||||
// (idle); health checks skip it instead of logging an outage.
|
// (idle); health checks skip it instead of logging an outage.
|
||||||
var errManagedDown = errors.New("managed upstream intentionally stopped")
|
var errManagedDown = errors.New("managed upstream intentionally stopped")
|
||||||
|
|
||||||
// gpuWatch records the latest VRAM reading from the game detector's poll
|
// gpuWatch records the latest VRAM reading and game-detector finding from
|
||||||
// loop, for the status channel. Known stays false when game detection is
|
// the game detector's poll loop, for the status channel. Enabled is false
|
||||||
// not configured (no nvidia-smi polling happens then).
|
// when game detection is not configured (no polling happens then); Known is
|
||||||
|
// false until the first successful nvidia-smi reading.
|
||||||
type gpuWatch struct {
|
type gpuWatch struct {
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
usedMB int
|
enabled bool
|
||||||
total int
|
usedMB int
|
||||||
known bool
|
total int
|
||||||
|
known bool
|
||||||
|
foreign string // last Check result: external holders, "" when none
|
||||||
|
at time.Time
|
||||||
}
|
}
|
||||||
|
|
||||||
func (g *gpuWatch) set(used, total int) {
|
func (g *gpuWatch) setVRAM(used, total int) {
|
||||||
g.mu.Lock()
|
g.mu.Lock()
|
||||||
g.usedMB, g.total, g.known = used, total, true
|
g.usedMB, g.total, g.known = used, total, true
|
||||||
g.mu.Unlock()
|
g.mu.Unlock()
|
||||||
}
|
}
|
||||||
|
|
||||||
func (g *gpuWatch) get() (used, total int, known bool) {
|
func (g *gpuWatch) setCheck(foreign string) {
|
||||||
|
g.mu.Lock()
|
||||||
|
g.foreign, g.at = foreign, time.Now()
|
||||||
|
g.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *gpuWatch) get() (used, total int, known bool, foreign string, ageS int64) {
|
||||||
g.mu.Lock()
|
g.mu.Lock()
|
||||||
defer g.mu.Unlock()
|
defer g.mu.Unlock()
|
||||||
return g.usedMB, g.total, g.known
|
ageS = -1
|
||||||
|
if !g.at.IsZero() {
|
||||||
|
ageS = int64(time.Since(g.at).Seconds())
|
||||||
|
}
|
||||||
|
return g.usedMB, g.total, g.known, g.foreign, ageS
|
||||||
}
|
}
|
||||||
|
|
||||||
// gameLoop polls for foreign GPU holders (a game, another ML job). While one
|
// gameLoop polls for foreign GPU holders (a game, another ML job). While one
|
||||||
@@ -877,8 +892,9 @@ func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *gam
|
|||||||
if err != nil && ctx.Err() == nil {
|
if err != nil && ctx.Err() == nil {
|
||||||
log.Warn("game detection failed", "err", err)
|
log.Warn("game detection failed", "err", err)
|
||||||
}
|
}
|
||||||
|
gw.setCheck(summarizeHolders(holders))
|
||||||
if used, total, verr := game.QueryVRAMMB(ctx); verr == nil {
|
if used, total, verr := game.QueryVRAMMB(ctx); verr == nil {
|
||||||
gw.set(used, total)
|
gw.setVRAM(used, total)
|
||||||
}
|
}
|
||||||
switch {
|
switch {
|
||||||
case len(holders) > 0 && !held:
|
case len(holders) > 0 && !held:
|
||||||
@@ -1012,12 +1028,21 @@ func updateLoop(ctx context.Context, interval time.Duration, log *slog.Logger, u
|
|||||||
// restarts. u is nil when AUTO_UPDATE=false.
|
// restarts. u is nil when AUTO_UPDATE=false.
|
||||||
func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string), status func() string, reload func() string) {
|
func serveControl(ctx context.Context, log *slog.Logger, u *update.Updater, exePath string, applyStaged func(to string), status func() string, reload func() string) {
|
||||||
var mu sync.Mutex
|
var mu sync.Mutex
|
||||||
var lastTrigger time.Time
|
var lastTrigger, lastReload time.Time
|
||||||
h := func(cmd string) string {
|
h := func(cmd string) string {
|
||||||
switch cmd {
|
switch cmd {
|
||||||
case control.CmdStatus:
|
case control.CmdStatus:
|
||||||
return "OK " + status()
|
return "OK " + status()
|
||||||
case control.CmdReloadEnv:
|
case control.CmdReloadEnv:
|
||||||
|
// reload re-reads the config file from disk; a short limiter
|
||||||
|
// keeps a local flood from turning into disk churn.
|
||||||
|
mu.Lock()
|
||||||
|
if wait := 2*time.Second - time.Since(lastReload); wait > 0 {
|
||||||
|
mu.Unlock()
|
||||||
|
return fmt.Sprintf("ERR rate limited: retry in %ds", int(wait.Seconds())+1)
|
||||||
|
}
|
||||||
|
lastReload = time.Now()
|
||||||
|
mu.Unlock()
|
||||||
return reload()
|
return reload()
|
||||||
case control.CmdUpdateNow:
|
case control.CmdUpdateNow:
|
||||||
default:
|
default:
|
||||||
@@ -1097,8 +1122,9 @@ type statusSnapshot struct {
|
|||||||
UptimeS int64 `json:"uptime_s"`
|
UptimeS int64 `json:"uptime_s"`
|
||||||
Downstreams []statusDownstream `json:"downstreams"`
|
Downstreams []statusDownstream `json:"downstreams"`
|
||||||
Lock statusLock `json:"lock"`
|
Lock statusLock `json:"lock"`
|
||||||
// GPU carries the latest VRAM reading; Known is false when game
|
// GPU carries the latest VRAM reading and detector finding; Enabled is
|
||||||
// detection (and with it nvidia-smi polling) is not configured.
|
// false when game detection (and with it nvidia-smi polling) is not
|
||||||
|
// configured.
|
||||||
GPU statusGPU `json:"gpu"`
|
GPU statusGPU `json:"gpu"`
|
||||||
// MonitorNote is set client-side (never over the wire) when the
|
// MonitorNote is set client-side (never over the wire) when the
|
||||||
// monitor's own binary differs from the service's version.
|
// monitor's own binary differs from the service's version.
|
||||||
@@ -1106,9 +1132,17 @@ type statusSnapshot struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type statusGPU struct {
|
type statusGPU struct {
|
||||||
|
// Enabled reports whether game detection is configured (and with it
|
||||||
|
// VRAM polling); when false the other fields carry no information.
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
UsedMB int `json:"used_mb"`
|
UsedMB int `json:"used_mb"`
|
||||||
TotalMB int `json:"total_mb"`
|
TotalMB int `json:"total_mb"`
|
||||||
Known bool `json:"known"`
|
Known bool `json:"known"`
|
||||||
|
// Foreign is the last detector finding (external GPU holders), empty
|
||||||
|
// when the last check found none.
|
||||||
|
Foreign string `json:"foreign,omitempty"`
|
||||||
|
// AgeS is how long ago the last check ran; -1 before the first check.
|
||||||
|
AgeS int64 `json:"age_s"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// reloadHandler re-reads and validates the service's config file for
|
// reloadHandler re-reads and validates the service's config file for
|
||||||
@@ -1130,15 +1164,19 @@ func reloadHandler(current config.Config, configPath string, restartWhenIdle fun
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// diffConfig lists the names of fields whose values differ between two
|
// diffConfig lists the env names of settings whose values differ between
|
||||||
// configs.
|
// two configs (from each field's env tag, so users recognize them).
|
||||||
func diffConfig(a, b config.Config) []string {
|
func diffConfig(a, b config.Config) []string {
|
||||||
va, vb := reflect.ValueOf(a), reflect.ValueOf(b)
|
va, vb := reflect.ValueOf(a), reflect.ValueOf(b)
|
||||||
t := va.Type()
|
t := va.Type()
|
||||||
var out []string
|
var out []string
|
||||||
for i := 0; i < t.NumField(); i++ {
|
for i := 0; i < t.NumField(); i++ {
|
||||||
if !reflect.DeepEqual(va.Field(i).Interface(), vb.Field(i).Interface()) {
|
if !reflect.DeepEqual(va.Field(i).Interface(), vb.Field(i).Interface()) {
|
||||||
out = append(out, t.Field(i).Name)
|
name := t.Field(i).Tag.Get("env")
|
||||||
|
if name == "" {
|
||||||
|
name = t.Field(i).Name
|
||||||
|
}
|
||||||
|
out = append(out, name)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return out
|
return out
|
||||||
@@ -1151,8 +1189,12 @@ func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Proces
|
|||||||
Version: version,
|
Version: version,
|
||||||
UptimeS: int64(time.Since(started).Seconds()),
|
UptimeS: int64(time.Since(started).Seconds()),
|
||||||
}
|
}
|
||||||
used, total, known := gw.get()
|
used, total, known, foreign, ageS := gw.get()
|
||||||
snap.GPU = statusGPU{UsedMB: used, TotalMB: total, Known: known}
|
snap.GPU = statusGPU{
|
||||||
|
Enabled: gw.enabled,
|
||||||
|
UsedMB: used, TotalMB: total, Known: known,
|
||||||
|
Foreign: foreign, AgeS: ageS,
|
||||||
|
}
|
||||||
if cfg.OllamaURL != "" {
|
if cfg.OllamaURL != "" {
|
||||||
snap.Downstreams = append(snap.Downstreams, statusDownstream{
|
snap.Downstreams = append(snap.Downstreams, statusDownstream{
|
||||||
Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
|
Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
|
||||||
|
|||||||
@@ -196,8 +196,8 @@ func renderMonitor(snap statusSnapshot, width int) string {
|
|||||||
b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n",
|
b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n",
|
||||||
cYellow, snap.Lock.ImageQueue, cReset))
|
cYellow, snap.Lock.ImageQueue, cReset))
|
||||||
}
|
}
|
||||||
if snap.GPU.Known {
|
if snap.GPU.Enabled {
|
||||||
b.WriteString(" GPU: " + renderVRAM(snap.GPU.UsedMB, snap.GPU.TotalMB) + "\x1b[K\n")
|
b.WriteString(renderGPU(snap.GPU) + "\x1b[K\n")
|
||||||
}
|
}
|
||||||
if snap.MonitorNote != "" {
|
if snap.MonitorNote != "" {
|
||||||
b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
|
b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
|
||||||
@@ -207,6 +207,26 @@ func renderMonitor(snap statusSnapshot, width int) string {
|
|||||||
return b.String()
|
return b.String()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// renderGPU renders the GPU line: VRAM usage (when nvidia-smi answered),
|
||||||
|
// the game detector's last finding, and how long ago it ran.
|
||||||
|
func renderGPU(g statusGPU) string {
|
||||||
|
s := " GPU: "
|
||||||
|
if g.Known {
|
||||||
|
s += renderVRAM(g.UsedMB, g.TotalMB)
|
||||||
|
} else {
|
||||||
|
s += cDim + "VRAM unknown (nvidia-smi not answering)" + cReset
|
||||||
|
}
|
||||||
|
if g.Foreign != "" {
|
||||||
|
s += " · external: " + cRed + g.Foreign + cReset
|
||||||
|
} else {
|
||||||
|
s += cDim + " · no external process" + cReset
|
||||||
|
}
|
||||||
|
if g.AgeS >= 0 {
|
||||||
|
s += cDim + fmt.Sprintf(" (checked %s ago)", fmtDur(g.AgeS)) + cReset
|
||||||
|
}
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
|
// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
|
||||||
func renderVRAM(used, total int) string {
|
func renderVRAM(used, total int) string {
|
||||||
format := func(mb int) string {
|
format := func(mb int) string {
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package main
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"gpu-turnstile/internal/config"
|
"gpu-turnstile/internal/config"
|
||||||
|
"reflect"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
)
|
)
|
||||||
@@ -30,6 +31,22 @@ func TestRenderMonitor(t *testing.T) {
|
|||||||
t.Errorf("frame missing %q:\n%s", want, frame)
|
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
snap.GPU = statusGPU{Enabled: true, Known: true, UsedMB: 4300, TotalMB: 16384, Foreign: "cyberpunk2077.exe (pid 1234)", AgeS: 12}
|
||||||
|
frame = renderMonitor(snap, 80)
|
||||||
|
for _, want := range []string{"GPU:", "4.2 GiB / 16.0 GiB used", "external:", "checked 12s ago"} {
|
||||||
|
if !strings.Contains(frame, want) {
|
||||||
|
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
snap.GPU = statusGPU{Enabled: true, AgeS: 4}
|
||||||
|
frame = renderMonitor(snap, 80)
|
||||||
|
for _, want := range []string{"VRAM unknown", "no external process"} {
|
||||||
|
if !strings.Contains(frame, want) {
|
||||||
|
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestFmtDur(t *testing.T) {
|
func TestFmtDur(t *testing.T) {
|
||||||
@@ -50,7 +67,19 @@ func TestDiffConfig(t *testing.T) {
|
|||||||
b.LogLevel = -4
|
b.LogLevel = -4
|
||||||
b.GameProcs = []string{"game.exe"}
|
b.GameProcs = []string{"game.exe"}
|
||||||
got := diffConfig(a, b)
|
got := diffConfig(a, b)
|
||||||
if len(got) != 2 || got[0] != "GameProcs" || got[1] != "LogLevel" {
|
if len(got) != 2 || got[0] != "GAME_PROCS" || got[1] != "LOGLEVEL" {
|
||||||
t.Fatalf("got %v, want [GameProcs LogLevel]", got)
|
t.Fatalf("got %v, want [GAME_PROCS LOGLEVEL]", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Every Config field must carry an env tag so user-facing output (the
|
||||||
|
// reload diff) can name the setting the user would actually change.
|
||||||
|
func TestConfigFieldsHaveEnvTags(t *testing.T) {
|
||||||
|
typ := reflect.TypeOf(config.Config{})
|
||||||
|
for i := 0; i < typ.NumField(); i++ {
|
||||||
|
f := typ.Field(i)
|
||||||
|
if f.Tag.Get("env") == "" {
|
||||||
|
t.Errorf("config.Config.%s has no env tag", f.Name)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+40
-38
@@ -13,43 +13,45 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Config holds every gpu-turnstile setting.
|
// Config holds every gpu-turnstile setting. Each field's env tag names the
|
||||||
|
// environment variable / config-file key that sets it; user-facing output
|
||||||
|
// (e.g. the reload diff) uses those names, never the Go field names.
|
||||||
type Config struct {
|
type Config struct {
|
||||||
ListenOllama string
|
ListenOllama string `env:"LISTEN_OLLAMA"`
|
||||||
ListenComfy string
|
ListenComfy string `env:"LISTEN_COMFY"`
|
||||||
OllamaURL string
|
OllamaURL string `env:"OLLAMA_URL"`
|
||||||
ComfyURL string
|
ComfyURL string `env:"COMFY_URL"`
|
||||||
UnloadTimeout time.Duration
|
UnloadTimeout time.Duration `env:"UNLOAD_TIMEOUT"`
|
||||||
JobTimeout time.Duration
|
JobTimeout time.Duration `env:"JOB_TIMEOUT"`
|
||||||
LLMWaitTimeout time.Duration
|
LLMWaitTimeout time.Duration `env:"LLM_WAIT_TIMEOUT"`
|
||||||
|
|
||||||
UnloadPollInterval time.Duration
|
UnloadPollInterval time.Duration `env:"UNLOAD_POLL_INTERVAL"`
|
||||||
HistoryPollInterval time.Duration
|
HistoryPollInterval time.Duration `env:"HISTORY_POLL_INTERVAL"`
|
||||||
ProbeTimeout time.Duration
|
ProbeTimeout time.Duration `env:"PROBE_TIMEOUT"`
|
||||||
HealthInterval time.Duration
|
HealthInterval time.Duration `env:"HEALTH_INTERVAL"`
|
||||||
FreeTimeout time.Duration
|
FreeTimeout time.Duration `env:"FREE_TIMEOUT"`
|
||||||
WarmTimeout time.Duration
|
WarmTimeout time.Duration `env:"WARM_TIMEOUT"`
|
||||||
ShutdownTimeout time.Duration
|
ShutdownTimeout time.Duration `env:"SHUTDOWN_TIMEOUT"`
|
||||||
BackoffInitial time.Duration
|
BackoffInitial time.Duration `env:"BACKOFF_INITIAL"`
|
||||||
BackoffMax time.Duration
|
BackoffMax time.Duration `env:"BACKOFF_MAX"`
|
||||||
PromptCaptureLimit int64
|
PromptCaptureLimit int64 `env:"PROMPT_CAPTURE_LIMIT"`
|
||||||
|
|
||||||
AutoUpdate bool
|
AutoUpdate bool `env:"AUTO_UPDATE"`
|
||||||
UpdateInterval time.Duration
|
UpdateInterval time.Duration `env:"UPDATE_INTERVAL"`
|
||||||
UpdateRepo string
|
UpdateRepo string `env:"UPDATE_REPO"`
|
||||||
UpdateAsset string
|
UpdateAsset string `env:"UPDATE_ASSET"`
|
||||||
|
|
||||||
// AppVersion is the version the user wants to run: "dev" disables
|
// AppVersion is the version the user wants to run: "dev" disables
|
||||||
// updates, "stable" tracks the latest release, anything else is an
|
// updates, "stable" tracks the latest release, anything else is an
|
||||||
// exact vX.Y.Z release to pin. From APP_VER; defaults to "stable".
|
// exact vX.Y.Z release to pin. From APP_VER; defaults to "stable".
|
||||||
AppVersion string
|
AppVersion string `env:"APP_VER"`
|
||||||
|
|
||||||
// LLMBusyMode is "wait" (hold requests until the lock is free or
|
// LLMBusyMode is "wait" (hold requests until the lock is free or
|
||||||
// LLMWaitTimeout expires) or "reject" (immediately answer with
|
// LLMWaitTimeout expires) or "reject" (immediately answer with
|
||||||
// LLMBusyStatus + Retry-After when an image job is active or pending).
|
// LLMBusyStatus + Retry-After when an image job is active or pending).
|
||||||
LLMBusyMode string
|
LLMBusyMode string `env:"LLM_BUSY_MODE"`
|
||||||
LLMBusyStatus int
|
LLMBusyStatus int `env:"LLM_BUSY_STATUS"`
|
||||||
BusyRetryAfter int
|
BusyRetryAfter int `env:"BUSY_RETRY_AFTER"`
|
||||||
|
|
||||||
// ComfyCmd spawns and supervises a ComfyUI server on demand. When
|
// ComfyCmd spawns and supervises a ComfyUI server on demand. When
|
||||||
// ComfyCmd is empty but ComfyDir is set, management is enabled with the
|
// ComfyCmd is empty but ComfyDir is set, management is enabled with the
|
||||||
@@ -59,10 +61,10 @@ type Config struct {
|
|||||||
// set explicitly. The managed server is stopped after ComfyIdleTimeout
|
// set explicitly. The managed server is stopped after ComfyIdleTimeout
|
||||||
// without requests, freeing its VRAM; ComfyStartTimeout bounds how long
|
// without requests, freeing its VRAM; ComfyStartTimeout bounds how long
|
||||||
// a request waits for it to come up.
|
// a request waits for it to come up.
|
||||||
ComfyCmd string
|
ComfyCmd string `env:"COMFY_CMD"`
|
||||||
ComfyDir string
|
ComfyDir string `env:"COMFY_DIR"`
|
||||||
ComfyIdleTimeout time.Duration
|
ComfyIdleTimeout time.Duration `env:"COMFY_IDLE_TIMEOUT"`
|
||||||
ComfyStartTimeout time.Duration
|
ComfyStartTimeout time.Duration `env:"COMFY_START_TIMEOUT"`
|
||||||
|
|
||||||
// GameProcs (GAME_PROCS) is a watch list of process names; while any of
|
// GameProcs (GAME_PROCS) is a watch list of process names; while any of
|
||||||
// them runs, the GPU is treated as held by a foreign process. The
|
// them runs, the GPU is treated as held by a foreign process. The
|
||||||
@@ -70,15 +72,15 @@ type Config struct {
|
|||||||
// when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than
|
// when a process not in GPUIgnoreProcs (GPU_IGNORE_PROCS) holds more than
|
||||||
// that many MiB of VRAM. GamePollInterval (GAME_POLL_INTERVAL) is how
|
// that many MiB of VRAM. GamePollInterval (GAME_POLL_INTERVAL) is how
|
||||||
// often both checks run.
|
// often both checks run.
|
||||||
GameProcs []string
|
GameProcs []string `env:"GAME_PROCS"`
|
||||||
GPUForeignVRAMMB int
|
GPUForeignVRAMMB int `env:"GPU_FOREIGN_VRAM_MB"`
|
||||||
GPUIgnoreProcs []string
|
GPUIgnoreProcs []string `env:"GPU_IGNORE_PROCS"`
|
||||||
GamePollInterval time.Duration
|
GamePollInterval time.Duration `env:"GAME_POLL_INTERVAL"`
|
||||||
|
|
||||||
WarmModel string
|
WarmModel string `env:"WARM_MODEL"`
|
||||||
LogLevel slog.Level
|
LogLevel slog.Level `env:"LOGLEVEL"`
|
||||||
LogJSON bool
|
LogJSON bool `env:"LOG_FORMAT"`
|
||||||
LogFile string
|
LogFile string `env:"LOG_FILE"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// Defaults returns the configuration used when neither the environment nor
|
// Defaults returns the configuration used when neither the environment nor
|
||||||
|
|||||||
@@ -7,6 +7,13 @@
|
|||||||
// triggers, so the worst a local user can cause is a cheap, throttled
|
// triggers, so the worst a local user can cause is a cheap, throttled
|
||||||
// check and a GPU-idle-gated restart onto a signed binary.
|
// check and a GPU-idle-gated restart onto a signed binary.
|
||||||
//
|
//
|
||||||
|
// Abuse hardening: the command read is capped (4 KiB), each connection is
|
||||||
|
// force-closed after connTimeout so a stalled client cannot pin a goroutine
|
||||||
|
// (or a Windows pipe instance) forever, and concurrently served connections
|
||||||
|
// are capped at maxConns — beyond that, connections are closed on arrival.
|
||||||
|
// On Windows the pipe's ACL additionally denies network logons, so the
|
||||||
|
// channel cannot be reached from another machine.
|
||||||
|
//
|
||||||
// Protocol: the client writes one command line, the server answers with
|
// Protocol: the client writes one command line, the server answers with
|
||||||
// one reply line ("OK ..." or "ERR ...") and hangs up.
|
// one reply line ("OK ..." or "ERR ...") and hangs up.
|
||||||
package control
|
package control
|
||||||
@@ -17,6 +24,7 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"strings"
|
"strings"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
|
// CmdUpdateNow asks the service to check for, stage and (once the GPU is
|
||||||
@@ -37,9 +45,50 @@ var ErrUnavailable = errors.New("control channel unavailable")
|
|||||||
// line. It must start with "OK " or "ERR ".
|
// line. It must start with "OK " or "ERR ".
|
||||||
type Handler func(cmd string) string
|
type Handler func(cmd string) string
|
||||||
|
|
||||||
|
// connTimeout bounds one connection's lifetime: a client that stops
|
||||||
|
// mid-command or never reads the reply would otherwise pin its goroutine
|
||||||
|
// (and on Windows one of the pipe instances) indefinitely. A var so tests
|
||||||
|
// can shrink it.
|
||||||
|
var connTimeout = 10 * time.Second
|
||||||
|
|
||||||
|
// maxConns caps concurrently served connections; beyond it, new
|
||||||
|
// connections are closed on arrival. Bound on the goroutines a local
|
||||||
|
// flood can pile up.
|
||||||
|
const maxConns = 32
|
||||||
|
|
||||||
|
var connSem = make(chan struct{}, maxConns)
|
||||||
|
|
||||||
|
// serve dispatches connection handling under the concurrency cap. It
|
||||||
|
// returns false when the cap is reached — the caller must then close the
|
||||||
|
// connection itself.
|
||||||
|
func serve(c io.ReadWriteCloser, h Handler) bool {
|
||||||
|
select {
|
||||||
|
case connSem <- struct{}{}:
|
||||||
|
go func() {
|
||||||
|
defer func() { <-connSem }()
|
||||||
|
serveConn(c, h)
|
||||||
|
}()
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// forceCloser is implemented by connections that can be torn down
|
||||||
|
// abortively, unblocking pending reads and writes (Windows pipe:
|
||||||
|
// DisconnectNamedPipe; unix socket: a deadline in the past). The
|
||||||
|
// connection watchdog uses it; normal closes still flush the reply.
|
||||||
|
type forceCloser interface {
|
||||||
|
ForceClose() error
|
||||||
|
}
|
||||||
|
|
||||||
// serveConn runs the line protocol on one accepted connection.
|
// serveConn runs the line protocol on one accepted connection.
|
||||||
func serveConn(c io.ReadWriteCloser, h Handler) {
|
func serveConn(c io.ReadWriteCloser, h Handler) {
|
||||||
defer c.Close()
|
defer c.Close()
|
||||||
|
if fc, ok := c.(forceCloser); ok {
|
||||||
|
timer := time.AfterFunc(connTimeout, func() { fc.ForceClose() })
|
||||||
|
defer timer.Stop()
|
||||||
|
}
|
||||||
line, err := bufio.NewReader(io.LimitReader(c, 4096)).ReadString('\n')
|
line, err := bufio.NewReader(io.LimitReader(c, 4096)).ReadString('\n')
|
||||||
cmd := strings.TrimSpace(line)
|
cmd := strings.TrimSpace(line)
|
||||||
if cmd == "" {
|
if cmd == "" {
|
||||||
|
|||||||
@@ -37,12 +37,24 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return // shutting down
|
return // shutting down
|
||||||
}
|
}
|
||||||
go serveConn(c, h)
|
uc := unixConn{c}
|
||||||
|
if !serve(uc, h) {
|
||||||
|
uc.ForceClose()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// unixConn adds an abortive ForceClose to net.Conn: a deadline in the
|
||||||
|
// past fails pending and future I/O immediately.
|
||||||
|
type unixConn struct{ net.Conn }
|
||||||
|
|
||||||
|
func (c unixConn) ForceClose() error {
|
||||||
|
c.SetDeadline(time.Now().Add(-time.Second)) //nolint:errcheck // best effort
|
||||||
|
return c.Conn.Close()
|
||||||
|
}
|
||||||
|
|
||||||
// Ask sends one command to the running service and returns its reply.
|
// Ask sends one command to the running service and returns its reply.
|
||||||
func Ask(cmd string) (string, error) {
|
func Ask(cmd string) (string, error) {
|
||||||
c, err := net.DialTimeout("unix", sockPath, 2*time.Second)
|
c, err := net.DialTimeout("unix", sockPath, 2*time.Second)
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestRoundTrip(t *testing.T) {
|
func TestRoundTrip(t *testing.T) {
|
||||||
@@ -44,3 +45,59 @@ func TestEmptyReplyIsUnavailable(t *testing.T) {
|
|||||||
t.Fatalf("err = %v, want ErrUnavailable", err)
|
t.Fatalf("err = %v, want ErrUnavailable", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestServeCap(t *testing.T) {
|
||||||
|
for i := 0; i < maxConns; i++ {
|
||||||
|
connSem <- struct{}{}
|
||||||
|
}
|
||||||
|
defer func() {
|
||||||
|
for i := 0; i < maxConns; i++ {
|
||||||
|
<-connSem
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
server, client := net.Pipe()
|
||||||
|
defer server.Close()
|
||||||
|
defer client.Close()
|
||||||
|
if serve(server, func(string) string { return "OK" }) {
|
||||||
|
t.Fatal("serve accepted a connection beyond the cap")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// forcePipe records ForceClose calls for the watchdog test.
|
||||||
|
type forcePipe struct {
|
||||||
|
net.Conn
|
||||||
|
forced chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c forcePipe) ForceClose() error {
|
||||||
|
err := c.Conn.Close()
|
||||||
|
close(c.forced)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConnWatchdog(t *testing.T) {
|
||||||
|
old := connTimeout
|
||||||
|
connTimeout = 50 * time.Millisecond
|
||||||
|
defer func() { connTimeout = old }()
|
||||||
|
|
||||||
|
server, client := net.Pipe()
|
||||||
|
defer client.Close()
|
||||||
|
fc := forcePipe{Conn: server, forced: make(chan struct{})}
|
||||||
|
done := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
serveConn(fc, func(string) string { return "OK" })
|
||||||
|
close(done)
|
||||||
|
}()
|
||||||
|
// The client never sends anything; the watchdog must tear the
|
||||||
|
// connection down instead of blocking forever.
|
||||||
|
select {
|
||||||
|
case <-fc.forced:
|
||||||
|
case <-time.After(5 * time.Second):
|
||||||
|
t.Fatal("watchdog did not force-close the stalled connection")
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
case <-time.After(5 * time.Second):
|
||||||
|
t.Fatal("serveConn still blocked after the force close")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -86,7 +86,10 @@ func Serve(ctx context.Context, h Handler, log *slog.Logger) error {
|
|||||||
windows.CloseHandle(pipe)
|
windows.CloseHandle(pipe)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
go serveConn(&pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}, h)
|
conn := &pipeConn{f: os.NewFile(uintptr(pipe), pipePath), h: pipe}
|
||||||
|
if !serve(conn, h) {
|
||||||
|
conn.ForceClose()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
return nil
|
return nil
|
||||||
@@ -114,6 +117,14 @@ func (c *pipeConn) Close() error {
|
|||||||
return c.f.Close()
|
return c.f.Close()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ForceClose aborts the connection without flushing: disconnecting
|
||||||
|
// unblocks pending reads and writes at the cost of possibly discarding an
|
||||||
|
// unread reply. Used by the connection watchdog; normal closes flush.
|
||||||
|
func (c *pipeConn) ForceClose() error {
|
||||||
|
windows.DisconnectNamedPipe(c.h) //nolint:errcheck // best effort
|
||||||
|
return c.f.Close()
|
||||||
|
}
|
||||||
|
|
||||||
// Ask sends one command to the running service and returns its reply.
|
// Ask sends one command to the running service and returns its reply.
|
||||||
func Ask(cmd string) (string, error) {
|
func Ask(cmd string) (string, error) {
|
||||||
name, err := windows.UTF16PtrFromString(pipePath)
|
name, err := windows.UTF16PtrFromString(pipePath)
|
||||||
|
|||||||
Reference in New Issue
Block a user