mirror of
https://github.com/jcreek/CosmicClash.git
synced 2026-09-13 12:02:02 +00:00
feat(multiplayer): add shared allocation quota
This commit is contained in:
@@ -0,0 +1,78 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
var ErrAllocationQuotaExceeded = errors.New("allocation quota exceeded")
|
||||
|
||||
type AllocationQuota struct {
|
||||
DB *sql.DB
|
||||
}
|
||||
|
||||
const allocationQuotaSelectSQL = `SELECT window_started_at, window_seconds, used_allocations, max_allocations
|
||||
FROM allocation_quotas WHERE region = $1 FOR UPDATE`
|
||||
|
||||
const allocationQuotaResetSQL = `UPDATE allocation_quotas
|
||||
SET window_started_at = $2, used_allocations = 1, updated_at = $2 WHERE region = $1`
|
||||
|
||||
const allocationQuotaIncrementSQL = `UPDATE allocation_quotas
|
||||
SET used_allocations = used_allocations + 1, updated_at = $2 WHERE region = $1`
|
||||
|
||||
const SetAllocationQuotaSQL = `INSERT INTO allocation_quotas
|
||||
(region, window_started_at, window_seconds, used_allocations, max_allocations, updated_at)
|
||||
VALUES ($1, $2, $3, 0, $4, $2)
|
||||
ON CONFLICT (region) DO UPDATE SET window_started_at = EXCLUDED.window_started_at,
|
||||
window_seconds = EXCLUDED.window_seconds, used_allocations = 0,
|
||||
max_allocations = EXCLUDED.max_allocations, updated_at = EXCLUDED.updated_at`
|
||||
|
||||
// SetAllocationQuota configures the optional shared regional quota. It is
|
||||
// intended for operator provisioning, not for a request path.
|
||||
func SetAllocationQuota(ctx context.Context, db *sql.DB, region string, maxAllocations int, window time.Duration, now time.Time) error {
|
||||
if db == nil || (region != "EU" && region != "NA") || maxAllocations < 1 || window <= 0 || window > 365*24*time.Hour || now.IsZero() {
|
||||
return fmt.Errorf("invalid allocation quota")
|
||||
}
|
||||
seconds := int(window / time.Second)
|
||||
if seconds < 1 {
|
||||
return fmt.Errorf("allocation quota window is too small")
|
||||
}
|
||||
_, err := db.ExecContext(ctx, SetAllocationQuotaSQL, region, now, seconds, maxAllocations)
|
||||
return err
|
||||
}
|
||||
|
||||
func (q AllocationQuota) Consume(ctx context.Context, region string, now time.Time) error {
|
||||
if q.DB == nil || (region != "EU" && region != "NA") || now.IsZero() {
|
||||
return fmt.Errorf("invalid allocation quota request")
|
||||
}
|
||||
return RunSerializable(ctx, q.DB, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
|
||||
return consumeAllocationQuotaTx(ctx, tx, region, now)
|
||||
})
|
||||
}
|
||||
|
||||
// consumeAllocationQuotaTx consumes one unit when a quota row exists. The
|
||||
// caller must already be inside the serializable allocation transaction; the
|
||||
// row lock makes this global across allocator replicas sharing PostgreSQL.
|
||||
func consumeAllocationQuotaTx(ctx context.Context, tx *sql.Tx, region string, now time.Time) error {
|
||||
var started time.Time
|
||||
var seconds, used, maximum int
|
||||
err := tx.QueryRowContext(ctx, allocationQuotaSelectSQL, region).Scan(&started, &seconds, &used, &maximum)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !now.Before(started.Add(time.Duration(seconds) * time.Second)) {
|
||||
_, err = tx.ExecContext(ctx, allocationQuotaResetSQL, region, now)
|
||||
return err
|
||||
}
|
||||
if used >= maximum {
|
||||
return ErrAllocationQuotaExceeded
|
||||
}
|
||||
_, err = tx.ExecContext(ctx, allocationQuotaIncrementSQL, region, now)
|
||||
return err
|
||||
}
|
||||
@@ -75,6 +75,9 @@ func ClaimAllocation(ctx context.Context, db *sql.DB, request domain.AllocationR
|
||||
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 {
|
||||
|
||||
@@ -21,6 +21,29 @@ func TestAllocatorSQLClaimsAndAuditsCompatibleReadyServers(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, fragment := range []string{"allocation_quotas", "ON CONFLICT (region)", "used_allocations"} {
|
||||
if !contains(SetAllocationQuotaSQL, fragment) {
|
||||
t.Fatalf("quota query missing %q", fragment)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAllocationQuotaRejectsInvalidArgumentsWithoutDatabase(t *testing.T) {
|
||||
if err := SetAllocationQuota(nil, nil, "EU", 1, time.Minute, time.Unix(1000, 0)); err == nil {
|
||||
t.Fatal("nil database accepted")
|
||||
}
|
||||
if err := SetAllocationQuota(nil, nil, "APAC", 1, time.Minute, time.Unix(1000, 0)); err == nil {
|
||||
t.Fatal("unknown region accepted")
|
||||
}
|
||||
if err := SetAllocationQuota(nil, nil, "EU", 0, time.Minute, time.Unix(1000, 0)); err == nil {
|
||||
t.Fatal("zero limit accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllocationQuotaConsumeRejectsInvalidArgumentsWithoutDatabase(t *testing.T) {
|
||||
if err := (AllocationQuota{}).Consume(nil, "EU", time.Unix(1000, 0)); err == nil {
|
||||
t.Fatal("nil database accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimAllocationRejectsInvalidRequestsWithoutDatabase(t *testing.T) {
|
||||
|
||||
@@ -44,7 +44,7 @@ func openIntegrationPostgres(t *testing.T) *sql.DB {
|
||||
|
||||
func applyIntegrationMigrations(t *testing.T, db *sql.DB) {
|
||||
t.Helper()
|
||||
if _, err := db.ExecContext(context.Background(), `DROP TABLE IF EXISTS schema_migrations, assignments, audit_events, outbox, result_receipts, ranked_season_rollovers, penalties, seasons, ratings, match_participants, matches, proposal_participants, proposals, allocations, game_servers, queue_tickets, idempotency_keys, sessions, identities CASCADE`); err != nil {
|
||||
if _, err := db.ExecContext(context.Background(), `DROP TABLE IF EXISTS schema_migrations, allocation_quotas, assignments, audit_events, outbox, result_receipts, ranked_season_rollovers, penalties, seasons, ratings, match_participants, matches, proposal_participants, proposals, allocations, game_servers, queue_tickets, idempotency_keys, sessions, identities CASCADE`); err != nil {
|
||||
t.Fatalf("reset PostgreSQL schema: %v", err)
|
||||
}
|
||||
if err := migrations.Apply(context.Background(), db, filepath.Join("..", "migrations")); err != nil {
|
||||
@@ -104,6 +104,41 @@ func TestPostgreSQLAllocatorClaimReplayAndCapacityFence(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgreSQLSharedAllocationQuotaFencesClaimsAndReplays(t *testing.T) {
|
||||
db := openIntegrationPostgres(t)
|
||||
applyIntegrationMigrations(t, db)
|
||||
now := time.Now().UTC().Truncate(time.Microsecond)
|
||||
ctx := context.Background()
|
||||
if err := SetAllocationQuota(ctx, db, "EU", 1, time.Minute, now); err != nil {
|
||||
t.Fatalf("set quota: %v", err)
|
||||
}
|
||||
for _, server := range []domain.ReadyServer{
|
||||
{ServerID: "quota-server-a", Region: "EU", Build: "build-1", Protocol: 1, Transport: "enet", State: domain.ServerReady},
|
||||
{ServerID: "quota-server-b", Region: "EU", Build: "build-1", Protocol: 1, Transport: "enet", State: domain.ServerReady},
|
||||
} {
|
||||
if err := RegisterReadyServer(ctx, db, server, now); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
first := domain.AllocationRequest{AllocationID: "quota-allocation-1", MatchID: "quota-match-1", Region: "EU", Build: "build-1", Protocol: 1, Transport: "enet"}
|
||||
if _, err := ClaimAllocation(ctx, db, first, now); err != nil {
|
||||
t.Fatalf("first claim: %v", err)
|
||||
}
|
||||
if _, err := ClaimAllocation(ctx, db, first, now.Add(time.Second)); err != nil {
|
||||
t.Fatalf("idempotent replay was fenced: %v", err)
|
||||
}
|
||||
second := domain.AllocationRequest{AllocationID: "quota-allocation-2", MatchID: "quota-match-2", Region: "EU", Build: "build-1", Protocol: 1, Transport: "enet"}
|
||||
if _, err := ClaimAllocation(ctx, db, second, now.Add(2*time.Second)); !errors.Is(err, ErrAllocationQuotaExceeded) {
|
||||
t.Fatalf("second claim err=%v, want shared quota fence", err)
|
||||
}
|
||||
if err := SetAllocationQuota(ctx, db, "EU", 1, time.Minute, now.Add(time.Minute)); err != nil {
|
||||
t.Fatalf("reset quota: %v", err)
|
||||
}
|
||||
if _, err := ClaimAllocation(ctx, db, second, now.Add(time.Minute)); err != nil {
|
||||
t.Fatalf("claim after quota window reset: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPostgreSQLConcurrentAllocationClaimNeverDoubleBooksAReadyServer is the
|
||||
// live counterpart to TestPostgreSQLAllocatorClaimReplayAndCapacityFence: that
|
||||
// test claims strictly one request at a time, so it cannot show what happens
|
||||
@@ -1249,8 +1284,8 @@ func TestPostgreSQLMigrationsRollBackAndReapplyCleanly(t *testing.T) {
|
||||
// Roll back every migration one at a time, in reverse, checking each
|
||||
// down file actually undoes what its forward file created — not just
|
||||
// that Rollback returns nil.
|
||||
if err := migrations.Rollback(context.Background(), db, dir, 1); err != nil {
|
||||
t.Fatalf("rollback 0006: %v", err)
|
||||
if err := migrations.Rollback(context.Background(), db, dir, 2); err != nil {
|
||||
t.Fatalf("rollback 0007 and 0006: %v", err)
|
||||
}
|
||||
var hasAllocationClaimColumn bool
|
||||
if err := db.QueryRow(`SELECT count(*) > 0 FROM information_schema.columns WHERE table_name = 'matches' AND column_name = 'allocation_id'`).Scan(&hasAllocationClaimColumn); err != nil {
|
||||
|
||||
Reference in New Issue
Block a user