49 lines
1.1 KiB
Go
49 lines
1.1 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
)
|
|
|
|
// Scheduler bounds concurrent backend LLM calls.
|
|
type Scheduler struct {
|
|
permits chan struct{}
|
|
}
|
|
|
|
// 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{
|
|
permits: make(chan struct{}, 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) {
|
|
select {
|
|
case s.permits <- struct{}{}:
|
|
var once sync.Once
|
|
return func() {
|
|
once.Do(func() {
|
|
<-s.permits
|
|
})
|
|
}, nil
|
|
case <-ctx.Done():
|
|
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 {
|
|
release, err := s.Acquire(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer release()
|
|
return fn(ctx)
|
|
}
|