diff --git a/flock.go b/flock.go index 52f119b..16ab4a9 100644 --- a/flock.go +++ b/flock.go @@ -128,6 +128,17 @@ func (f *Flock) pick(nest Nest, n int) []string { return out } +// reserve takes a flock slot, honoring ctx: a saturated flock waits for a +// slot, but gives up when the context ends instead of wedging the caller. +func (f *Flock) reserve(ctx context.Context) (func(), bool) { + select { + case f.sem <- struct{}{}: + return func() { <-f.sem }, true + case <-ctx.Done(): + return nil, false + } +} + // ask is the only way any crow touches the model: budget, flock cap, timeout, log. func (f *Flock) ask(ctx context.Context, w Watcher, crow, system, prompt string, budgets ...*Budget) (string, error) { for _, b := range budgets { @@ -136,8 +147,11 @@ func (f *Flock) ask(ctx context.Context, w Watcher, crow, system, prompt string, return "", err } } - f.sem <- struct{}{} - defer func() { <-f.sem }() + release, slot := f.reserve(ctx) + if !slot { + return "", fmt.Errorf("flock at cap: %w", ctx.Err()) + } + defer release() if w != nil { w.Flying(crow) diff --git a/jev.go b/jev.go index ea7343d..80b8660 100644 --- a/jev.go +++ b/jev.go @@ -200,8 +200,11 @@ func (f *Flock) ClassifyMode(ctx context.Context, text string, budgets ...*Budge return "", false } } - f.sem <- struct{}{} - defer func() { <-f.sem }() + release, slot := f.reserve(ctx) + if !slot { + return "", false + } + defer release() cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout) defer cancel() start := time.Now() @@ -295,8 +298,11 @@ func (f *Flock) ClassifyDivergence(ctx context.Context, state any, budgets ...*B return 0, false } } - f.sem <- struct{}{} - defer func() { <-f.sem }() + release, slot := f.reserve(ctx) + if !slot { + return 0, false + } + defer release() cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout) defer cancel() start := time.Now() @@ -318,8 +324,11 @@ func (f *Flock) ClassifyKind(ctx context.Context, text string, budgets ...*Budge return "", nil, false } } - f.sem <- struct{}{} - defer func() { <-f.sem }() + release, slot := f.reserve(ctx) + if !slot { + return "", nil, false + } + defer release() cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout) defer cancel() start := time.Now() @@ -346,8 +355,11 @@ func (f *Flock) classify(ctx context.Context, w Watcher, item string, budgets .. return JevCall{}, err } } - f.sem <- struct{}{} - defer func() { <-f.sem }() + release, slot := f.reserve(ctx) + if !slot { + return JevCall{}, fmt.Errorf("flock at cap: %w", ctx.Err()) + } + defer release() if w != nil { w.Flying("Sentry") }