CosmicClash/server/domain/allocator.go

package domain

import (
	"crypto/sha256"
	"fmt"
	"sort"
	"sync"
	"time"
)

type ServerLifecycle string

const (
	ServerReady     ServerLifecycle = "READY"
	ServerAllocated ServerLifecycle = "ALLOCATED"
)

type ReadyServer struct {
	ServerID  string
	Region    string
	Build     string
	Protocol  int
	Transport string
	State     ServerLifecycle
}

type AllocationRequest struct {
	AllocationID string
	MatchID      string
	Playlist     Playlist
	Region       string
	Build        string
	Protocol     int
	ArenaPath    string
	Transport    string
}

type Allocation struct {
	AllocationID string
	MatchID      string
	ServerID     string
	Region       string
	Build        string
	Protocol     int
	ArenaPath    string
	Transport    string
	State        ServerLifecycle
	AllocatedAt  time.Time
	// Endpoint is the client-facing address the provider returned. It is
	// persisted so a worker that crashes between allocating and publishing the
	// assignment roster can recover it instead of stranding the match.
	Endpoint string
}

type Allocator struct {
	mu            sync.Mutex
	servers       map[string]ReadyServer
	allocations   map[string]Allocation
	assignments   map[string]Assignment
	requestHashes map[string][32]byte
}

var (
	ErrNoCapacity         = fmt.Errorf("no compatible ready server")
	ErrAllocationInput    = fmt.Errorf("invalid allocation request")
	ErrAllocationNotFound = fmt.Errorf("allocation not found")
)

func NewAllocator(servers []ReadyServer) (*Allocator, error) {
	a := &Allocator{servers: make(map[string]ReadyServer, len(servers)), allocations: make(map[string]Allocation), assignments: make(map[string]Assignment), requestHashes: make(map[string][32]byte)}
	for _, server := range servers {
		if server.ServerID == "" || server.Region == "" || server.Build == "" || server.Protocol <= 0 || (server.Transport != "enet" && server.Transport != "steam_sdr") || server.State != ServerReady {
			return nil, fmt.Errorf("%w: invalid ready server", ErrAllocationInput)
		}
		if _, exists := a.servers[server.ServerID]; exists {
			return nil, fmt.Errorf("%w: duplicate server", ErrAllocationInput)
		}
		a.servers[server.ServerID] = server
	}
	return a, nil
}

// Allocate is the in-process equivalent of a GameServerAllocation. The mutex
// represents the durable allocator transaction; the PostgreSQL/Agones adapter
// must preserve this claim-before-assignment ordering across replicas.
func (a *Allocator) Allocate(request AllocationRequest, now time.Time) (Allocation, error) {
	if err := validateAllocationRequest(request); err != nil {
		return Allocation{}, err
	}
	digest := allocationDigest(request)
	a.mu.Lock()
	defer a.mu.Unlock()
	if prior, ok := a.allocations[request.AllocationID]; ok {
		if a.requestHashes[request.AllocationID] != digest {
			return Allocation{}, ErrConflict
		}
		return prior, nil
	}
	ids := make([]string, 0)
	for id, server := range a.servers {
		if server.State == ServerReady && server.Region == request.Region && server.Build == request.Build && server.Protocol == request.Protocol && server.Transport == request.Transport {
			ids = append(ids, id)
		}
	}
	if len(ids) == 0 {
		return Allocation{}, ErrNoCapacity
	}
	sort.Strings(ids)
	server := a.servers[ids[0]]
	server.State = ServerAllocated
	a.servers[server.ServerID] = server
	allocation := Allocation{AllocationID: request.AllocationID, MatchID: request.MatchID, ServerID: server.ServerID, Region: server.Region, Build: server.Build, Protocol: server.Protocol, ArenaPath: request.ArenaPath, Transport: server.Transport, State: ServerAllocated, AllocatedAt: now}
	a.allocations[request.AllocationID] = allocation
	a.requestHashes[request.AllocationID] = digest
	return allocation, nil
}

// PublishAssignment is the allocation-to-client boundary. It holds the same
// allocator lock as the claim and exposes no assignment until the allocated
// server, complete compatibility tuple, endpoint, and manifest signature all
// verify. The returned assignment is stable across an identical retry.
func (a *Allocator) PublishAssignment(allocationID string, manifest AllocationManifest, endpoint string, signature []byte, verify func([]byte, []byte) bool) (Assignment, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	allocation, ok := a.allocations[allocationID]
	if !ok {
		return Assignment{}, ErrAllocationNotFound
	}
	assignment, err := VerifyAssignment(allocation, manifest, endpoint, signature, verify)
	if err != nil {
		return Assignment{}, err
	}
	if prior, exists := a.assignments[allocationID]; exists {
		if prior != assignment {
			return Assignment{}, ErrConflict
		}
		return prior, nil
	}
	a.assignments[allocationID] = assignment
	return assignment, nil
}

func validateAllocationRequest(request AllocationRequest) error {
	if request.AllocationID == "" || request.MatchID == "" || request.Region == "" || request.Build == "" || request.Protocol <= 0 || (request.Transport != "enet" && request.Transport != "steam_sdr") || (request.ArenaPath != "" && !IsRankedArenaPath(request.ArenaPath)) {
		return ErrAllocationInput
	}
	return nil
}

func allocationDigest(request AllocationRequest) [32]byte {
	return sha256.Sum256([]byte(fmt.Sprintf("%s\x00%s\x00%s\x00%s\x00%d\x00%s\x00%s", request.AllocationID, request.MatchID, request.Region, request.Build, request.Protocol, request.Transport, request.ArenaPath)))
}