CosmicClash/server/cmd/allocator/main.go

package main

import (
	"context"
	"database/sql"
	"flag"
	"log"
	"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")
	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")
	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 {
		fatalf("--transport must be enet or steam_sdr and --interval must be positive")
	}
	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)
	}
	client := agones.Client{BaseURL: *agonesURL, Namespace: *namespace, 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,
			Now:      now,
		},
		Now: now,
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	ticker := time.NewTicker(*interval)
	defer ticker.Stop()
	for {
		servers, err := client.ListReadyServers(ctx)
		if err != nil && ctx.Err() == nil {
			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 {
					log.Printf("allocator: register Ready GameServer %s: %v", server.ServerID, err)
				}
			}
		}
		if _, err := worker.RunOnce(ctx); err != nil && ctx.Err() == nil {
			log.Printf("allocator: run once: %v", err)
		}
		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)
}