glpi-neural-brain/internal/workqueue/limiter.go

package workqueue

import (
	"context"
	"errors"
	"strings"
	"sync"
)

var ErrQueueFull = errors.New("shared research/ollama queue is full")

// Limiter bounds concurrent expensive/outbound work and the number of callers
// waiting for a slot. Unlike a channel semaphore the limits are runtime
// adjustable, which lets the Brain switch between conservative and Speed mode
// without restarting workers or losing queue accounting.
type Limiter struct {
	mu          sync.Mutex
	maxInflight int
	maxWaiting  int
	waiting     int
	active      int
	admitted    uint64
	rejected    uint64
	kinds       map[string]*kindCounters
	changed     chan struct{}
}

type kindCounters struct {
	active   int
	waiting  int
	admitted uint64
	rejected uint64
}

type KindStatus struct {
	Active   int    `json:"active"`
	Waiting  int    `json:"waiting"`
	Admitted uint64 `json:"admitted"`
	Rejected uint64 `json:"rejected"`
}

type Status struct {
	MaxInflight int                   `json:"max_inflight"`
	QueueSize   int                   `json:"queue_size"`
	Active      int                   `json:"active"`
	Waiting     int                   `json:"waiting"`
	Admitted    uint64                `json:"admitted"`
	Rejected    uint64                `json:"rejected"`
	Kinds       map[string]KindStatus `json:"kinds,omitempty"`
}

func New(maxInflight, queueSize int) *Limiter {
	if maxInflight < 1 {
		maxInflight = 1
	}
	if queueSize < 1 {
		queueSize = 1
	}
	return &Limiter{maxInflight: maxInflight, maxWaiting: queueSize, kinds: map[string]*kindCounters{}, changed: make(chan struct{})}
}

// SetLimits changes the live concurrency/queue limits. Existing active work is
// never cancelled when shrinking; new callers wait until active drops below the
// new ceiling. Raising the limit wakes all waiters immediately.
func (l *Limiter) SetLimits(maxInflight, queueSize int) {
	if l == nil {
		return
	}
	if maxInflight < 1 {
		maxInflight = 1
	}
	if queueSize < 1 {
		queueSize = 1
	}
	l.mu.Lock()
	changed := l.maxInflight != maxInflight || l.maxWaiting != queueSize
	l.maxInflight = maxInflight
	l.maxWaiting = queueSize
	if changed {
		l.notifyLocked()
	}
	l.mu.Unlock()
}

func (l *Limiter) Acquire(ctx context.Context) (func(), error) {
	return l.AcquireKind(ctx, "unspecified")
}

func (l *Limiter) AcquireKind(ctx context.Context, kind string) (func(), error) {
	if l == nil {
		return func() {}, nil
	}
	kind = normalizeKind(kind)
	registeredWaiting := false
	for {
		l.mu.Lock()
		if l.active < l.maxInflight {
			if registeredWaiting {
				l.waiting--
				c := l.kindLocked(kind)
				if c.waiting > 0 {
					c.waiting--
				}
			}
			l.active++
			l.admitted++
			c := l.kindLocked(kind)
			c.active++
			c.admitted++
			l.mu.Unlock()
			return l.releaseFunc(kind), nil
		}
		if !registeredWaiting {
			if l.waiting >= l.maxWaiting {
				l.rejected++
				l.kindLocked(kind).rejected++
				l.mu.Unlock()
				return nil, ErrQueueFull
			}
			l.waiting++
			l.kindLocked(kind).waiting++
			registeredWaiting = true
		}
		changed := l.changed
		l.mu.Unlock()

		select {
		case <-ctx.Done():
			l.mu.Lock()
			if registeredWaiting {
				if l.waiting > 0 {
					l.waiting--
				}
				c := l.kindLocked(kind)
				if c.waiting > 0 {
					c.waiting--
				}
			}
			l.mu.Unlock()
			return nil, ctx.Err()
		case <-changed:
		}
	}
}

func (l *Limiter) releaseFunc(kind string) func() {
	var once sync.Once
	return func() {
		once.Do(func() {
			l.mu.Lock()
			if l.active > 0 {
				l.active--
			}
			c := l.kindLocked(kind)
			if c.active > 0 {
				c.active--
			}
			l.notifyLocked()
			l.mu.Unlock()
		})
	}
}

func (l *Limiter) notifyLocked() {
	if l.changed == nil {
		l.changed = make(chan struct{})
		return
	}
	close(l.changed)
	l.changed = make(chan struct{})
}

func (l *Limiter) Status() Status {
	if l == nil {
		return Status{}
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	kinds := make(map[string]KindStatus, len(l.kinds))
	for kind, counters := range l.kinds {
		if counters.active == 0 && counters.waiting == 0 && counters.admitted == 0 && counters.rejected == 0 {
			continue
		}
		kinds[kind] = KindStatus{Active: counters.active, Waiting: counters.waiting, Admitted: counters.admitted, Rejected: counters.rejected}
	}
	return Status{
		MaxInflight: l.maxInflight,
		QueueSize:   l.maxWaiting,
		Active:      l.active,
		Waiting:     l.waiting,
		Admitted:    l.admitted,
		Rejected:    l.rejected,
		Kinds:       kinds,
	}
}

func (l *Limiter) kindLocked(kind string) *kindCounters {
	if l.kinds == nil {
		l.kinds = map[string]*kindCounters{}
	}
	c := l.kinds[kind]
	if c == nil {
		c = &kindCounters{}
		l.kinds[kind] = c
	}
	return c
}

func normalizeKind(kind string) string {
	kind = strings.ToLower(strings.TrimSpace(kind))
	if kind == "" {
		return "unspecified"
	}
	return kind
}