feat(multiplayer): add shared allocation quota

server/store/allocation_quota_sql.go

@@ -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
+}