Ch1

internal/workqueue/limiter.go

@@ -0,0 +1,112 @@
+package workqueue
+
+import (
+	"context"
+	"errors"
+	"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.
+type Limiter struct {
+	slots chan struct{}
+
+	mu         sync.Mutex
+	maxWaiting int
+	waiting    int
+	active     int
+	admitted   uint64
+	rejected   uint64
+}
+
+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"`
+}
+
+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}
+}
+
+func (l *Limiter) Acquire(ctx context.Context) (func(), error) {
+	if l == nil {
+		return func() {}, nil
+	}
+
+	select {
+	case l.slots <- struct{}{}:
+		l.mu.Lock()
+		l.active++
+		l.admitted++
+		l.mu.Unlock()
+		return l.releaseFunc(), nil
+	default:
+	}
+
+	l.mu.Lock()
+	if l.waiting >= l.maxWaiting {
+		l.rejected++
+		l.mu.Unlock()
+		return nil, ErrQueueFull
+	}
+	l.waiting++
+	l.mu.Unlock()
+
+	select {
+	case l.slots <- struct{}{}:
+		l.mu.Lock()
+		l.waiting--
+		l.active++
+		l.admitted++
+		l.mu.Unlock()
+		return l.releaseFunc(), nil
+	case <-ctx.Done():
+		l.mu.Lock()
+		l.waiting--
+		l.mu.Unlock()
+		return nil, ctx.Err()
+	}
+}
+
+func (l *Limiter) releaseFunc() func() {
+	var once sync.Once
+	return func() {
+		once.Do(func() {
+			<-l.slots
+			l.mu.Lock()
+			if l.active > 0 {
+				l.active--
+			}
+			l.mu.Unlock()
+		})
+	}
+}
+
+func (l *Limiter) Status() Status {
+	if l == nil {
+		return Status{}
+	}
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return Status{
+		MaxInflight: cap(l.slots),
+		QueueSize:   l.maxWaiting,
+		Active:      l.active,
+		Waiting:     l.waiting,
+		Admitted:    l.admitted,
+		Rejected:    l.rejected,
+	}
+}

internal/workqueue/limiter_test.go

@@ -0,0 +1,43 @@
+package workqueue
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestLimiterBoundsQueue(t *testing.T) {
+	l := New(1, 1)
+	release, err := l.Acquire(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer release()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	waiterDone := make(chan error, 1)
+	go func() {
+		r, err := l.Acquire(ctx)
+		if err == nil {
+			r()
+		}
+		waiterDone <- err
+	}()
+
+	deadline := time.Now().Add(time.Second)
+	for l.Status().Waiting != 1 && time.Now().Before(deadline) {
+		time.Sleep(time.Millisecond)
+	}
+	if l.Status().Waiting != 1 {
+		t.Fatalf("expected one waiter: %+v", l.Status())
+	}
+	if _, err := l.Acquire(context.Background()); !errors.Is(err, ErrQueueFull) {
+		t.Fatalf("expected ErrQueueFull, got %v", err)
+	}
+	release()
+	if err := <-waiterDone; err != nil {
+		t.Fatalf("waiter failed: %v", err)
+	}
+}