gpu-turnstile/internal/config/config.go

// Package config loads gpu-turnstile's configuration from environment
// variables and an optional .env-style config file.
package config

import (
	"bufio"
	"fmt"
	"io"
	"log/slog"
	"strconv"
	"strings"
	"time"
)

// Config holds every gpu-turnstile setting.
type Config struct {
	ListenOllama   string
	ListenComfy    string
	OllamaURL      string
	ComfyURL       string
	UnloadTimeout  time.Duration
	JobTimeout     time.Duration
	LLMWaitTimeout time.Duration

	UnloadPollInterval  time.Duration
	HistoryPollInterval time.Duration
	ProbeTimeout        time.Duration
	FreeTimeout         time.Duration
	WarmTimeout         time.Duration
	ShutdownTimeout     time.Duration
	BackoffInitial      time.Duration
	BackoffMax          time.Duration
	PromptCaptureLimit  int64

	AutoUpdate     bool
	UpdateInterval time.Duration
	UpdateRepo     string
	UpdateAsset    string

	WarmModel string
	LogLevel  slog.Level
	LogJSON   bool
	LogFile   string
}

// Defaults returns the configuration used when neither the environment nor
// a config file sets a value.
func Defaults() Config {
	return Config{
		ListenOllama:   ":11434",
		ListenComfy:    ":8188",
		OllamaURL:      "http://127.0.0.1:11435",
		ComfyURL:       "http://127.0.0.1:8189",
		UnloadTimeout:  time.Minute,
		JobTimeout:     15 * time.Minute,
		LLMWaitTimeout: 10 * time.Minute,

		UnloadPollInterval:  500 * time.Millisecond,
		HistoryPollInterval: time.Second,
		ProbeTimeout:        5 * time.Second,
		FreeTimeout:         30 * time.Second,
		WarmTimeout:         2 * time.Minute,
		ShutdownTimeout:     10 * time.Second,
		BackoffInitial:      time.Second,
		BackoffMax:          time.Minute,
		PromptCaptureLimit:  64 * 1024,

		AutoUpdate:     true,
		UpdateInterval: 6 * time.Hour,
		UpdateRepo:     "https://git.rambossek.at/PUBLIC/gpu-turnstile",
		UpdateAsset:    "gpu-turnstile.exe",

		LogLevel: slog.LevelWarn,
	}
}

// ParseEnvFile parses a .env-style file: KEY=VALUE lines, blank lines and
// #-comments are ignored, no quoting. A line without '=' is an error.
func ParseEnvFile(r io.Reader) (map[string]string, error) {
	values := make(map[string]string)
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 64*1024), 1024*1024)
	lineNo := 0
	for scanner.Scan() {
		lineNo++
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		key, value, ok := strings.Cut(line, "=")
		if !ok {
			return nil, fmt.Errorf("line %d: expected KEY=VALUE", lineNo)
		}
		key = strings.TrimSpace(key)
		if key == "" {
			return nil, fmt.Errorf("line %d: empty key", lineNo)
		}
		values[key] = strings.TrimSpace(value)
	}
	return values, scanner.Err()
}

func envDuration(getenv func(string) string, name string, dst *time.Duration) error {
	v := getenv(name)
	if v == "" {
		return nil
	}
	d, err := time.ParseDuration(v)
	if err != nil {
		return fmt.Errorf("%s: %w", name, err)
	}
	*dst = d
	return nil
}

// Load overlays values from getenv onto the Defaults. Unknown keys are
// ignored. Invalid values are fatal.
func Load(getenv func(string) string) (Config, error) {
	cfg := Defaults()
	for _, e := range []struct {
		name string
		dst  *string
	}{
		{"LISTEN_OLLAMA", &cfg.ListenOllama},
		{"LISTEN_COMFY", &cfg.ListenComfy},
		{"OLLAMA_URL", &cfg.OllamaURL},
		{"COMFY_URL", &cfg.ComfyURL},
		{"WARM_MODEL", &cfg.WarmModel},
		{"UPDATE_REPO", &cfg.UpdateRepo},
		{"UPDATE_ASSET", &cfg.UpdateAsset},
		{"LOG_FILE", &cfg.LogFile},
	} {
		if v := getenv(e.name); v != "" {
			*e.dst = v
		}
	}
	for _, e := range []struct {
		name string
		dst  *time.Duration
	}{
		{"UNLOAD_TIMEOUT", &cfg.UnloadTimeout},
		{"JOB_TIMEOUT", &cfg.JobTimeout},
		{"LLM_WAIT_TIMEOUT", &cfg.LLMWaitTimeout},
		{"UNLOAD_POLL_INTERVAL", &cfg.UnloadPollInterval},
		{"HISTORY_POLL_INTERVAL", &cfg.HistoryPollInterval},
		{"PROBE_TIMEOUT", &cfg.ProbeTimeout},
		{"FREE_TIMEOUT", &cfg.FreeTimeout},
		{"WARM_TIMEOUT", &cfg.WarmTimeout},
		{"SHUTDOWN_TIMEOUT", &cfg.ShutdownTimeout},
		{"BACKOFF_INITIAL", &cfg.BackoffInitial},
		{"BACKOFF_MAX", &cfg.BackoffMax},
		{"UPDATE_INTERVAL", &cfg.UpdateInterval},
	} {
		if err := envDuration(getenv, e.name, e.dst); err != nil {
			return cfg, err
		}
	}
	if v := getenv("PROMPT_CAPTURE_LIMIT"); v != "" {
		n, err := strconv.ParseInt(v, 10, 64)
		if err != nil || n < 0 {
			return cfg, fmt.Errorf("PROMPT_CAPTURE_LIMIT: must be a non-negative integer (bytes)")
		}
		cfg.PromptCaptureLimit = n
	}
	if v := getenv("AUTO_UPDATE"); v != "" {
		b, err := strconv.ParseBool(v)
		if err != nil {
			return cfg, fmt.Errorf("AUTO_UPDATE: must be a boolean (true/false)")
		}
		cfg.AutoUpdate = b
	}
	// LOGLEVEL is the canonical spelling; LOG_LEVEL is kept as an alias.
	logLevelValue := getenv("LOGLEVEL")
	if logLevelValue == "" {
		logLevelValue = getenv("LOG_LEVEL")
	}
	if logLevelValue != "" {
		var level slog.Level
		if err := level.UnmarshalText([]byte(logLevelValue)); err != nil {
			return cfg, fmt.Errorf("LOGLEVEL: %w", err)
		}
		cfg.LogLevel = level
	}
	switch strings.ToLower(getenv("LOG_FORMAT")) {
	case "", "text":
	case "json":
		cfg.LogJSON = true
	default:
		return cfg, fmt.Errorf("LOG_FORMAT: must be \"text\" or \"json\"")
	}
	return cfg, nil
}