egress-proxies-operator/internal/lease/store.go

// Package lease implements the in-memory lease store behind the discovery
// API: TTL-based proxy assignment with server-side usage tracking and
// per-(proxy, target) cooldowns. Accepted prototype limitation, documented
// in the README: state is per-process, so an operator restart drops all
// leases and cooldowns — clients must tolerate a lease vanishing (their
// requests still work; they just re-lease).
package lease

import (
	"cmp"
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"slices"
	"strings"
	"sync"
	"time"
)

// Result is a client's report of how a leased proxy behaved against a
// target. ResultRateLimited and ResultBanned record a cooldown; ResultOK is
// an acknowledgement and records nothing.
type Result string

const (
	ResultOK          Result = "ok"
	ResultRateLimited Result = "rate_limited"
	ResultBanned      Result = "banned"
)

// ParseResult maps a wire value to a Result; ok is false for anything
// unknown, which the API layer turns into a 400.
func ParseResult(s string) (Result, bool) {
	switch r := Result(s); r {
	case ResultOK, ResultRateLimited, ResultBanned:
		return r, true
	default:
		return "", false
	}
}

var (
	// ErrNoMatch means no candidate could take a lease; AcquireStats says
	// why, and the API layer turns both into the 409 body.
	ErrNoMatch = errors.New("lease: no candidate available")
	// ErrUnknownLease means the lease ID does not resolve (404). Reports on
	// recently expired leases do NOT hit this — see the retention note on
	// Store.
	ErrUnknownLease = errors.New("lease: unknown lease id")
)

// Candidate is one leasable proxy as seen by the caller at selection time.
// The store itself knows nothing about Proxy objects — the discovery layer
// filters for health/attributes and passes what selection needs.
type Candidate struct {
	// Proxy is the opaque proxy key ("namespace/name").
	Proxy string
	// MaxLeases caps concurrent leases; 0 means unleasable.
	MaxLeases int32
	// Latency is the proxy's last reported latency, used as the tie-break.
	Latency time.Duration
}

// Lease is a granted assignment. Values returned by the store are copies;
// mutating them does not affect the store.
type Lease struct {
	ID        string
	Proxy     string
	Target    string
	ExpiresAt time.Time
}

// AcquireRequest carries the candidate set and lease parameters. Acquire
// deliberately takes the whole candidate set, not a pre-chosen proxy:
// selection and insertion must happen under one lock, or two concurrent
// requests both see "3 of 5 used" and overcommit.
type AcquireRequest struct {
	Candidates []Candidate
	// Target scopes the cooldown check; empty means the global pool.
	Target string
	TTL    time.Duration
}

// AcquireStats explains an ErrNoMatch (and is returned on success too):
// every candidate is either leased, at capacity, or in cooldown.
type AcquireStats struct {
	Considered int
	AtCapacity int
	InCooldown int
}

type cooldownKey struct{ proxy, target string }

// Store is the in-memory lease store. One mutex guards everything: at tens
// of proxies and human-rate QPS, sharding would be premature complexity.
//
// Retention: an expired lease is kept for CooldownWindow past its TTL so a
// Report arriving just after expiry still resolves — which matters most
// exactly when a proxy is being rate-limited. Acquire and the counts ignore
// retained leases; only the sweep finally drops them.
type Store struct {
	// CooldownWindow is how long a reported proxy/target pair is excluded
	// from selection (default 15m; --lease-cooldown in Step 10).
	CooldownWindow time.Duration
	// SweepInterval is how often the expiry sweep runs (default 30s).
	SweepInterval time.Duration

	now func() time.Time

	mu        sync.Mutex
	byID      map[string]*Lease
	byProxy   map[string]map[string]*Lease
	cooldowns map[cooldownKey]time.Time
}

// NewStore returns a ready Store. A non-positive cooldownWindow selects the
// 15-minute default.
func NewStore(cooldownWindow time.Duration) *Store {
	if cooldownWindow <= 0 {
		cooldownWindow = 15 * time.Minute
	}
	return &Store{
		CooldownWindow: cooldownWindow,
		SweepInterval:  30 * time.Second,
		now:            time.Now,
		byID:           map[string]*Lease{},
		byProxy:        map[string]map[string]*Lease{},
		cooldowns:      map[cooldownKey]time.Time{},
	}
}

