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