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

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)

internal/choose/choose.go

@@ -0,0 +1,63 @@
+// Package choose picks the host for a new lease: among the healthy, non-draining, known hosts it
+// prefers those that have the model loaded over those that would only be able to serve it, then the
+// one with the most free slots times weight, breaking ties by shortest queue and finally list
+// order.
+package choose
+
+// Info is one candidate host's view of itself for a single model.
+type Info struct {
+	Healthy  bool // answered its last poll
+	Draining bool // operator is draining it; no new leases
+	Loaded   bool // the model is resident here
+	CanServe bool // config lists the model, so we may load it
+	Free     int
+	Queued   int
+	Weight   float64
+}
+
+// Best returns the best host for a new lease, or ok=false when nothing is eligible. It runs two
+// passes over the candidates in order: the first over those that have the model loaded, the second,
+// only if the first found nothing, over those that can serve it. A host must be healthy, not
+// draining, and known (info reported ok) to be eligible in either pass; a host with zero free slots
+// is still eligible, since it will queue. Among the eligible ones the highest Free*Weight wins, ties
+// go to the lowest queue, and a remaining tie keeps the earlier candidate.
+func Best(candidates []string, info func(host string) (Info, bool)) (string, bool) {
+	const (
+		passLoaded   = 0
+		passCanServe = 1
+	)
+	best := ""
+	var (
+		bestScore  float64
+		bestQueued int
+		found      bool
+	)
+	for pass := passLoaded; pass <= passCanServe; pass++ {
+		for _, host := range candidates {
+			v, ok := info(host)
+			if !ok || !v.Healthy || v.Draining {
+				continue
+			}
+			switch pass {
+			case passLoaded:
+				if !v.Loaded {
+					continue
+				}
+			case passCanServe:
+				if !v.CanServe {
+					continue
+				}
+			}
+			score := float64(v.Free) * v.Weight
+			// Replace only when strictly better on score, or equal score with a shorter queue; a
+			// further tie keeps the earlier candidate because we scan in order and use "<".
+			if !found || score > bestScore || (score == bestScore && v.Queued < bestQueued) {
+				best, bestScore, bestQueued, found = host, score, v.Queued, true
+			}
+		}
+		if found {
+			break
+		}
+	}
+	return best, found
+}

internal/choose/choose_test.go

@@ -0,0 +1,83 @@
+package choose_test
+
+import (
+	"testing"
+
+	"git.wntrmute.dev/kyle/crossbar/internal/choose"
+)
+
+func infoFor(m map[string]choose.Info) func(string) (choose.Info, bool) {
+	return func(name string) (choose.Info, bool) { i, ok := m[name]; return i, ok }
+}
+
+func TestMostFreeSlotsTimesWeightWins(t *testing.T) {
+	info := infoFor(map[string]choose.Info{
+		"alpha": {Healthy: true, Loaded: true, CanServe: true, Free: 3, Weight: 1.0},
+		"beta":  {Healthy: true, Loaded: true, CanServe: true, Free: 2, Weight: 2.0}, // 4 > 3
+		"gamma": {Healthy: true, Loaded: true, CanServe: true, Free: 4, Weight: 0.5}, // 2
+	})
+	got, ok := choose.Best([]string{"alpha", "beta", "gamma"}, info)
+	if !ok || got != "beta" {
+		t.Errorf("got %q %v, want beta", got, ok)
+	}
+}
+
+func TestTieGoesToShortestQueueThenListOrder(t *testing.T) {
+	info := infoFor(map[string]choose.Info{
+		"alpha": {Healthy: true, Loaded: true, CanServe: true, Free: 2, Weight: 1, Queued: 3},
+		"beta":  {Healthy: true, Loaded: true, CanServe: true, Free: 2, Weight: 1, Queued: 1},
+		"gamma": {Healthy: true, Loaded: true, CanServe: true, Free: 2, Weight: 1, Queued: 1},
+	})
+	if got, _ := choose.Best([]string{"alpha", "beta", "gamma"}, info); got != "beta" {
+		t.Errorf("tie on score: shortest queue wins, then list order; got %q", got)
+	}
+	if got, _ := choose.Best([]string{"gamma", "beta"}, info); got != "gamma" {
+		t.Errorf("full tie: first in list order wins; got %q", got)
+	}
+}
+
+func TestLoadedBeatsMerelyCapable(t *testing.T) {
+	info := infoFor(map[string]choose.Info{
+		"alpha": {Healthy: true, Loaded: false, CanServe: true, Free: 8, Weight: 4},
+		"beta":  {Healthy: true, Loaded: true, CanServe: true, Free: 1, Weight: 1},
+	})
+	got, ok := choose.Best([]string{"alpha", "beta"}, info)
+	if !ok || got != "beta" {
+		t.Errorf("a host that has the model loaded wins over one that would have to load it; got %q", got)
+	}
+}
+
+func TestFallsBackToCapableHost(t *testing.T) {
+	info := infoFor(map[string]choose.Info{
+		"alpha": {Healthy: true, Loaded: false, CanServe: true, Free: 1, Weight: 1},
+		"beta":  {Healthy: true, Loaded: false, CanServe: false, Free: 9, Weight: 9},
+	})
+	got, ok := choose.Best([]string{"beta", "alpha"}, info)
+	if !ok || got != "alpha" {
+		t.Errorf("only a host configured to serve the model may load it; got %q %v", got, ok)
+	}
+}
+
+func TestSkipsUnhealthyDrainingUnknownAndFull(t *testing.T) {
+	info := infoFor(map[string]choose.Info{
+		"down":  {Healthy: false, Loaded: true, CanServe: true, Free: 9, Weight: 9},
+		"drain": {Healthy: true, Draining: true, Loaded: true, CanServe: true, Free: 9, Weight: 9},
+		"full":  {Healthy: true, Loaded: true, CanServe: true, Free: 0, Weight: 9, Queued: 0},
+		"ok":    {Healthy: true, Loaded: true, CanServe: true, Free: 1, Weight: 1},
+	})
+	got, ok := choose.Best([]string{"down", "drain", "missing", "full", "ok"}, info)
+	if !ok || got != "ok" {
+		t.Errorf("got %q %v, want ok", got, ok)
+	}
+	// A full host is still better than nothing: it gets the request (it will queue).
+	got, ok = choose.Best([]string{"down", "full"}, info)
+	if !ok || got != "full" {
+		t.Errorf("with only a full host left it must still be chosen; got %q %v", got, ok)
+	}
+	if _, ok := choose.Best([]string{"down", "drain", "missing"}, info); ok {
+		t.Errorf("nothing usable must give ok=false")
+	}
+	if _, ok := choose.Best(nil, info); ok {
+		t.Errorf("empty candidates must give ok=false")
+	}
+}

internal/limiter/limiter.go

@@ -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
+}

internal/limiter/limiter_test.go

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