// Acquire selects the least-loaded eligible candidate (ties: lowest
// latency, then name, so selection is deterministic and testable) and
// grants a lease on it.
func (s *Store) Acquire(_ context.Context, req AcquireRequest) (*Lease, AcquireStats, error) {
	stats := AcquireStats{Considered: len(req.Candidates)}
	if req.TTL <= 0 {
		return nil, stats, fmt.Errorf("lease: non-positive TTL %v", req.TTL)
	}
	now := s.now()

	s.mu.Lock()
	defer s.mu.Unlock()

	type eligible struct {
		cand   Candidate
		active int
	}
	var elig []eligible
	for _, c := range req.Candidates {
		if s.inCooldownLocked(c.Proxy, req.Target, now) {
			stats.InCooldown++
			continue
		}
		active := s.activeCountLocked(c.Proxy, now)
		if int32(active) >= c.MaxLeases {
			stats.AtCapacity++
			continue
		}
		elig = append(elig, eligible{cand: c, active: active})
	}
	if len(elig) == 0 {
		return nil, stats, ErrNoMatch
	}

	slices.SortFunc(elig, func(a, b eligible) int {
		if c := cmp.Compare(a.active, b.active); c != 0 {
			return c
		}
		if c := cmp.Compare(a.cand.Latency, b.cand.Latency); c != 0 {
			return c
		}
		return strings.Compare(a.cand.Proxy, b.cand.Proxy)
	})

	l := &Lease{
		ID:        rand.Text(),
		Proxy:     elig[0].cand.Proxy,
		Target:    req.Target,
		ExpiresAt: now.Add(req.TTL),
	}
	s.byID[l.ID] = l
	if s.byProxy[l.Proxy] == nil {
		s.byProxy[l.Proxy] = map[string]*Lease{}
	}
	s.byProxy[l.Proxy][l.ID] = l

	granted := *l
	return &granted, stats, nil
}

// Release drops a lease early. Idempotent: releasing an unknown or already
// expired lease is a no-op, so the API's DELETE can always answer 204.
func (s *Store) Release(_ context.Context, id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.dropLocked(id)
}

// Report records the outcome of using a lease. Rate-limited and banned
// results put the (proxy, target) pair in cooldown — target taken from the
// report, falling back to the lease's own target, falling back to the
// global pool. Reports on recently expired leases still resolve (see the
// retention note on Store).
func (s *Store) Report(_ context.Context, id string, result Result, target string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	l, ok := s.byID[id]
	if !ok {
		return ErrUnknownLease
	}
	if result == ResultOK {
		return nil
	}
	if target == "" {
		target = l.Target
	}
	s.cooldowns[cooldownKey{proxy: l.Proxy, target: target}] = s.now().Add(s.CooldownWindow)
	return nil
}

// ActiveCount returns the number of unexpired leases held on one proxy.
func (s *Store) ActiveCount(proxy string) int {
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.activeCountLocked(proxy, now)
}

// Counts returns the active-lease count per proxy, for the discovery list
// endpoint and the metrics collector. Proxies with no active leases are
// absent from the map.
func (s *Store) Counts() map[string]int {
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()
	counts := make(map[string]int, len(s.byProxy))
	for proxy := range s.byProxy {
		if n := s.activeCountLocked(proxy, now); n > 0 {
			counts[proxy] = n
		}
	}
	return counts
}

// Start runs the expiry sweep until ctx ends; it satisfies
// manager.Runnable so cmd/main.go can mgr.Add the store directly.
func (s *Store) Start(ctx context.Context) error {
	ticker := time.NewTicker(s.SweepInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
			s.sweep(s.now())
		}
	}
}

// NeedLeaderElection is false: lease state is per-process and the discovery
// API serves wherever this process runs, so the sweep must run there too.
func (s *Store) NeedLeaderElection() bool { return false }

// sweep drops leases past their retention window and elapsed cooldowns.
// Correctness never depends on sweep timing — every read path checks
// expiry against the clock — so this is purely garbage collection.
func (s *Store) sweep(now time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for id, l := range s.byID {
		if now.After(l.ExpiresAt.Add(s.CooldownWindow)) {
			s.dropLocked(id)
		}
	}
	for k, until := range s.cooldowns {
		if now.After(until) {
			delete(s.cooldowns, k)
		}
	}
}

func (s *Store) dropLocked(id string) {
	l, ok := s.byID[id]
	if !ok {
		return
	}
	delete(s.byID, id)
	delete(s.byProxy[l.Proxy], id)
	if len(s.byProxy[l.Proxy]) == 0 {
		delete(s.byProxy, l.Proxy)
	}
}

func (s *Store) activeCountLocked(proxy string, now time.Time) int {
	n := 0
	for _, l := range s.byProxy[proxy] {
		if now.Before(l.ExpiresAt) {
			n++
		}
	}
	return n
}

// inCooldownLocked: the global cooldown (empty target) always applies; a
// target-scoped cooldown additionally applies to acquisitions for that
// target. An acquisition without a target sees only the global pool — a
// proxy rate-limited by one site is still fine for everyone else.
func (s *Store) inCooldownLocked(proxy, target string, now time.Time) bool {
	if until, ok := s.cooldowns[cooldownKey{proxy: proxy}]; ok && now.Before(until) {
		return true
	}
	if target == "" {
		return false
	}
	until, ok := s.cooldowns[cooldownKey{proxy: proxy, target: target}]
	return ok && now.Before(until)
}