promptkit/internal/capacity/manager.go

// 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)
	}
}