RC-4

internal/workqueue/limiter.go

@@ -10,19 +10,19 @@ import (
 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.
+// 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 {
-	slots chan struct{}
-
-	mu         sync.Mutex
-	maxWaiting int
-	waiting    int
-	active     int
-	admitted   uint64
-	rejected   uint64
-	kinds      map[string]*kindCounters
+	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 {
@@ -56,7 +56,30 @@ func New(maxInflight, queueSize int) *Limiter {
 	if queueSize < 1 {
 		queueSize = 1
 	}
-	return &Limiter{slots: make(chan struct{}, maxInflight), maxWaiting: queueSize, kinds: map[string]*kindCounters{}}
+	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) {
@@ -68,56 +91,55 @@ func (l *Limiter) AcquireKind(ctx context.Context, kind string) (func(), error)
 		return func() {}, nil
 	}
 	kind = normalizeKind(kind)
-
-	select {
-	case l.slots <- struct{}{}:
+	registeredWaiting := false
+	for {
 		l.mu.Lock()
-		l.active++
-		l.admitted++
-		c := l.kindLocked(kind)
-		c.active++
-		c.admitted++
+		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()
-		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--
+		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:
 		}
-		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()
 	}
 }
 
@@ -125,7 +147,6 @@ 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--
@@ -134,11 +155,21 @@ func (l *Limiter) releaseFunc(kind string) func() {
 			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{}
@@ -153,7 +184,7 @@ func (l *Limiter) Status() Status {
 		kinds[kind] = KindStatus{Active: counters.active, Waiting: counters.waiting, Admitted: counters.admitted, Rejected: counters.rejected}
 	}
 	return Status{
-		MaxInflight: cap(l.slots),
+		MaxInflight: l.maxInflight,
 		QueueSize:   l.maxWaiting,
 		Active:      l.active,
 		Waiting:     l.waiting,

internal/workqueue/limiter_test.go

@@ -64,3 +64,33 @@ func TestLimiterTracksKindsWithoutDuplicateCounters(t *testing.T) {
 		t.Fatalf("release duplicated or leaked active counters: %+v", status)
 	}
 }
+
+func TestLimiterCanRaiseRuntimeLimit(t *testing.T) {
+	l := New(1, 4)
+	release1, err := l.Acquire(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	acquired := make(chan func(), 1)
+	go func() {
+		release, err := l.Acquire(context.Background())
+		if err == nil {
+			acquired <- release
+		}
+	}()
+	time.Sleep(10 * time.Millisecond)
+	if got := l.Status(); got.Active != 1 || got.Waiting != 1 {
+		t.Fatalf("unexpected before resize: %+v", got)
+	}
+	l.SetLimits(2, 8)
+	select {
+	case release2 := <-acquired:
+		release2()
+	case <-time.After(time.Second):
+		t.Fatal("raising limiter did not wake waiter")
+	}
+	release1()
+	if got := l.Status(); got.MaxInflight != 2 || got.QueueSize != 8 || got.Active != 0 {
+		t.Fatalf("unexpected resized status: %+v", got)
+	}
+}