CosmicClash/server/cmd/allocator/main.go

package main

import (
	"context"
	"database/sql"
	"flag"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/cosmic-clash/cosmic-clash/server/agones"
	"github.com/cosmic-clash/cosmic-clash/server/allocator"
	"github.com/cosmic-clash/cosmic-clash/server/migrations"
	"github.com/cosmic-clash/cosmic-clash/server/store"
	_ "github.com/jackc/pgx/v5/stdlib"
)

func main() {
	dsn := flag.String("dsn", os.Getenv("COSMIC_CLASH_POSTGRES_DSN"), "PostgreSQL connection string")
	migrationDir := flag.String("migrations", "migrations", "directory containing numbered SQL migrations")
	agonesURL := flag.String("agones-url", os.Getenv("COSMIC_CLASH_AGONES_URL"), "Agones allocation API base URL")
	namespace := flag.String("agones-namespace", envOrDefault("COSMIC_CLASH_AGONES_NAMESPACE", "default"), "Agones namespace")
	kubernetesTokenPath := flag.String("kubernetes-token-path", envOrDefault("COSMIC_CLASH_KUBERNETES_TOKEN_PATH", "/var/run/secrets/kubernetes.io/serviceaccount/token"), "rotating Kubernetes service-account bearer token")
	kubernetesCAPath := flag.String("kubernetes-ca-path", envOrDefault("COSMIC_CLASH_KUBERNETES_CA_PATH", "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"), "Kubernetes API cluster CA bundle")
	providerTimeout := flag.Duration("provider-timeout", 10*time.Second, "timeout for each Kubernetes/Agones API request")
	readinessMaxStale := flag.Duration("readiness-max-stale", 30*time.Second, "maximum age of the last fully successful allocator cycle")
	transport := flag.String("transport", envOrDefault("COSMIC_CLASH_TRANSPORT", "enet"), "game transport: enet or steam_sdr")
	interval := flag.Duration("interval", time.Second, "allocation poll interval")
	workloadSecret := flag.String("workload-secret", os.Getenv("COSMIC_CLASH_WORKLOAD_SECRET"), "HMAC secret for control-plane-issued workload tokens (see workload/signed_token.go); must match cmd/control-plane's own --workload-secret. Unset skips minting a cosmic-clash.io/workload-token annotation entirely")
	allocationQuota := flag.Int("allocation-quota", 0, "optional per-replica allocation attempts per region per quota window; zero disables this local guard")
	allocationQuotaWindow := flag.Duration("allocation-quota-window", time.Minute, "window for --allocation-quota")
	metricsAddr := flag.String("metrics-addr", envOrDefault("COSMIC_CLASH_ALLOCATOR_METRICS_ADDR", ":9091"), "allocator Prometheus metrics address; empty disables metrics")
	flag.Parse()
	if *dsn == "" || *agonesURL == "" {
		fatalf("--dsn/COSMIC_CLASH_POSTGRES_DSN and --agones-url/COSMIC_CLASH_AGONES_URL are required")
	}
	if (*transport != "enet" && *transport != "steam_sdr") || *interval <= 0 || *providerTimeout <= 0 || *readinessMaxStale <= 0 {
		fatalf("--transport must be enet or steam_sdr and --interval/--provider-timeout/--readiness-max-stale must be positive")
	}
	if *readinessMaxStale < *interval+*providerTimeout {
		fatalf("--readiness-max-stale must be at least --interval plus --provider-timeout")
	}
	if *allocationQuota < 0 || *allocationQuotaWindow <= 0 {
		fatalf("--allocation-quota must be non-negative and --allocation-quota-window must be positive")
	}
	db, err := sql.Open("pgx", *dsn)
	if err != nil {
		fatalf("open PostgreSQL: %v", err)
	}
	defer db.Close()
	startupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if err := db.PingContext(startupCtx); err != nil {
		fatalf("ping PostgreSQL: %v", err)
	}
	if err := migrations.Apply(startupCtx, db, *migrationDir); err != nil {
		fatalf("apply migrations: %v", err)
	}
	if *workloadSecret == "" {
		log.Printf("allocator: warning: --workload-secret / COSMIC_CLASH_WORKLOAD_SECRET is unset; allocated GameServers will receive no cosmic-clash.io/workload-token annotation, and control-plane registration will fail unless a --workload-token-path is separately configured on the supervisor")
	}
	now := func() time.Time { return time.Now().UTC() }
	var budget allocator.AllocationBudget
	if *allocationQuota > 0 {
		budget, err = allocator.NewFixedWindowBudget(*allocationQuota, *allocationQuotaWindow)
		if err != nil {
			fatalf("allocation quota: %v", err)
		}
		log.Printf("allocator: enabled per-replica regional allocation quota=%d window=%s", *allocationQuota, *allocationQuotaWindow)
	}
	metrics := allocator.NewMetrics()
	health := &allocator.Health{}
	providerHTTP, err := agones.NewKubernetesHTTPClient(*agonesURL, *kubernetesTokenPath, *kubernetesCAPath, *providerTimeout)
	if err != nil {
		fatalf("configure Kubernetes API client: %v", err)
	}
	client := agones.Client{BaseURL: *agonesURL, Namespace: *namespace, HTTP: providerHTTP, WorkloadSecret: []byte(*workloadSecret)}
	worker := allocator.Worker{
		Claims: store.AllocatingMatchClaims{DB: db, Transport: *transport},
		Service: allocator.Service{
			Provider: client,
			Durable:  store.AllocationRegistry{DB: db},
			Quota:    store.AllocationQuota{DB: db},
			Budget:   budget,
			Metrics:  metrics,
			Now:      now,
		},
		Now: now,
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	var metricsServer *http.Server
	if *metricsAddr != "" {
		metricsServer = &http.Server{Addr: *metricsAddr, Handler: allocator.RoleHandler(metrics, health, *readinessMaxStale, now), ReadHeaderTimeout: 5 * time.Second}
		go func() {
			if err := metricsServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
				log.Printf("allocator: metrics server: %v", err)
			}
		}()
		log.Printf("allocator: metrics listening on %s", *metricsAddr)
	}
	defer func() {
		if metricsServer != nil {
			shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_ = metricsServer.Shutdown(shutdownCtx)
		}
	}()
	ticker := time.NewTicker(*interval)
	defer ticker.Stop()
	for {
		cycleHealthy := true
		servers, err := client.ListReadyServers(ctx)
		if err != nil && ctx.Err() == nil {
			cycleHealthy = false
			log.Printf("allocator: list Ready GameServers: %v", err)
		} else {
			for _, server := range servers {
				if err := store.RegisterReadyServer(ctx, db, server, now()); err != nil && ctx.Err() == nil {
					cycleHealthy = false
					log.Printf("allocator: register Ready GameServer %s: %v", server.ServerID, err)
				}
			}
		}
		if _, err := worker.RunOnce(ctx); err != nil && ctx.Err() == nil {
			cycleHealthy = false
			log.Printf("allocator: run once: %v", err)
		}
		if cycleHealthy && ctx.Err() == nil {
			health.ObserveSuccessfulCycle(now())
		}
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

func envOrDefault(name, fallback string) string {
	if value := os.Getenv(name); value != "" {
		return value
	}
	return fallback
}

func fatalf(format string, args ...any) {
	log.Printf("allocator: "+format, args...)
	os.Exit(1)
}