Add backend capacity manager

internal/capacity/manager.go

@@ -0,0 +1,160 @@
+// Package capacity coordinates engine-local run admission and model-generation
+// concurrency for configured backends.
+package capacity
+
+import (
+	"container/list"
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"sync"
+
+	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
+)
+
+// ErrCapacityExceeded identifies an admission rejected because a backend's
+// configured run capacity is full.
+var ErrCapacityExceeded = errors.New("backend capacity exceeded")
+
+// Manager owns independent backend capacity pools with immutable limits.
+type Manager struct {
+	pools map[string]*pool
+}
+
+type pool struct {
+	mu               sync.Mutex
+	concurrencyLimit int
+	totalCapacity    int
+	admitted         int
+	active           int
+	waiters          list.List
+}
+
+type waiter struct {
+	ready   chan struct{}
+	element *list.Element
+	granted bool
+}
+
+// NewManager constructs independent pools from normalized backend policies.
+func NewManager(policies map[string]domain.BackendCapacityPolicy) (*Manager, error) {
+	manager := &Manager{
+		pools: make(map[string]*pool, len(policies)),
+	}
+	maxInt := int(^uint(0) >> 1)
+	for id, policy := range policies {
+		if strings.TrimSpace(id) == "" {
+			return nil, errors.New("backend capacity policy ID must not be blank")
+		}
+		if policy.ConcurrencyLimit <= 0 {
+			return nil, fmt.Errorf(
+				"backend %q concurrency limit must be positive",
+				id,
+			)
+		}
+		if policy.QueueCapacity < 0 {
+			return nil, fmt.Errorf(
+				"backend %q queue capacity must not be negative",
+				id,
+			)
+		}
+		if policy.QueueCapacity > maxInt-policy.ConcurrencyLimit {
+			return nil, fmt.Errorf("backend %q total capacity overflows int", id)
+		}
+		manager.pools[id] = &pool{
+			concurrencyLimit: policy.ConcurrencyLimit,
+			totalCapacity:    policy.ConcurrencyLimit + policy.QueueCapacity,
+		}
+	}
+	return manager, nil
+}
+
+// Admit immediately reserves one configured backend run slot. Backends without
+// a configured pool are unlimited.
+func (m *Manager) Admit(ctx context.Context, backendID string) (func(), error) {
+	pool := m.getPool(backendID)
+	if pool == nil {
+		return releaseNothing, nil
+	}
+
+	pool.mu.Lock()
+	defer pool.mu.Unlock()
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	if pool.admitted >= pool.totalCapacity {
+		return nil, ErrCapacityExceeded
+	}
+	pool.admitted++
+
+	var once sync.Once
+	return func() {
+		once.Do(func() {
+			pool.mu.Lock()
+			pool.admitted--
+			pool.mu.Unlock()
+		})
+	}, nil
+}
+
+func releaseNothing() {}
+
+func (m *Manager) getPool(backendID string) *pool {
+	if m == nil || backendID == "" {
+		return nil
+	}
+	return m.pools[backendID]
+}
+
+func (p *pool) acquire(ctx context.Context) error {
+	p.mu.Lock()
+	if err := ctx.Err(); err != nil {
+		p.mu.Unlock()
+		return err
+	}
+	if p.active < p.concurrencyLimit && p.waiters.Len() == 0 {
+		p.active++
+		p.mu.Unlock()
+		return nil
+	}
+
+	waiter := &waiter{ready: make(chan struct{})}
+	waiter.element = p.waiters.PushBack(waiter)
+	p.mu.Unlock()
+
+	select {
+	case <-waiter.ready:
+		return nil
+	case <-ctx.Done():
+		p.mu.Lock()
+		if !waiter.granted {
+			p.waiters.Remove(waiter.element)
+			waiter.element = nil
+			p.mu.Unlock()
+			return ctx.Err()
+		}
+		p.mu.Unlock()
+		return nil
+	}
+}
+
+func (p *pool) releaseActive() {
+	var ready chan struct{}
+
+	p.mu.Lock()
+	if element := p.waiters.Front(); element != nil {
+		waiter := element.Value.(*waiter)
+		p.waiters.Remove(element)
+		waiter.element = nil
+		waiter.granted = true
+		ready = waiter.ready
+	} else {
+		p.active--
+	}
+	p.mu.Unlock()
+
+	if ready != nil {
+		close(ready)
+	}
+}