Add the per-host-model limiter and the host chooser

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
2026-09-25 03:53:38 -07:00
parent 463cea18de
commit 7e0dbb4a6f
5 changed files with 505 additions and 0 deletions
@@ -0,0 +1,180 @@
+// Package limiter hands out at most `parallel` concurrent slots per (host, model) and lets at
+// most `queue_max` requests wait in a FIFO. A request that finds the queue full is refused at
+// once so the caller can retry elsewhere; a waiting request can cancel and leave without leaking
+// a slot or a queue place.
+package limiter
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+)
+
+// ErrQueueFull is returned by Acquire when the queue is already at queue_max; the caller may try a
+// different host.
+var ErrQueueFull = errors.New("queue full")
+
+// defaults is what an unconfigured (host, model) behaves as: one slot, no waiting room.
+const (
+	defaultParallel = 1
+	defaultQueueMax = 0
+)
+
+// pair holds the live state for one (host, model): how many slots exist, how many are taken, and
+// the FIFO of waiters. All fields are guarded by Limiter.mu.
+type pair struct {
+	parallel   int
+	queueMax   int
+	inflight   int
+	waiters    []chan struct{}
+	configured bool
+}
+
+// Limiter tracks one pair per (host, model). Safe for concurrent use.
+type Limiter struct {
+	mu    sync.Mutex
+	pairs map[pairKey]*pair
+}
+
+type pairKey struct {
+	host  string
+	model string
+}
+
+// New returns an empty Limiter.
+func New() *Limiter {
+	return &Limiter{pairs: make(map[pairKey]*pair)}
+}
+
+// Configure sets the slot and queue limits for one (host, model). It may be called before any
+// request or after one has created the pair with the defaults; either way the limits apply.
+func (l *Limiter) Configure(host, model string, parallel, queueMax int) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	p := l.pairLocked(host, model)
+	p.parallel = parallel
+	p.queueMax = queueMax
+	p.configured = true
+}
+
+// pairLocked returns the pair for (host, model), creating it with the unconfigured defaults if it
+// does not exist yet. The caller holds l.mu.
+func (l *Limiter) pairLocked(host, model string) *pair {
+	k := pairKey{host, model}
+	p := l.pairs[k]
+	if p == nil {
+		p = &pair{parallel: defaultParallel, queueMax: defaultQueueMax}
+		l.pairs[k] = p
+	}
+	return p
+}
+
+// Acquire blocks until a slot is held. It returns a release that gives the slot back exactly once
+// (a second call is a no-op), how long the caller spent in the queue, and an error: ErrQueueFull
+// when the queue is already full (returned immediately, without waiting), or ctx.Err() when the
+// context ends while waiting.
+func (l *Limiter) Acquire(ctx context.Context, host, model string) (release func(), waited time.Duration, err error) {
+	l.mu.Lock()
+	start := time.Now()
+	p := l.pairLocked(host, model)
+	if p.inflight < p.parallel {
+		p.inflight++
+		l.mu.Unlock()
+		return l.release(p), time.Since(start), nil
+	}
+	if len(p.waiters) >= p.queueMax {
+		l.mu.Unlock()
+		return nil, time.Since(start), ErrQueueFull
+	}
+	waiter := make(chan struct{})
+	p.waiters = append(p.waiters, waiter)
+	l.mu.Unlock()
+
+	select {
+	case <-ctx.Done():
+		// The slot may have been handed to us the instant the context ended; release it either
+		// way so neither a slot nor a queue place leaks.
+		l.mu.Lock()
+		if !p.dropWaiter(waiter) {
+			l.mu.Unlock()
+			l.release(p)
+			return nil, time.Since(start), ctx.Err()
+		}
+		l.mu.Unlock()
+		return nil, time.Since(start), ctx.Err()
+	case <-waiter:
+		return l.release(p), time.Since(start), nil
+	}
+}
+
+// release returns the function the caller holds for a slot: it hands the slot to the next waiter
+// if one is waiting, otherwise it frees the slot. It is safe to call through the sync.Once that
+// Acquire wrapped it in.
+func (l *Limiter) release(p *pair) func() {
+	var once sync.Once
+	return func() {
+		once.Do(func() {
+			l.mu.Lock()
+			defer l.mu.Unlock()
+			if len(p.waiters) > 0 {
+				next := p.waiters[0]
+				p.waiters = p.waiters[1:]
+				close(next)
+				return
+			}
+			p.inflight--
+		})
+	}
+}
+
+// dropWaiter removes w from the middle of the queue. It reports whether w was there; false means
+// the slot was already handed to w (its channel closed) and the caller must give it back.
+func (p *pair) dropWaiter(w chan struct{}) bool {
+	for i, cw := range p.waiters {
+		if cw == w {
+			p.waiters = append(p.waiters[:i], p.waiters[i+1:]...)
+			return true
+		}
+	}
+	return false
+}
+
+// InFlight returns the number of held slots for (host, model).
+func (l *Limiter) InFlight(host, model string) int {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	p := l.pairs[pairKey{host, model}]
+	if p == nil {
+		return 0
+	}
+	return p.inflight
+}
+
+// Queued returns the number of requests waiting for (host, model).
+func (l *Limiter) Queued(host, model string) int {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	p := l.pairs[pairKey{host, model}]
+	if p == nil {
+		return 0
+	}
+	return len(p.waiters)
+}
+
+// FreeSlots sums the unused slots over the host's configured models, never counting below zero for
+// one. An unknown host has none.
+func (l *Limiter) FreeSlots(host string) int {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	total := 0
+	for k, p := range l.pairs {
+		if k.host != host || !p.configured {
+			continue
+		}
+		if free := p.parallel - p.inflight; free > 0 {
+			total += free
+		}
+	}
+	return total
+}
@@ -0,0 +1,178 @@
+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 {
+			defer rel()
+			if waited < 40*time.Millisecond {
+				err = errors.New("third acquire did not wait")
+			}
+		}
+		got3 <- err
+	}()
+	time.Sleep(20 * time.Millisecond)
+	if l.Queued("alpha", "m") != 1 {
+		t.Errorf("queued = %d, want 1", l.Queued("alpha", "m"))
+	}
+	// 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(30 * time.Millisecond)
+	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 }()
+	time.Sleep(20 * time.Millisecond)
+	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)
+		time.Sleep(15 * time.Millisecond) // stagger arrivals so the order is defined
+	}
+	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")
+	}
+}