homelab-frontend/internal/proxy/scheduler_test.go

package proxy

import (
	"context"
	"errors"
	"runtime"
	"sync"
	"testing"
	"time"
)

// TestBackendAcquireBlocksThenUnblocksOnRelease verifies that a second Acquire
// blocks while a slot is held, then unblocks when Release is called.
func TestBackendAcquireBlocksThenUnblocksOnRelease(t *testing.T) {
	b := newBackend("host:1", 1, 1)

	// G1: acquire and hold the only slot.
	if err := b.Acquire(context.Background()); err != nil {
		t.Fatalf("G1 Acquire: unexpected error: %v", err)
	}

	g2Done := make(chan error, 1)
	go func() {
		g2Done <- b.Acquire(context.Background())
	}()

	// Wait until G2 enters the queue channel (proves Acquire is blocking on slots).
	for len(b.queue) == 0 {
		runtime.Gosched()
	}

	select {
	case err := <-g2Done:
		t.Fatalf("G2 returned before Release: %v", err)
	default:
	}

	// G1 releases; G2 should unblock.
	b.Release()

	select {
	case err := <-g2Done:
		if err != nil {
			t.Fatalf("G2 Acquire after Release: unexpected error: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("G2 did not unblock within 1s after Release")
	}

	// Cleanup: release G2's slot.
	b.Release()
}

// TestBackendAcquireReturnsErrQueueFullWhenQueueFull verifies that Acquire returns
// ErrQueueFull immediately when no queue capacity exists (maxQueue=0).
// With a zero-capacity queue channel the non-blocking select always hits default.
func TestBackendAcquireReturnsErrQueueFullWhenQueueFull(t *testing.T) {
	b := newBackend("host:1", 1, 0)

	err := b.Acquire(context.Background())
	if !errors.Is(err, ErrQueueFull) {
		t.Fatalf("expected ErrQueueFull, got %v", err)
	}
}

// TestSchedulerSelectRendezvousHashDeterministic verifies that the same cache key
// always selects the same backend (rendezvous / highest-random-weight hashing).
func TestSchedulerSelectRendezvousHashDeterministic(t *testing.T) {
	addrs := []string{"a:80", "b:80", "c:80"}
	s := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})

	first := s.Select("same-cache-key")
	if first == nil {
		t.Fatal("Select returned nil")
	}
	for i := 1; i < 10; i++ {
		b := s.Select("same-cache-key")
		if b == nil || b.Address != first.Address {
			t.Fatalf("call %d: want %s, got %v", i, first.Address, b)
		}
	}
}

// TestSchedulerSelectNoKeyRoundRobin verifies that round-robin distributes
// requests evenly across all backends when no cache key is given.
func TestSchedulerSelectNoKeyRoundRobin(t *testing.T) {
	addrs := []string{"a:80", "b:80", "c:80"}
	s := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})

	counts := make(map[string]int)
	for i := 0; i < 9; i++ {
		b := s.Select("")
		if b == nil {
			t.Fatalf("call %d: Select returned nil", i)
		}
		counts[b.Address]++
	}

	for _, addr := range addrs {
		if counts[addr] != 3 {
			t.Errorf("addr %s: want 3 selections, got %d (all counts: %v)", addr, counts[addr], counts)
		}
	}
}

// TestBackendAcquireCancelledContext verifies that Acquire returns ctx.Err() when the
// context is cancelled while the caller is blocked waiting for a concurrency slot,
// and that the queue depth returns to 0 after cancellation.
func TestBackendAcquireCancelledContext(t *testing.T) {
	b := newBackend("host:1", 1, 1)

	// G1: acquire and hold the only slot indefinitely.
	if err := b.Acquire(context.Background()); err != nil {
		t.Fatalf("G1 Acquire: unexpected error: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	g2Done := make(chan error, 1)
	go func() {
		g2Done <- b.Acquire(ctx)
	}()

	// Wait until G2 enters the queue channel.
	for len(b.queue) == 0 {
		runtime.Gosched()
	}

	// Verify G2 is queued (queue channel has one item).
	if qd := len(b.queue); qd != 1 {
		t.Errorf("expected queue depth 1 while G2 blocked, got %d", qd)
	}

	// Cancel ctx while G2 is waiting for the slot.
	cancel()

	select {
	case err := <-g2Done:
		if err == nil {
			t.Fatal("G2 Acquire: expected ctx.Err(), got nil")
		}
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("G2 Acquire: expected context.Canceled, got %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("G2 did not return within 1s after context cancel")
	}

	// Queue depth must return to 0 after G2 is drained.
	if qd := len(b.queue); qd != 0 {
		t.Errorf("expected queue depth 0 after cancellation, got %d", qd)
	}

	// Cleanup: release G1's slot.
	b.Release()
}

// TestSchedulerMarkFailedExcludesBackend verifies that a failed backend is not
// selected within the fail window (30 s default).
func TestSchedulerMarkFailedExcludesBackend(t *testing.T) {
	addrs := []string{"a:80", "b:80"}
	s := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})

	s.MarkFailed("a:80")

	var mu sync.Mutex
	seen := make(map[string]int)
	for i := 0; i < 10; i++ {
		b := s.Select("any-key")
		if b == nil {
			t.Fatalf("call %d: Select returned nil", i)
		}
		mu.Lock()
		seen[b.Address]++
		mu.Unlock()
	}

	if seen["a:80"] > 0 {
		t.Errorf("failed backend a:80 was selected %d times; expected 0 within fail window", seen["a:80"])
	}
}