CosmicClash/server/store/allocation_match_sql.go

package store

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

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

const AllocationClaimLease = time.Minute

type PendingAllocation struct {
	Request domain.AllocationRequest
}

const ClaimAllocatingMatchSQL = `WITH candidate AS (
    SELECT match_id FROM matches
    WHERE state = 'ALLOCATING' AND server_id IS NULL
      AND (allocation_id IS NULL OR allocation_claimed_at <= $1)
    ORDER BY created_at, match_id
    LIMIT 1
    FOR UPDATE SKIP LOCKED
)
UPDATE matches m
SET allocation_id = 'allocation-' || candidate.match_id, allocation_claimed_at = $2
FROM candidate
WHERE m.match_id = candidate.match_id
RETURNING m.match_id, m.playlist, m.region, m.protocol_version, m.arena_path, m.allocation_id`

const AllocatingMatchBuildSQL = `SELECT client_build
FROM queue_tickets q
JOIN match_participants mp ON mp.ticket_id = q.ticket_id AND mp.player_id = q.player_id
WHERE mp.match_id = $1
ORDER BY q.client_build`

const BindAllocatedMatchParticipantsSQL = `WITH bound AS (
	UPDATE matches
	SET server_id = $3, revision = revision + 1
    WHERE match_id = $1 AND state = 'ALLOCATING' AND allocation_id = $2 AND server_id IS NULL
      AND EXISTS (
          SELECT 1 FROM allocations
          WHERE allocation_id = $2 AND match_id = $1 AND server_id = $3 AND state = 'ALLOCATED'
      )
    RETURNING match_id, revision
), participants AS (
    SELECT mp.ticket_id, mp.player_id
    FROM match_participants mp
    JOIN bound ON bound.match_id = mp.match_id
), advanced AS (
    UPDATE queue_tickets q
    SET state = 'ALLOCATING', revision = revision + 1
    FROM participants p
    WHERE q.ticket_id = p.ticket_id AND q.player_id = p.player_id AND q.state = 'ACCEPTED'
    RETURNING q.ticket_id
)
	SELECT (SELECT count(*) FROM participants), (SELECT count(*) FROM advanced), COALESCE((SELECT revision FROM bound), -1)`

const ReleaseAllocatedMatchClaimSQL = `UPDATE matches
SET allocation_id = NULL, allocation_claimed_at = NULL
WHERE match_id = $1 AND state = 'ALLOCATING' AND allocation_id = $2 AND server_id IS NULL`

const AdvanceServerRegistrationSQL = `WITH matched AS (
    UPDATE matches
    SET state = $4,
        initial_connect_ready_at = CASE WHEN $4 = 'ASSIGNMENT_READY' THEN $6 ELSE initial_connect_ready_at END,
        revision = revision + 1
    WHERE match_id = $1 AND server_id = $2 AND state = $3 AND protocol_version = $7
      AND EXISTS (SELECT 1 FROM allocations WHERE match_id = $1 AND server_id = $2 AND allocation_id = $5 AND protocol_version = $7 AND state = 'ALLOCATED')
      AND ($4 <> 'ASSIGNMENT_READY' OR (SELECT count(*) FROM assignments WHERE match_id = $1 AND expires_at > $6) = (SELECT count(*) FROM match_participants WHERE match_id = $1))
    RETURNING match_id
), advanced AS (
    UPDATE queue_tickets q
    SET state = $4, revision = revision + 1
    FROM match_participants mp JOIN matched m ON m.match_id = mp.match_id
    WHERE q.ticket_id = mp.ticket_id AND q.player_id = mp.player_id AND q.state = $3
    RETURNING q.ticket_id
)
SELECT (SELECT count(*) FROM matched), (SELECT count(*) FROM match_participants WHERE match_id = $1), (SELECT count(*) FROM advanced), COALESCE((SELECT revision FROM matched), -1)`

const serverRegistrationParticipantIDsSQL = `SELECT player_id FROM match_participants WHERE match_id = $1 ORDER BY player_id`

const serverRegistrationOutboxSQL = `INSERT INTO outbox
    (event_id, aggregate_type, aggregate_id, revision, event_type, payload)
VALUES ($1, 'match', $2, $3, 'state_changed', $4)`

const ServerRegistrationIdempotencyScope = "server.register"

const ServerRegistrationIdempotencyInsertSQL = `INSERT INTO idempotency_keys
    (scope, idempotency_key, payload_digest, result)
VALUES ($1, $2, $3, '{}')
ON CONFLICT (scope, idempotency_key) DO NOTHING`

