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

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")
	}
}