113 lines
2.0 KiB
Go
113 lines
2.0 KiB
Go
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,
|
|
}
|
|
}
|