Add the in-memory lease store: least-loaded selection, cooldowns, TTL retention
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
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