74 lines
1.9 KiB
Go
74 lines
1.9 KiB
Go
|
package requestbudget
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
const GraphMaxSends = 3
|
||
|
|
const GraphTimeout = 180 * time.Second
|
||
|
|
const GraphPhase = "graph_extraction"
|
||
|
|
|
||
|
|
// GraphExecution is shared by extraction and its one correction. Callers must
|
||
|
|
// consume a durable single-use receipt before binding this in-memory allowance.
|
||
|
|
type GraphExecution struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
key string
|
||
|
|
model string
|
||
|
|
reserved int
|
||
|
|
stopped error
|
||
|
|
cancel context.CancelCauseFunc
|
||
|
|
}
|
||
|
|
|
||
|
|
func BindGraphExecution(ctx context.Context, key, model string) (context.Context, *GraphExecution, context.CancelFunc, error) {
|
||
|
|
if key == "" || model == "" || FromContext(ctx) != nil {
|
||
|
|
return ctx, nil, func() {}, fmt.Errorf("%w: invalid graph execution scope", ErrStopped)
|
||
|
|
}
|
||
|
|
if err := ctx.Err(); err != nil {
|
||
|
|
return ctx, nil, func() {}, err
|
||
|
|
}
|
||
|
|
ctx, deadlineCancel := context.WithTimeout(ctx, GraphTimeout)
|
||
|
|
ctx, cancel := context.WithCancelCause(ctx)
|
||
|
|
g := &GraphExecution{key: "graph:" + key, model: model, cancel: cancel}
|
||
|
|
ctx = WithPhase(WithGuard(ctx, g), GraphPhase)
|
||
|
|
return ctx, g, func() { deadlineCancel(); cancel(nil) }, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *GraphExecution) Key() string { return g.key }
|
||
|
|
|
||
|
|
func (g *GraphExecution) checkLocked(ctx context.Context) error {
|
||
|
|
if g.stopped != nil {
|
||
|
|
return g.stopped
|
||
|
|
}
|
||
|
|
return ctx.Err()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *GraphExecution) Check(ctx context.Context) error {
|
||
|
|
g.mu.Lock()
|
||
|
|
defer g.mu.Unlock()
|
||
|
|
return g.checkLocked(ctx)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *GraphExecution) Reserve(ctx context.Context, model, phase string) error {
|
||
|
|
g.mu.Lock()
|
||
|
|
defer g.mu.Unlock()
|
||
|
|
if err := g.checkLocked(ctx); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if model != g.model || phase != GraphPhase || g.reserved >= GraphMaxSends {
|
||
|
|
g.stopped = fmt.Errorf("%w: graph send limit or scope mismatch", ErrStopped)
|
||
|
|
g.cancel(g.stopped)
|
||
|
|
return g.stopped
|
||
|
|
}
|
||
|
|
g.reserved++
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *GraphExecution) Reservations() int {
|
||
|
|
g.mu.Lock()
|
||
|
|
defer g.mu.Unlock()
|
||
|
|
return g.reserved
|
||
|
|
}
|