CosmicClash/server/allocator/worker.go

package allocator

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"github.com/cosmic-clash/cosmic-clash/server/agones"
	"github.com/cosmic-clash/cosmic-clash/server/domain"
)

// MatchClaimSource is the durable allocator work queue. Implementations must
// lease a match before returning it and fence binding by allocation ID.
type MatchClaimSource interface {
	ClaimAllocatingMatch(context.Context, time.Time) (domain.AllocationRequest, bool, error)
	FindProviderAllocation(context.Context, domain.AllocationRequest) (domain.Allocation, bool, error)
	BindAllocatedMatch(context.Context, domain.Allocation) error
}

// Worker consumes one leased match at a time. Provider failures deliberately
// retain the lease: an HTTP/provider failure can be ambiguous after an external
// allocation, so releasing it could allocate two GameServers for one match.
type Worker struct {
	Claims  MatchClaimSource
	Service Service
	Now     func() time.Time
	// Roster and Keys wire the assignment hand-off. Without them the worker
	// binds an allocation and stops, nothing ever writes the assignments
	// table, and the allocated supervisor's roster fetch fails -- so every
	// real allocation dies before the game process launches. They are optional
	// only so existing allocation-only tests need no key material.
	Roster AssignmentRosterSource
	Keys   JoinSigningKeys
}

// RunOnce returns whether it found a claimed match. It never exposes an
// endpoint itself; Service first records the provider allocation durably and
// BindAllocatedMatch then attaches that already-recorded allocation to the
// fenced match claim.
func (w Worker) RunOnce(ctx context.Context) (bool, error) {
	if w.Claims == nil || w.Now == nil {
		return false, errNotConfigured
	}
	request, found, err := w.Claims.ClaimAllocatingMatch(ctx, w.Now())
	if err != nil || !found {
		return found, err
	}
	allocation, recorded, err := w.Claims.FindProviderAllocation(ctx, request)
	if err != nil {
		return true, fmt.Errorf("recover allocation for match %s: %w", request.MatchID, err)
	}
	if !recorded {
		if recoverer, ok := w.Service.Provider.(ProviderRecoverer); ok {
			recovered, found, err := recoverer.RecoverAllocation(ctx, request, w.Now())
			if err != nil {
				return true, fmt.Errorf("recover provider allocation for match %s: %w", request.MatchID, err)
			}
			if found {
				if err := validateProviderAllocation(request, recovered); err != nil {
					return true, fmt.Errorf("recovered provider allocation for match %s: %w", request.MatchID, err)
				}
				recorded, err := w.Service.RecordProviderAllocation(ctx, recovered, w.Now())
				if err != nil {
					return true, fmt.Errorf("record recovered allocation for match %s: %w", request.MatchID, err)
				}
				allocation = recorded
			} else {
				result, err := w.Service.Allocate(ctx, request, AllocationLabels(request))
				if err != nil {
					return true, fmt.Errorf("allocate claimed match %s: %w", request.MatchID, err)
				}
				allocation = result.Allocation
			}
		} else {
			result, err := w.Service.Allocate(ctx, request, AllocationLabels(request))
			if err != nil {
				return true, fmt.Errorf("allocate claimed match %s: %w", request.MatchID, err)
			}
			allocation = result.Allocation
		}
	}
	if err := w.Claims.BindAllocatedMatch(ctx, allocation); err != nil {
		return true, fmt.Errorf("bind allocated match %s: %w", request.MatchID, err)
	}
	if err := w.publishAssignmentRoster(ctx, allocation); err != nil {
		return true, fmt.Errorf("publish assignment roster for match %s: %w", request.MatchID, err)
	}
	return true, nil
}

// publishAssignmentRoster completes the hand-off from allocation to a joinable
// match. The supervisor fetches a non-empty roster before it launches the game
// child, so skipping this leaves the match stuck short of ASSIGNMENT_READY
// forever.
//
// It is safe to retry: SaveVerifiedAssignmentRoster upserts by (match, player)
// and re-validates every claim against the durable participants, so a worker
// that crashed after binding but before publishing simply republishes on the
// next pass.
func (w Worker) publishAssignmentRoster(ctx context.Context, allocation domain.Allocation) error {
	if w.Roster == nil {
		// Allocation-only deployments (and the allocation-focused tests) leave
		// this unset deliberately.
		return nil
	}
	if err := w.Keys.validate(); err != nil {
		return err
	}
	participants, err := w.Roster.LoadAssignmentParticipants(ctx, allocation)
	if err != nil {
		return err
	}
	assignment, roster, err := BuildSignedRoster(allocation, participants, w.Keys, w.Now())
	if err != nil {
		return err
	}
	return w.Service.PublishRoster(ctx, assignment, roster, domain.VerifyJoinAuthorisationHMAC(w.Keys.Keys))
}

func validateProviderAllocation(request domain.AllocationRequest, result agones.AllocatedServer) error {
	allocation := result.Allocation
	if result.Endpoint == "" || allocation.State != domain.ServerAllocated || allocation.AllocationID != request.AllocationID || allocation.MatchID != request.MatchID || allocation.ServerID == "" || allocation.Region != request.Region || allocation.Build != request.Build || allocation.Protocol != request.Protocol || allocation.Transport != request.Transport || allocation.ArenaPath != request.ArenaPath {
		return fmt.Errorf("provider allocation does not match request")
	}
	return nil
}

// AllocationLabels are the compatibility selectors shared with the Fleet
// template. They are derived only from the durable match plan, never client
// input or mutable worker configuration.
func AllocationLabels(request domain.AllocationRequest) map[string]string {
	labels := map[string]string{
		"cosmic-clash.io/region":    request.Region,
		"cosmic-clash.io/build":     request.Build,
		"cosmic-clash.io/protocol":  strconv.Itoa(request.Protocol),
		"cosmic-clash.io/transport": request.Transport,
	}
	if request.Playlist != "" {
		labels["cosmic-clash.io/playlist"] = string(request.Playlist)
	}
	return labels
}