diff --git a/docs/architecture.md b/docs/architecture.md index 6159760..ae490c3 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -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 diff --git a/docs/plans-executions/2026-08-07-1747-proxy-operator.md b/docs/plans-executions/2026-08-07-1747-proxy-operator.md index 2938982..7795d1a 100644 --- a/docs/plans-executions/2026-08-07-1747-proxy-operator.md +++ b/docs/plans-executions/2026-08-07-1747-proxy-operator.md @@ -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). diff --git a/internal/lease/store.go b/internal/lease/store.go new file mode 100644 index 0000000..4be66f8 --- /dev/null +++ b/internal/lease/store.go @@ -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) +} diff --git a/internal/lease/store_test.go b/internal/lease/store_test.go new file mode 100644 index 0000000..e2c2300 --- /dev/null +++ b/internal/lease/store_test.go @@ -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") + } +}