Add backend capacity manager

internal/capacity/manager_test.go

@@ -0,0 +1,158 @@
+package capacity
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
+)
+
+func TestNewManagerRejectsInvalidPolicies(t *testing.T) {
+	maxInt := int(^uint(0) >> 1)
+	tests := []struct {
+		name   string
+		id     string
+		policy domain.BackendCapacityPolicy
+	}{
+		{
+			name:   "blank ID",
+			id:     " \t ",
+			policy: domain.BackendCapacityPolicy{ConcurrencyLimit: 1},
+		},
+		{
+			name:   "zero concurrency",
+			id:     "backend",
+			policy: domain.BackendCapacityPolicy{},
+		},
+		{
+			name:   "negative concurrency",
+			id:     "backend",
+			policy: domain.BackendCapacityPolicy{ConcurrencyLimit: -1},
+		},
+		{
+			name: "negative queue",
+			id:   "backend",
+			policy: domain.BackendCapacityPolicy{
+				ConcurrencyLimit: 1,
+				QueueCapacity:    -1,
+			},
+		},
+		{
+			name: "total overflow",
+			id:   "backend",
+			policy: domain.BackendCapacityPolicy{
+				ConcurrencyLimit: maxInt,
+				QueueCapacity:    1,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := NewManager(map[string]domain.BackendCapacityPolicy{
+				tc.id: tc.policy,
+			})
+			if err == nil {
+				t.Fatal("expected invalid policy error")
+			}
+		})
+	}
+}
+
+func TestManagerAdmissionIsBoundedAndReleaseIsIdempotent(t *testing.T) {
+	policies := map[string]domain.BackendCapacityPolicy{
+		"limited": {
+			ConcurrencyLimit: 2,
+			QueueCapacity:    1,
+		},
+		"independent": {
+			ConcurrencyLimit: 1,
+		},
+	}
+	manager, err := NewManager(policies)
+	if err != nil {
+		t.Fatalf("construct manager: %v", err)
+	}
+	policies["limited"] = domain.BackendCapacityPolicy{
+		ConcurrencyLimit: 100,
+		QueueCapacity:    100,
+	}
+
+	releases := make([]func(), 0, 3)
+	for range 3 {
+		release, err := manager.Admit(context.Background(), "limited")
+		if err != nil {
+			t.Fatalf("admit within configured capacity: %v", err)
+		}
+		releases = append(releases, release)
+	}
+	if release, err := manager.Admit(context.Background(), "limited"); release != nil ||
+		!errors.Is(err, ErrCapacityExceeded) {
+		t.Fatalf("admission beyond capacity=(release=%t, err=%v), want ErrCapacityExceeded",
+			release != nil, err)
+	}
+	independentRelease, err := manager.Admit(context.Background(), "independent")
+	if err != nil {
+		t.Fatalf("admit independent backend while first is full: %v", err)
+	}
+	independentRelease()
+
+	releases[0]()
+	releases[0]()
+	replacement, err := manager.Admit(context.Background(), "limited")
+	if err != nil {
+		t.Fatalf("admit after release: %v", err)
+	}
+	replacement()
+	releases[1]()
+	releases[2]()
+
+	pool := manager.pools["limited"]
+	pool.mu.Lock()
+	admitted := pool.admitted
+	pool.mu.Unlock()
+	if admitted != 0 {
+		t.Fatalf("admitted runs after releases=%d, want 0", admitted)
+	}
+}
+
+func TestManagerAdmissionHonorsContextAndUnlimitedBackends(t *testing.T) {
+	manager, err := NewManager(map[string]domain.BackendCapacityPolicy{
+		"limited": {ConcurrencyLimit: 1},
+	})
+	if err != nil {
+		t.Fatalf("construct manager: %v", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if release, err := manager.Admit(ctx, "limited"); release != nil ||
+		!errors.Is(err, context.Canceled) {
+		t.Fatalf("canceled limited admission=(release=%t, err=%v), want context cancellation",
+			release != nil, err)
+	}
+
+	var nilManager *Manager
+	for _, tc := range []struct {
+		name      string
+		manager   *Manager
+		backendID string
+	}{
+		{name: "nil manager", manager: nilManager, backendID: "limited"},
+		{name: "blank ID", manager: manager},
+		{name: "unknown ID", manager: manager, backendID: "unknown"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			release, err := tc.manager.Admit(ctx, tc.backendID)
+			if err != nil {
+				t.Fatalf("unlimited admission: %v", err)
+			}
+			if release == nil {
+				t.Fatal("unlimited admission returned nil release")
+			}
+			release()
+			release()
+		})
+	}
+}