Add the in-memory lease store: least-loaded selection, cooldowns, TTL retention
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
319
internal/lease/store.go
Normal file
319
internal/lease/store.go
Normal file
@@ -0,0 +1,319 @@
|
||||
// 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)
|
||||
}
|
||||
Reference in New Issue
Block a user