126 lines
2.4 KiB
Go
126 lines
2.4 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
)
|
|
|
|
// Scheduler bounds concurrent LLM backend calls.
|
|
type Scheduler struct {
|
|
maxConcurrency int
|
|
mu sync.Mutex
|
|
inFlight int
|
|
queue []*waiter
|
|
}
|
|
|
|
type waiter struct {
|
|
ready chan struct{}
|
|
queued bool
|
|
granted bool
|
|
}
|
|
|
|
// NewScheduler creates a scheduler with a fixed concurrency limit.
|
|
func NewScheduler(maxConcurrency int) (*Scheduler, error) {
|
|
if maxConcurrency <= 0 {
|
|
return nil, fmt.Errorf("max concurrency must be greater than zero")
|
|
}
|
|
return &Scheduler{
|
|
maxConcurrency: maxConcurrency,
|
|
}, nil
|
|
}
|
|
|
|
// Acquire blocks until a permit is available or the context is canceled.
|
|
// The returned release function is safe to call multiple times.
|
|
func (s *Scheduler) Acquire(ctx context.Context) (func(), error) {
|
|
if s == nil {
|
|
return nil, fmt.Errorf("scheduler must not be nil")
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
s.mu.Lock()
|
|
if s.inFlight < s.maxConcurrency && len(s.queue) == 0 {
|
|
s.inFlight++
|
|
s.mu.Unlock()
|
|
return s.releaseFunc(), nil
|
|
}
|
|
|
|
w := &waiter{
|
|
ready: make(chan struct{}),
|
|
queued: true,
|
|
}
|
|
s.queue = append(s.queue, w)
|
|
s.mu.Unlock()
|
|
|
|
select {
|
|
case <-w.ready:
|
|
return s.releaseFunc(), nil
|
|
case <-ctx.Done():
|
|
s.mu.Lock()
|
|
if w.queued {
|
|
s.removeQueuedWaiterLocked(w)
|
|
s.mu.Unlock()
|
|
return nil, ctx.Err()
|
|
}
|
|
if w.granted {
|
|
s.inFlight--
|
|
s.grantQueuedLocked()
|
|
}
|
|
s.mu.Unlock()
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
|
|
// Run acquires a permit, executes fn, and releases the permit.
|
|
func (s *Scheduler) Run(ctx context.Context, fn func(context.Context) error) error {
|
|
if fn == nil {
|
|
return fmt.Errorf("scheduler function must not be nil")
|
|
}
|
|
release, err := s.Acquire(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer release()
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
return fn(ctx)
|
|
}
|
|
|
|
func (s *Scheduler) releaseFunc() func() {
|
|
var once sync.Once
|
|
return func() {
|
|
once.Do(func() {
|
|
s.mu.Lock()
|
|
if s.inFlight > 0 {
|
|
s.inFlight--
|
|
s.grantQueuedLocked()
|
|
}
|
|
s.mu.Unlock()
|
|
})
|
|
}
|
|
}
|
|
|
|
func (s *Scheduler) grantQueuedLocked() {
|
|
for s.inFlight < s.maxConcurrency && len(s.queue) > 0 {
|
|
w := s.queue[0]
|
|
s.queue = s.queue[1:]
|
|
w.queued = false
|
|
w.granted = true
|
|
s.inFlight++
|
|
close(w.ready)
|
|
}
|
|
}
|
|
|
|
func (s *Scheduler) removeQueuedWaiterLocked(target *waiter) {
|
|
for i, w := range s.queue {
|
|
if w == target {
|
|
w.queued = false
|
|
s.queue = append(s.queue[:i], s.queue[i+1:]...)
|
|
return
|
|
}
|
|
}
|
|
}
|