const ServerRegistrationIdempotencySelectSQL = `SELECT payload_digest
FROM idempotency_keys
WHERE scope = $1 AND idempotency_key = $2
FOR UPDATE`

func AdvanceServerRegistration(ctx context.Context, db *sql.DB, binding domain.WorkloadBinding, protocol int, assignmentReady bool, idempotencyKey string, now time.Time) error {
	if db == nil || binding.MatchID == "" || binding.ServerID == "" || binding.AllocationID == "" || protocol < 1 || len(idempotencyKey) < 16 || len(idempotencyKey) > 128 || now.IsZero() {
		return fmt.Errorf("invalid server registration")
	}
	from, to := domain.Allocating, domain.ProcessReady
	if assignmentReady {
		from, to = domain.ProcessReady, domain.AssignmentReady
	}
	digest := sha256.Sum256([]byte(fmt.Sprintf("%s\x00%s\x00%s\x00%d\x00%t", binding.AllocationID, binding.MatchID, binding.ServerID, protocol, assignmentReady)))
	return RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		inserted, err := tx.ExecContext(ctx, ServerRegistrationIdempotencyInsertSQL, ServerRegistrationIdempotencyScope, idempotencyKey, digest[:])
		if err != nil {
			return err
		}
		changed, err := inserted.RowsAffected()
		if err != nil {
			return err
		}
		if changed == 0 {
			var prior []byte
			if err := tx.QueryRowContext(ctx, ServerRegistrationIdempotencySelectSQL, ServerRegistrationIdempotencyScope, idempotencyKey).Scan(&prior); err != nil {
				return err
			}
			if !bytes.Equal(prior, digest[:]) {
				return domain.ErrConflict
			}
			return nil
		}
		var matched, participants, advanced int
		var revision int64
		if err := tx.QueryRowContext(ctx, AdvanceServerRegistrationSQL, binding.MatchID, binding.ServerID, from, to, binding.AllocationID, now, protocol).Scan(&matched, &participants, &advanced, &revision); err != nil {
			return err
		}
		if matched != 1 || participants == 0 || advanced != participants {
			return domain.ErrConflict
		}
		rows, err := tx.QueryContext(ctx, serverRegistrationParticipantIDsSQL, binding.MatchID)
		if err != nil {
			return err
		}
		playerIDs := make([]string, 0, participants)
		for rows.Next() {
			var playerID string
			if err := rows.Scan(&playerID); err != nil {
				rows.Close()
				return err
			}
			playerIDs = append(playerIDs, playerID)
		}
		if err := rows.Err(); err != nil {
			rows.Close()
			return err
		}
		if err := rows.Close(); err != nil {
			return err
		}
		payload, err := json.Marshal(map[string]any{
			"event": "state_changed", "revision": revision, "resource_id": binding.MatchID,
			"occurred_at": now, "state": string(to), "match_id": binding.MatchID, "player_ids": playerIDs,
		})
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, serverRegistrationOutboxSQL, fmt.Sprintf("match:%s:%d", binding.MatchID, revision), binding.MatchID, revision, payload)
		return err
	})
}

// FindProviderAllocation verifies whether a recovered lease has already
// crossed the durable provider boundary. A worker can then bind it without
// issuing a second external allocation request after a crash.
func FindProviderAllocation(ctx context.Context, db *sql.DB, request domain.AllocationRequest) (domain.Allocation, bool, error) {
	if db == nil || request.AllocationID == "" || request.MatchID == "" || request.Region == "" || request.Build == "" || request.Protocol <= 0 || (request.Transport != "enet" && request.Transport != "steam_sdr") {
		return domain.Allocation{}, false, fmt.Errorf("invalid provider allocation lookup")
	}
	var allocation domain.Allocation
	var digest []byte
	err := db.QueryRowContext(ctx, SelectAllocationSQL, request.AllocationID).Scan(&allocation.AllocationID, &allocation.MatchID, &allocation.ServerID, &allocation.Region, &allocation.Build, &allocation.Protocol, &allocation.ArenaPath, &allocation.Transport, &allocation.AllocatedAt, &digest)
	if err == sql.ErrNoRows {
		return domain.Allocation{}, false, nil
	}
	if err != nil {
		return domain.Allocation{}, false, err
	}
	want := allocationRequestDigest(request)
	if !bytes.Equal(digest, want[:]) || allocation.MatchID != request.MatchID || allocation.Region != request.Region || allocation.Build != request.Build || allocation.Protocol != request.Protocol || allocation.ArenaPath != request.ArenaPath || allocation.Transport != request.Transport {
		return domain.Allocation{}, false, domain.ErrConflict
	}
	allocation.State = domain.ServerAllocated
	return allocation, true, nil
}

