Ch2

internal/workqueue/limiter.go

@@ -3,14 +3,16 @@ 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. It is intentionally small and dependency-free so the
-// same limiter can be shared by SearXNG/fetch work and Ollama calls.
+// 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{}
 
@@ -20,15 +22,31 @@ type Limiter struct {
 	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"`
+	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 {
@@ -38,31 +56,41 @@ func New(maxInflight, queueSize int) *Limiter {
 	if queueSize < 1 {
 		queueSize = 1
 	}
-	return &Limiter{slots: make(chan struct{}, maxInflight), maxWaiting: queueSize}
+	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(), nil
+		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 {
@@ -71,17 +99,29 @@ func (l *Limiter) Acquire(ctx context.Context) (func(), error) {
 		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(), nil
+		return l.releaseFunc(kind), nil
 	case <-ctx.Done():
 		l.mu.Lock()
-		l.waiting--
+		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() func() {
+func (l *Limiter) releaseFunc(kind string) func() {
 	var once sync.Once
 	return func() {
 		once.Do(func() {
@@ -90,6 +130,10 @@ func (l *Limiter) releaseFunc() func() {
 			if l.active > 0 {
 				l.active--
 			}
+			c := l.kindLocked(kind)
+			if c.active > 0 {
+				c.active--
+			}
 			l.mu.Unlock()
 		})
 	}
@@ -101,6 +145,13 @@ func (l *Limiter) Status() 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,
@@ -108,5 +159,26 @@ func (l *Limiter) Status() Status {
 		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
+}

internal/workqueue/limiter_test.go

@@ -41,3 +41,26 @@ func TestLimiterBoundsQueue(t *testing.T) {
 		t.Fatalf("waiter failed: %v", err)
 	}
 }
+
+func TestLimiterTracksKindsWithoutDuplicateCounters(t *testing.T) {
+	l := New(2, 2)
+	releaseSearch, err := l.AcquireKind(context.Background(), "searxng.search")
+	if err != nil {
+		t.Fatal(err)
+	}
+	releaseChat, err := l.AcquireKind(context.Background(), "ollama.chat")
+	if err != nil {
+		t.Fatal(err)
+	}
+	status := l.Status()
+	if status.Active != 2 || status.Kinds["searxng.search"].Active != 1 || status.Kinds["ollama.chat"].Active != 1 {
+		t.Fatalf("unexpected typed queue status: %+v", status)
+	}
+	releaseSearch()
+	releaseSearch() // release must be idempotent
+	releaseChat()
+	status = l.Status()
+	if status.Active != 0 || status.Kinds["searxng.search"].Active != 0 || status.Kinds["ollama.chat"].Active != 0 {
+		t.Fatalf("release duplicated or leaked active counters: %+v", status)
+	}
+}