CosmicClash/server/store/allocation_quota_sql.go

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
}