CosmicClash/server/allocator/budget.go

package allocator

import (
	"fmt"
	"sync"
	"time"
)

// ErrAllocationBudgetExceeded is deliberately generic: callers should not
// learn quota internals, and the allocator can safely retry the leased match
// after the current window expires.
var ErrAllocationBudgetExceeded = fmt.Errorf("allocation budget exceeded")

// FixedWindowBudget is a process-local denial-of-wallet guard. It limits the
// number of provider allocation attempts per region in a time window. The
// production deployment must use the same policy behind a shared durable
// counter for a global quota; this type prevents one allocator replica from
// spending without bound and is useful in tests and single-replica setups.
type FixedWindowBudget struct {
	mu          sync.Mutex
	limit       int
	window      time.Duration
	windowStart time.Time
	counts      map[string]int
}

func NewFixedWindowBudget(limit int, window time.Duration) (*FixedWindowBudget, error) {
	if limit < 1 || window <= 0 {
		return nil, fmt.Errorf("invalid allocation budget")
	}
	return &FixedWindowBudget{limit: limit, window: window, counts: make(map[string]int)}, nil
}

func (b *FixedWindowBudget) Allow(region string, now time.Time) error {
	if b == nil || (region != "EU" && region != "NA") || now.IsZero() {
		return fmt.Errorf("invalid allocation budget request")
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.windowStart.IsZero() || !now.Before(b.windowStart.Add(b.window)) {
		b.windowStart = now
		b.counts = make(map[string]int)
	}
	if b.counts[region] >= b.limit {
		return ErrAllocationBudgetExceeded
	}
	b.counts[region]++
	return nil
}