crossbar/docs/plans/v1/_files/internal/limiter/limiter_test.go

package limiter_test

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

	"git.wntrmute.dev/kyle/crossbar/internal/limiter"
)

func TestParallelAndQueue(t *testing.T) {
	l := limiter.New()
	l.Configure("alpha", "m", 2, 1) // two slots, one waiting place
	ctx := context.Background()

	rel1, w1, err := l.Acquire(ctx, "alpha", "m")
	if err != nil || w1 > 50*time.Millisecond {
		t.Fatalf("first acquire: err %v waited %v", err, w1)
	}
	rel2, _, err := l.Acquire(ctx, "alpha", "m")
	if err != nil {
		t.Fatalf("second acquire: %v", err)
	}
	if l.InFlight("alpha", "m") != 2 || l.FreeSlots("alpha") != 0 {
		t.Errorf("in flight %d free %d, want 2 and 0", l.InFlight("alpha", "m"), l.FreeSlots("alpha"))
	}

	// Third waits in the queue.
	got3 := make(chan error, 1)
	go func() {
		rel, waited, err := l.Acquire(ctx, "alpha", "m")
		if err == nil {
			if waited < 40*time.Millisecond {
				err = errors.New("third acquire did not wait")
			}
			rel() // release before reporting, so the final count check cannot race it
		}
		got3 <- err
	}()
	waitUntil(t, func() bool { return l.Queued("alpha", "m") == 1 })
	// Fourth finds the queue full and is refused at once.
	start := time.Now()
	_, _, err = l.Acquire(ctx, "alpha", "m")
	if !errors.Is(err, limiter.ErrQueueFull) {
		t.Fatalf("fourth acquire: %v, want ErrQueueFull", err)
	}
	if time.Since(start) > 50*time.Millisecond {
		t.Errorf("a full queue must refuse immediately, took %v", time.Since(start))
	}
	time.Sleep(50 * time.Millisecond) // a lower bound on the third's wait, checked above as >= 40 ms
	rel1()                            // frees a slot: the queued third proceeds
	select {
	case err := <-got3:
		if err != nil {
			t.Fatalf("third: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("queued acquire did not proceed after a release")
	}
	rel2()
	if l.InFlight("alpha", "m") != 0 || l.Queued("alpha", "m") != 0 {
		t.Errorf("after releases: inflight %d queued %d", l.InFlight("alpha", "m"), l.Queued("alpha", "m"))
	}
}

func TestReleaseIsIdempotent(t *testing.T) {
	l := limiter.New()
	l.Configure("h", "m", 1, 0)
	rel, _, err := l.Acquire(context.Background(), "h", "m")
	if err != nil {
		t.Fatal(err)
	}
	rel()
	rel() // a second call must not free a slot that was never taken
	if l.InFlight("h", "m") != 0 {
		t.Errorf("in flight %d after double release", l.InFlight("h", "m"))
	}
	if _, _, err := l.Acquire(context.Background(), "h", "m"); err != nil {
		t.Errorf("slot must be free again: %v", err)
	}
}

func TestCancelWhileQueuedLeaksNothing(t *testing.T) {
	l := limiter.New()
	l.Configure("h", "m", 1, 2)
	rel, _, err := l.Acquire(context.Background(), "h", "m")
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { _, _, err := l.Acquire(ctx, "h", "m"); done <- err }()
	waitUntil(t, func() bool { return l.Queued("h", "m") == 1 })
	cancel()
	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("cancelled acquire returned %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("cancelled acquire did not return")
	}
	if l.Queued("h", "m") != 0 {
		t.Errorf("queued = %d after cancel", l.Queued("h", "m"))
	}
	rel()
	if l.InFlight("h", "m") != 0 {
		t.Errorf("in flight %d, the cancelled waiter must not have taken the slot", l.InFlight("h", "m"))
	}
}

func TestQueueIsFIFO(t *testing.T) {
	l := limiter.New()
	l.Configure("h", "m", 1, 8)
	rel, _, err := l.Acquire(context.Background(), "h", "m")
	if err != nil {
		t.Fatal(err)
	}
	var mu sync.Mutex
	var order []int
	var wg sync.WaitGroup
	for i := 1; i <= 4; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			r, _, err := l.Acquire(context.Background(), "h", "m")
			if err != nil {
				t.Errorf("waiter %d: %v", i, err)
				return
			}
			mu.Lock()
			order = append(order, i)
			mu.Unlock()
			time.Sleep(5 * time.Millisecond)
			r()
		}(i)
		waitUntil(t, func() bool { return l.Queued("h", "m") == i }) // arrivals in order, by observation
	}
	rel()
	wg.Wait()
	if len(order) != 4 || order[0] != 1 || order[1] != 2 || order[2] != 3 || order[3] != 4 {
		t.Errorf("waiters proceeded in order %v, want [1 2 3 4]", order)
	}
}

func TestUnconfiguredPairIsOneSlotNoQueue(t *testing.T) {
	l := limiter.New()
	rel, _, err := l.Acquire(context.Background(), "x", "y")
	if err != nil {
		t.Fatal(err)
	}
	defer rel()
	if _, _, err := l.Acquire(context.Background(), "x", "y"); !errors.Is(err, limiter.ErrQueueFull) {
		t.Errorf("second acquire on an unconfigured pair: %v, want ErrQueueFull", err)
	}
}

func TestFreeSlotsSumsModels(t *testing.T) {
	l := limiter.New()
	l.Configure("h", "a", 4, 0)
	l.Configure("h", "b", 2, 0)
	if got := l.FreeSlots("h"); got != 6 {
		t.Fatalf("free = %d, want 6", got)
	}
	rel, _, _ := l.Acquire(context.Background(), "h", "a")
	defer rel()
	if got := l.FreeSlots("h"); got != 5 {
		t.Errorf("free = %d, want 5", got)
	}
	if l.FreeSlots("nobody") != 0 {
		t.Errorf("unknown host has no slots")
	}
}

// waitUntil polls cond every millisecond for up to two seconds and fails the test if it never holds.
func waitUntil(t *testing.T, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatal("condition not reached within two seconds")
}