CosmicClash/server/store/allocator_sql.go

package store

import (
	"bytes"
	"context"
	"crypto/sha256"
	"database/sql"
	"fmt"
	"time"

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

const RegisterReadyServerSQL = `INSERT INTO game_servers
    (server_id, region, build, protocol_version, transport, state, updated_at)
VALUES ($1, $2, $3, $4, $5, 'READY', $6)
ON CONFLICT (server_id) DO UPDATE SET region = EXCLUDED.region,
    build = EXCLUDED.build, protocol_version = EXCLUDED.protocol_version,
    transport = EXCLUDED.transport, updated_at = EXCLUDED.updated_at
WHERE game_servers.state = 'READY'`

const ClaimReadyServerSQL = `UPDATE game_servers SET state = 'ALLOCATED', updated_at = $5
WHERE server_id = (
    SELECT server_id FROM game_servers
    WHERE state = 'READY' AND region = $1 AND build = $2
      AND protocol_version = $3 AND transport = $4
    ORDER BY server_id
    LIMIT 1
    FOR UPDATE SKIP LOCKED
)
RETURNING server_id`

const InsertAllocationSQL = `INSERT INTO allocations
    (allocation_id, match_id, server_id, region, build, protocol_version, arena_path, transport, request_digest, state, allocated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 'ALLOCATED', $10)`

const SelectAllocationSQL = `SELECT allocation_id, match_id, server_id, region, build,
    protocol_version, arena_path, transport, allocated_at, request_digest
FROM allocations WHERE allocation_id = $1`

const ProviderServerClaimSQL = `UPDATE game_servers SET state = 'ALLOCATED', updated_at = $6
WHERE server_id = $1 AND state = 'READY' AND region = $2 AND build = $3
  AND protocol_version = $4 AND transport = $5
RETURNING server_id`

const ServerAllocationConflictSQL = `SELECT allocation_id FROM allocations
WHERE server_id = $1 FOR UPDATE`

func RegisterReadyServer(ctx context.Context, db *sql.DB, server domain.ReadyServer, now time.Time) error {
	if db == nil || server.ServerID == "" || (server.Region != "EU" && server.Region != "NA") || server.Build == "" || server.Protocol <= 0 || (server.Transport != "enet" && server.Transport != "steam_sdr") || server.State != domain.ServerReady || now.IsZero() {
		return fmt.Errorf("invalid ready server registration")
	}
	_, err := db.ExecContext(ctx, RegisterReadyServerSQL, server.ServerID, server.Region, server.Build, server.Protocol, server.Transport, now)
	return err
}

func ClaimAllocation(ctx context.Context, db *sql.DB, request domain.AllocationRequest, now time.Time) (domain.Allocation, error) {
	if !validAllocationInput(db, request, now) {
		return domain.Allocation{}, domain.ErrAllocationInput
	}
	digest := allocationRequestDigest(request)
	var allocation domain.Allocation
	err := RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		var prior domain.Allocation
		var priorDigest []byte
		err := tx.QueryRowContext(ctx, SelectAllocationSQL, request.AllocationID).Scan(&prior.AllocationID, &prior.MatchID, &prior.ServerID, &prior.Region, &prior.Build, &prior.Protocol, &prior.ArenaPath, &prior.Transport, &prior.AllocatedAt, &priorDigest)
		if err == nil {
			if !bytes.Equal(priorDigest, digest[:]) {
				return domain.ErrConflict
			}
			allocation = prior
			allocation.State = domain.ServerAllocated
			return nil
		}
		if err != sql.ErrNoRows {
			return err
		}
		if err := consumeAllocationQuotaTx(ctx, tx, request.Region, now); err != nil {
			return err
		}
		var serverID string
		if err := tx.QueryRowContext(ctx, ClaimReadyServerSQL, request.Region, request.Build, request.Protocol, request.Transport, now).Scan(&serverID); err != nil {
			if err == sql.ErrNoRows {
				return domain.ErrNoCapacity
			}
			return err
		}
		allocation = domain.Allocation{AllocationID: request.AllocationID, MatchID: request.MatchID, ServerID: serverID, Region: request.Region, Build: request.Build, Protocol: request.Protocol, ArenaPath: request.ArenaPath, Transport: request.Transport, State: domain.ServerAllocated, AllocatedAt: now}
		_, err = tx.ExecContext(ctx, InsertAllocationSQL, request.AllocationID, request.MatchID, serverID, request.Region, request.Build, request.Protocol, request.ArenaPath, request.Transport, digest[:], now)
		return err
	})
	return allocation, err
}

// RecordProviderAllocation reconciles a provider-side Agones claim with the
// durable registry. It is deliberately separate from ClaimAllocation because
// Agones has already selected the server; no client-facing assignment may use
// the result until this exact tuple is durably recorded.
func RecordProviderAllocation(ctx context.Context, db *sql.DB, allocation domain.Allocation, now time.Time) (domain.Allocation, error) {
	request := domain.AllocationRequest{AllocationID: allocation.AllocationID, MatchID: allocation.MatchID, Region: allocation.Region, Build: allocation.Build, Protocol: allocation.Protocol, ArenaPath: allocation.ArenaPath, Transport: allocation.Transport}
	if !validAllocationInput(db, request, now) || allocation.State != domain.ServerAllocated || allocation.ServerID == "" {
		return domain.Allocation{}, domain.ErrAllocationInput
	}
	digest := allocationRequestDigest(request)
	var recorded domain.Allocation
	err := RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		var prior domain.Allocation
		var priorDigest []byte
		err := tx.QueryRowContext(ctx, SelectAllocationSQL, allocation.AllocationID).Scan(&prior.AllocationID, &prior.MatchID, &prior.ServerID, &prior.Region, &prior.Build, &prior.Protocol, &prior.ArenaPath, &prior.Transport, &prior.AllocatedAt, &priorDigest)
		if err == nil {
			if !bytes.Equal(priorDigest, digest[:]) || prior.ServerID != allocation.ServerID {
				return domain.ErrConflict
			}
			recorded = prior
			recorded.State = domain.ServerAllocated
			return nil
		}
		if err != sql.ErrNoRows {
			return err
		}
		var existing string
		if err := tx.QueryRowContext(ctx, ServerAllocationConflictSQL, allocation.ServerID).Scan(&existing); err == nil {
			return domain.ErrConflict
		} else if err != sql.ErrNoRows {
			return err
		}
		var serverID string
		if err := tx.QueryRowContext(ctx, ProviderServerClaimSQL, allocation.ServerID, allocation.Region, allocation.Build, allocation.Protocol, allocation.Transport, now).Scan(&serverID); err != nil {
			if err == sql.ErrNoRows {
				return domain.ErrNoCapacity
			}
			return err
		}
		recorded = allocation
		recorded.AllocatedAt = now
		_, err = tx.ExecContext(ctx, InsertAllocationSQL, allocation.AllocationID, allocation.MatchID, serverID, allocation.Region, allocation.Build, allocation.Protocol, allocation.ArenaPath, allocation.Transport, digest[:], now)
		return err
	})
	return recorded, err
}

func validAllocationInput(db *sql.DB, request domain.AllocationRequest, now time.Time) bool {
	return db != nil && request.AllocationID != "" && request.MatchID != "" && (request.Region == "EU" || request.Region == "NA") && request.Build != "" && request.Protocol > 0 && (request.Transport == "enet" || request.Transport == "steam_sdr") && (request.ArenaPath == "" || domain.IsRankedArenaPath(request.ArenaPath)) && !now.IsZero()
}

func allocationRequestDigest(request domain.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)))
}