proxy-operator: Kubernetes operator for crawling-proxy fleets #1
@@ -1,10 +1,12 @@
|
||||
# Architecture
|
||||
|
||||
> **Status:** the operator is built through Step 5 (health engine) of
|
||||
> **Status:** the operator is built through Step 6 (lease store) of
|
||||
> [docs/plans/2026-08-07-1747-proxy-operator.md](plans/2026-08-07-1747-proxy-operator.md).
|
||||
> This document currently covers the event/reconcile flow; the components
|
||||
> table and the Decisions section arrive with Step 10, and the diagrams
|
||||
> below grow as the lease store, discovery API, and orphan GC land.
|
||||
> table and the Decisions section arrive with Step 10. The lease store
|
||||
> (`internal/lease/`) is HTTP-driven, not cluster-event-driven, so its
|
||||
> diagram lands together with the discovery API in Step 7; the orphan-GC
|
||||
> flow lands with Step 9.
|
||||
|
||||
## Event flow: cluster events → reconciler functions
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ Pairs with [docs/plans/2026-08-07-1747-proxy-operator.md](../plans/2026-08-07-17
|
||||
- [x] Step 3 — Kubernetes pod provider (`internal/provider/kubernetes/`; first built as an in-memory mock, then replaced — see the two Step 3 sections below)
|
||||
- [x] Step 4 — Reconciler (`internal/controller/`)
|
||||
- [x] Step 5 — Health engine (`internal/health/`)
|
||||
- [ ] Step 6 — Lease store (`internal/lease/`)
|
||||
- [x] Step 6 — Lease store (`internal/lease/`)
|
||||
- [ ] Step 7 — Discovery API (`internal/discovery/`)
|
||||
- [ ] Step 8 — GCP provider (`internal/provider/gcp/`)
|
||||
- [ ] Step 9 — Orphan GC + metrics
|
||||
@@ -694,3 +694,67 @@ engine deliberately knows nothing about conditions except reading one at
|
||||
seed time, keeping the state/representation split honest. The
|
||||
`hint`-driven `wg.Go` idiom (Go 1.25+) replaced the classic
|
||||
`wg.Add/defer wg.Done` in the worker pool.
|
||||
|
||||
## Step 6 — Lease store (`internal/lease/`)
|
||||
|
||||
Implemented `store.go` per the plan: `Acquire` takes the whole candidate
|
||||
set so selection and insertion happen under the one store mutex (no
|
||||
overcommit between concurrent requests), selection is a linear scan +
|
||||
`slices.SortFunc` on `(activeLeases asc, latency asc, name asc)`,
|
||||
`AcquireStats{Considered, AtCapacity, InCooldown}` feeds Step 7's 409
|
||||
body, cooldowns live in a `map[{proxy, target}]time.Time` (empty target =
|
||||
global pool), and expired leases are retained for `CooldownWindow` past
|
||||
their TTL so a late `Report` — arriving exactly when a proxy is being
|
||||
rate-limited — still resolves and records its cooldown.
|
||||
|
||||
Semantics pinned against the spec (§8) rather than guessed:
|
||||
|
||||
- Report results are exactly `ok | rate_limited | banned` (`ParseResult`
|
||||
gives the API layer its 400 check). `rate_limited` and `banned` both
|
||||
record a cooldown for the same window; `ok` records nothing.
|
||||
Distinguishing ban duration from rate-limit duration would be a second
|
||||
knob the spec doesn't ask for — noted for the Decisions section.
|
||||
- Cooldown scoping: the global cooldown (empty target) always applies; a
|
||||
target-scoped cooldown additionally blocks acquisitions for that target;
|
||||
acquisitions without a target see only the global pool ("a proxy
|
||||
rate-limited by one site is still fine for everyone else").
|
||||
- A `Report` without a target falls back to the lease's own target before
|
||||
falling back to global — so a client that leased with a target doesn't
|
||||
accidentally poison the whole proxy by omitting it in the report.
|
||||
|
||||
Design notes:
|
||||
|
||||
- **Correctness never depends on the sweep.** Every read path
|
||||
(`Acquire`/`ActiveCount`/`Counts`) compares `ExpiresAt` against the
|
||||
injected clock, so TTL expiry frees capacity immediately even if the
|
||||
background loop hasn't run; the sweep is purely garbage collection. The
|
||||
plan's `ExpireLoop` became `Start(ctx)` + `NeedLeaderElection() false`
|
||||
so the store satisfies `manager.Runnable` directly — Step 10 just
|
||||
`mgr.Add(store)`s it. Not leader-elected because lease state is
|
||||
per-process and must expire wherever the discovery API is serving.
|
||||
- The store knows nothing about Proxy objects — `Candidate` carries the
|
||||
opaque key, `MaxLeases`, and latency; the discovery layer does the
|
||||
health/attribute filtering. The spec's `LeaseStore` interface will be
|
||||
defined consumer-side in `internal/discovery` (Step 7), per Go idiom;
|
||||
this package exports only the concrete in-memory `*Store`.
|
||||
- Lease IDs come from `crypto/rand.Text()` (Go 1.24+); returned `Lease`
|
||||
values are copies so callers can't mutate store internals.
|
||||
|
||||
Tests (94.8% coverage, `-race -count=2` clean): capacity + release
|
||||
freeing slots, `MaxLeases=0` unleasable, least-loaded/latency/name
|
||||
selection order, target-scoped vs global cooldown scoping, cooldown
|
||||
expiry via the injected fake clock, TTL freeing capacity with no sweep,
|
||||
report-on-expired-but-retained lease (then `ErrUnknownLease` after
|
||||
retention), `ok` recording nothing, idempotent release, `ParseResult`,
|
||||
40 concurrent acquires against `MaxLeases=5` granting exactly 5, and the
|
||||
`Start` loop sweeping then stopping cleanly on cancel.
|
||||
|
||||
```bash
|
||||
go test -race -count=2 ./internal/lease/
|
||||
make test # whole repo green, other packages' coverage unchanged
|
||||
```
|
||||
|
||||
Worth noting: `docs/architecture.md` was not extended this step — the
|
||||
lease store is HTTP-driven, not cluster-event-driven, so its diagram
|
||||
belongs with the discovery API and lands in Step 7 (banner updated to say
|
||||
so).
|
||||
|
||||
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)
|
||||
}
|
||||
365
internal/lease/store_test.go
Normal file
365
internal/lease/store_test.go
Normal file
@@ -0,0 +1,365 @@
|
||||
package lease
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// fakeClock is an injectable, manually advanced clock.
|
||||
type fakeClock struct {
|
||||
mu sync.Mutex
|
||||
cur time.Time
|
||||
}
|
||||
|
||||
func newFakeClock() *fakeClock {
|
||||
return &fakeClock{cur: time.Date(2026, 8, 9, 12, 0, 0, 0, time.UTC)}
|
||||
}
|
||||
|
||||
func (c *fakeClock) Now() time.Time {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.cur
|
||||
}
|
||||
|
||||
func (c *fakeClock) Advance(d time.Duration) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.cur = c.cur.Add(d)
|
||||
}
|
||||
|
||||
func newTestStore() (*Store, *fakeClock) {
|
||||
s := NewStore(15 * time.Minute)
|
||||
clock := newFakeClock()
|
||||
s.now = clock.Now
|
||||
return s, clock
|
||||
}
|
||||
|
||||
func candidate(proxy string, maxLeases int32, latency time.Duration) Candidate {
|
||||
return Candidate{Proxy: proxy, MaxLeases: maxLeases, Latency: latency}
|
||||
}
|
||||
|
||||
func mustAcquire(t *testing.T, s *Store, req AcquireRequest) *Lease {
|
||||
t.Helper()
|
||||
l, _, err := s.Acquire(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("Acquire: %v", err)
|
||||
}
|
||||
return l
|
||||
}
|
||||
|
||||
func TestAcquire_capacity(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
req := AcquireRequest{Candidates: []Candidate{candidate("ns/p1", 2, 0)}, TTL: time.Minute}
|
||||
|
||||
l1 := mustAcquire(t, s, req)
|
||||
l2 := mustAcquire(t, s, req)
|
||||
if l1.ID == l2.ID {
|
||||
t.Fatal("two leases share an ID")
|
||||
}
|
||||
if got := s.ActiveCount("ns/p1"); got != 2 {
|
||||
t.Fatalf("ActiveCount = %d, want 2", got)
|
||||
}
|
||||
|
||||
_, stats, err := s.Acquire(context.Background(), req)
|
||||
if !errors.Is(err, ErrNoMatch) {
|
||||
t.Fatalf("third acquire error = %v, want ErrNoMatch", err)
|
||||
}
|
||||
want := AcquireStats{Considered: 1, AtCapacity: 1}
|
||||
if stats != want {
|
||||
t.Errorf("stats = %+v, want %+v", stats, want)
|
||||
}
|
||||
|
||||
// Early release frees the slot again.
|
||||
s.Release(context.Background(), l1.ID)
|
||||
mustAcquire(t, s, req)
|
||||
}
|
||||
|
||||
func TestAcquire_maxLeasesZeroIsUnleasable(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
_, stats, err := s.Acquire(context.Background(), AcquireRequest{
|
||||
Candidates: []Candidate{candidate("ns/p1", 0, 0)},
|
||||
TTL: time.Minute,
|
||||
})
|
||||
if !errors.Is(err, ErrNoMatch) {
|
||||
t.Fatalf("err = %v, want ErrNoMatch", err)
|
||||
}
|
||||
if stats.AtCapacity != 1 {
|
||||
t.Errorf("stats = %+v, want the unleasable proxy counted AtCapacity", stats)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcquire_selectionOrder(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("least loaded wins", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
mustAcquire(t, s, AcquireRequest{
|
||||
Candidates: []Candidate{candidate("ns/a", 5, 10*time.Millisecond)}, TTL: time.Minute,
|
||||
})
|
||||
l := mustAcquire(t, s, AcquireRequest{
|
||||
Candidates: []Candidate{
|
||||
candidate("ns/a", 5, 10*time.Millisecond), // 1 active, lower latency
|
||||
candidate("ns/b", 5, 90*time.Millisecond), // 0 active
|
||||
},
|
||||
TTL: time.Minute,
|
||||
})
|
||||
if l.Proxy != "ns/b" {
|
||||
t.Errorf("chose %s, want the least-loaded ns/b", l.Proxy)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("latency breaks the load tie", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
l := mustAcquire(t, s, AcquireRequest{
|
||||
Candidates: []Candidate{
|
||||
candidate("ns/a", 5, 90*time.Millisecond),
|
||||
candidate("ns/b", 5, 10*time.Millisecond),
|
||||
},
|
||||
TTL: time.Minute,
|
||||
})
|
||||
if l.Proxy != "ns/b" {
|
||||
t.Errorf("chose %s, want the lower-latency ns/b", l.Proxy)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("name breaks a full tie deterministically", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
l := mustAcquire(t, s, AcquireRequest{
|
||||
Candidates: []Candidate{
|
||||
candidate("ns/b", 5, 10*time.Millisecond),
|
||||
candidate("ns/a", 5, 10*time.Millisecond),
|
||||
},
|
||||
TTL: time.Minute,
|
||||
})
|
||||
if l.Proxy != "ns/a" {
|
||||
t.Errorf("chose %s, want ns/a (lexicographic tie-break)", l.Proxy)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestAcquire_cooldownScoping(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
cands := []Candidate{candidate("ns/p1", 5, 0)}
|
||||
|
||||
l := mustAcquire(t, s, AcquireRequest{Candidates: cands, Target: "example.com", TTL: time.Minute})
|
||||
if err := s.Report(context.Background(), l.ID, ResultRateLimited, "example.com"); err != nil {
|
||||
t.Fatalf("Report: %v", err)
|
||||
}
|
||||
|
||||
// Same target: excluded.
|
||||
_, stats, err := s.Acquire(context.Background(), AcquireRequest{
|
||||
Candidates: cands, Target: "example.com", TTL: time.Minute,
|
||||
})
|
||||
if !errors.Is(err, ErrNoMatch) || stats.InCooldown != 1 {
|
||||
t.Errorf("same-target acquire = (%v, %+v), want ErrNoMatch with InCooldown=1", err, stats)
|
||||
}
|
||||
|
||||
// Different target: fine.
|
||||
mustAcquire(t, s, AcquireRequest{Candidates: cands, Target: "other.org", TTL: time.Minute})
|
||||
|
||||
// No target (global pool): a target-scoped cooldown does not apply.
|
||||
mustAcquire(t, s, AcquireRequest{Candidates: cands, TTL: time.Minute})
|
||||
}
|
||||
|
||||
func TestAcquire_globalCooldownBlocksEverything(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
cands := []Candidate{candidate("ns/p1", 5, 0)}
|
||||
|
||||
// A lease without a target, reported banned without a target: the
|
||||
// cooldown lands on the global pool.
|
||||
l := mustAcquire(t, s, AcquireRequest{Candidates: cands, TTL: time.Minute})
|
||||
if err := s.Report(context.Background(), l.ID, ResultBanned, ""); err != nil {
|
||||
t.Fatalf("Report: %v", err)
|
||||
}
|
||||
|
||||
for _, target := range []string{"", "example.com"} {
|
||||
_, stats, err := s.Acquire(context.Background(), AcquireRequest{
|
||||
Candidates: cands, Target: target, TTL: time.Minute,
|
||||
})
|
||||
if !errors.Is(err, ErrNoMatch) || stats.InCooldown != 1 {
|
||||
t.Errorf("acquire(target=%q) = (%v, %+v), want global cooldown to block", target, err, stats)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcquire_cooldownExpires(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, clock := newTestStore()
|
||||
cands := []Candidate{candidate("ns/p1", 5, 0)}
|
||||
|
||||
l := mustAcquire(t, s, AcquireRequest{Candidates: cands, TTL: time.Minute})
|
||||
if err := s.Report(context.Background(), l.ID, ResultRateLimited, ""); err != nil {
|
||||
t.Fatalf("Report: %v", err)
|
||||
}
|
||||
if _, _, err := s.Acquire(context.Background(), AcquireRequest{Candidates: cands, TTL: time.Minute}); !errors.Is(err, ErrNoMatch) {
|
||||
t.Fatal("expected cooldown to block immediately after the report")
|
||||
}
|
||||
|
||||
clock.Advance(15*time.Minute + time.Second)
|
||||
mustAcquire(t, s, AcquireRequest{Candidates: cands, TTL: time.Minute})
|
||||
}
|
||||
|
||||
func TestExpiry_freesCapacityWithoutSweep(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, clock := newTestStore()
|
||||
req := AcquireRequest{Candidates: []Candidate{candidate("ns/p1", 1, 0)}, TTL: time.Minute}
|
||||
|
||||
mustAcquire(t, s, req)
|
||||
if _, _, err := s.Acquire(context.Background(), req); !errors.Is(err, ErrNoMatch) {
|
||||
t.Fatal("capacity 1 not enforced")
|
||||
}
|
||||
|
||||
clock.Advance(2 * time.Minute)
|
||||
// No sweep has run; expiry must still free capacity and zero the counts.
|
||||
if got := s.ActiveCount("ns/p1"); got != 0 {
|
||||
t.Fatalf("ActiveCount after TTL = %d, want 0", got)
|
||||
}
|
||||
if counts := s.Counts(); len(counts) != 0 {
|
||||
t.Fatalf("Counts after TTL = %v, want empty", counts)
|
||||
}
|
||||
mustAcquire(t, s, req)
|
||||
}
|
||||
|
||||
func TestReport_expiredButRetainedLease(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, clock := newTestStore()
|
||||
cands := []Candidate{candidate("ns/p1", 5, 0)}
|
||||
|
||||
l := mustAcquire(t, s, AcquireRequest{Candidates: cands, Target: "example.com", TTL: time.Minute})
|
||||
|
||||
// TTL lapses; the report arrives late — exactly when the proxy is being
|
||||
// rate-limited, which is when the cooldown matters most.
|
||||
clock.Advance(5 * time.Minute)
|
||||
s.sweep(clock.Now())
|
||||
if err := s.Report(context.Background(), l.ID, ResultRateLimited, ""); err != nil {
|
||||
t.Fatalf("Report on an expired-but-retained lease: %v", err)
|
||||
}
|
||||
// The cooldown fell back to the lease's own target.
|
||||
_, stats, err := s.Acquire(context.Background(), AcquireRequest{
|
||||
Candidates: cands, Target: "example.com", TTL: time.Minute,
|
||||
})
|
||||
if !errors.Is(err, ErrNoMatch) || stats.InCooldown != 1 {
|
||||
t.Errorf("acquire = (%v, %+v), want cooldown from the late report", err, stats)
|
||||
}
|
||||
|
||||
// Past the retention window the sweep finally drops it.
|
||||
clock.Advance(15 * time.Minute)
|
||||
s.sweep(clock.Now())
|
||||
if err := s.Report(context.Background(), l.ID, ResultRateLimited, ""); !errors.Is(err, ErrUnknownLease) {
|
||||
t.Errorf("Report after retention = %v, want ErrUnknownLease", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReport_okRecordsNothing(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
cands := []Candidate{candidate("ns/p1", 5, 0)}
|
||||
l := mustAcquire(t, s, AcquireRequest{Candidates: cands, Target: "example.com", TTL: time.Minute})
|
||||
|
||||
if err := s.Report(context.Background(), l.ID, ResultOK, "example.com"); err != nil {
|
||||
t.Fatalf("Report(ok): %v", err)
|
||||
}
|
||||
mustAcquire(t, s, AcquireRequest{Candidates: cands, Target: "example.com", TTL: time.Minute})
|
||||
}
|
||||
|
||||
func TestRelease_isIdempotent(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
l := mustAcquire(t, s, AcquireRequest{Candidates: []Candidate{candidate("ns/p1", 1, 0)}, TTL: time.Minute})
|
||||
|
||||
s.Release(context.Background(), l.ID)
|
||||
s.Release(context.Background(), l.ID)
|
||||
s.Release(context.Background(), "never-existed")
|
||||
if got := s.ActiveCount("ns/p1"); got != 0 {
|
||||
t.Errorf("ActiveCount = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseResult(t *testing.T) {
|
||||
t.Parallel()
|
||||
for _, valid := range []string{"ok", "rate_limited", "banned"} {
|
||||
if _, ok := ParseResult(valid); !ok {
|
||||
t.Errorf("ParseResult(%q) rejected a valid value", valid)
|
||||
}
|
||||
}
|
||||
for _, invalid := range []string{"", "OK", "throttled", "rate-limited"} {
|
||||
if _, ok := ParseResult(invalid); ok {
|
||||
t.Errorf("ParseResult(%q) accepted an invalid value", invalid)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcquire_concurrentNeverOvercommits(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, _ := newTestStore()
|
||||
req := AcquireRequest{Candidates: []Candidate{candidate("ns/p1", 5, 0)}, TTL: time.Minute}
|
||||
|
||||
const attempts = 40
|
||||
var wg sync.WaitGroup
|
||||
granted := make(chan *Lease, attempts)
|
||||
for range attempts {
|
||||
wg.Go(func() {
|
||||
if l, _, err := s.Acquire(context.Background(), req); err == nil {
|
||||
granted <- l
|
||||
}
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
close(granted)
|
||||
|
||||
var n int
|
||||
for range granted {
|
||||
n++
|
||||
}
|
||||
if n != 5 {
|
||||
t.Errorf("%d of %d concurrent acquires granted, want exactly MaxLeases=5", n, attempts)
|
||||
}
|
||||
if got := s.ActiveCount("ns/p1"); got != 5 {
|
||||
t.Errorf("ActiveCount = %d, want 5", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStart_sweepsAndStops(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, clock := newTestStore()
|
||||
s.SweepInterval = time.Millisecond
|
||||
|
||||
l := mustAcquire(t, s, AcquireRequest{Candidates: []Candidate{candidate("ns/p1", 5, 0)}, TTL: time.Minute})
|
||||
clock.Advance(20 * time.Minute) // past TTL + retention
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- s.Start(ctx) }()
|
||||
|
||||
deadline := time.After(5 * time.Second)
|
||||
for {
|
||||
if err := s.Report(context.Background(), l.ID, ResultOK, ""); errors.Is(err, ErrUnknownLease) {
|
||||
break
|
||||
}
|
||||
select {
|
||||
case <-deadline:
|
||||
t.Fatal("sweep never dropped the lease")
|
||||
case <-time.After(5 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
|
||||
cancel()
|
||||
select {
|
||||
case err := <-done:
|
||||
if err != nil {
|
||||
t.Errorf("Start returned %v, want nil", err)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Start did not stop on cancel")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user