// ClaimAllocatingMatch returns one durable provider work item. The fixed
// allocation ID is retained across a lease recovery, allowing every later
// reconciliation step to reject a different server for the same match.
func ClaimAllocatingMatch(ctx context.Context, db *sql.DB, transport string, now time.Time) (PendingAllocation, bool, error) {
	if db == nil || (transport != "enet" && transport != "steam_sdr") || now.IsZero() {
		return PendingAllocation{}, false, fmt.Errorf("invalid allocation claim arguments")
	}
	var item PendingAllocation
	found := false
	err := RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		var matchID, playlist, region string
		var protocol int
		var arenaPath sql.NullString
		var claimedID string
		err := tx.QueryRowContext(ctx, ClaimAllocatingMatchSQL, now.Add(-AllocationClaimLease), now).Scan(&matchID, &playlist, &region, &protocol, &arenaPath, &claimedID)
		if err == sql.ErrNoRows {
			return nil
		}
		if err != nil {
			return err
		}
		rows, err := tx.QueryContext(ctx, AllocatingMatchBuildSQL, matchID)
		if err != nil {
			return err
		}
		defer rows.Close()
		build := ""
		for rows.Next() {
			var candidate string
			if err := rows.Scan(&candidate); err != nil {
				return err
			}
			if build == "" {
				build = candidate
			} else if build != candidate {
				return fmt.Errorf("allocating match has mixed client builds")
			}
		}
		if err := rows.Err(); err != nil {
			return err
		}
		if build == "" {
			return fmt.Errorf("allocating match has no participants")
		}
		if domain.Playlist(playlist) == domain.Ranked && (!arenaPath.Valid || !domain.IsRankedArenaPath(arenaPath.String)) {
			return fmt.Errorf("ranked allocating match has invalid arena")
		}
		item.Request = domain.AllocationRequest{AllocationID: claimedID, MatchID: matchID, Playlist: domain.Playlist(playlist), Region: region, Build: build, Protocol: protocol, ArenaPath: arenaPath.String, Transport: transport}
		found = true
		return nil
	})
	return item, found, err
}

func BindAllocatedMatch(ctx context.Context, db *sql.DB, allocation domain.Allocation) error {
	if db == nil || allocation.MatchID == "" || allocation.AllocationID == "" || allocation.ServerID == "" || allocation.State != domain.ServerAllocated || allocation.AllocatedAt.IsZero() {
		return fmt.Errorf("invalid allocated match binding")
	}
	return RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		var participants, advanced int
		var revision int64
		if err := tx.QueryRowContext(ctx, BindAllocatedMatchParticipantsSQL, allocation.MatchID, allocation.AllocationID, allocation.ServerID).Scan(&participants, &advanced, &revision); err != nil {
			return err
		}
		if participants == 0 || participants != advanced {
			return domain.ErrConflict
		}
		rows, err := tx.QueryContext(ctx, serverRegistrationParticipantIDsSQL, allocation.MatchID)
		if err != nil {
			return err
		}
		playerIDs := make([]string, 0, participants)
		for rows.Next() {
			var playerID string
			if err := rows.Scan(&playerID); err != nil {
				rows.Close()
				return err
			}
			playerIDs = append(playerIDs, playerID)
		}
		if err := rows.Err(); err != nil {
			rows.Close()
			return err
		}
		if err := rows.Close(); err != nil {
			return err
		}
		payload, err := json.Marshal(map[string]any{
			"event": "state_changed", "revision": revision, "resource_id": allocation.MatchID,
			"occurred_at": allocation.AllocatedAt, "state": string(domain.Allocating), "match_id": allocation.MatchID, "player_ids": playerIDs,
		})
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, serverRegistrationOutboxSQL, fmt.Sprintf("match:%s:%d", allocation.MatchID, revision), allocation.MatchID, revision, payload)
		return err
	})
}

func ReleaseAllocatedMatchClaim(ctx context.Context, db *sql.DB, matchID, allocationID string) error {
	if db == nil || matchID == "" || allocationID == "" {
		return fmt.Errorf("invalid allocated match claim release")
	}
	result, err := db.ExecContext(ctx, ReleaseAllocatedMatchClaimSQL, matchID, allocationID)
	if err != nil {
		return err
	}
	changed, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if changed != 1 {
		return domain.ErrConflict
	}
	return nil
}