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. Callers can attach a kind so status output distinguishes
// SearXNG, web fetches, Ollama chat and embeddings instead of exposing one
// opaque aggregate counter.
type Limiter struct {
	slots chan struct{}

	mu         sync.Mutex
	maxWaiting int
	waiting    int
	active     int
	admitted   uint64
	rejected   uint64
	kinds      map[string]*kindCounters
}

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{slots: make(chan struct{}, maxInflight), maxWaiting: queueSize, kinds: map[string]*kindCounters{}}
}

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)

	select {
	case l.slots <- struct{}{}:
		l.mu.Lock()
		l.active++
		l.admitted++
		c := l.kindLocked(kind)
		c.active++
		c.admitted++
		l.mu.Unlock()
		return l.releaseFunc(kind), nil
	default:
	}

	l.mu.Lock()
	if l.waiting >= l.maxWaiting {
		l.rejected++
		l.kindLocked(kind).rejected++
		l.mu.Unlock()
		return nil, ErrQueueFull
	}
	l.waiting++
	l.kindLocked(kind).waiting++
	l.mu.Unlock()

	select {
	case l.slots <- struct{}{}:
		l.mu.Lock()
		l.waiting--
		l.active++
		l.admitted++
		c := l.kindLocked(kind)
		if c.waiting > 0 {
			c.waiting--
		}
		c.active++
		c.admitted++
		l.mu.Unlock()
		return l.releaseFunc(kind), nil
	case <-ctx.Done():
		l.mu.Lock()
		if l.waiting > 0 {
			l.waiting--
		}
		c := l.kindLocked(kind)
		if c.waiting > 0 {
			c.waiting--
		}
		l.mu.Unlock()
		return nil, ctx.Err()
	}
}

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

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: cap(l.slots),
		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
}