74 lines
2.1 KiB
Go
74 lines
2.1 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
"strconv"
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/requestbudget"
|
||
|
|
)
|
||
|
|
|
||
|
|
func probeSendLimit(requested *int, environment string, envSet bool) (int, error) {
|
||
|
|
limit := requestbudget.GraphMaxSends
|
||
|
|
if requested != nil {
|
||
|
|
limit = *requested
|
||
|
|
}
|
||
|
|
if limit < 1 || limit > requestbudget.GraphMaxSends {
|
||
|
|
return 0, fmt.Errorf("invalid probe HTTP send cap")
|
||
|
|
}
|
||
|
|
if envSet && environment != strconv.Itoa(limit) {
|
||
|
|
return 0, fmt.Errorf("probe HTTP send cap differs between request and environment")
|
||
|
|
}
|
||
|
|
return limit, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// This tighter probe allowance delegates model, phase and deadline enforcement
|
||
|
|
// to the existing graph guard; worker construction and defaults are untouched.
|
||
|
|
type probeExecutionBudget struct {
|
||
|
|
*requestbudget.GraphExecution
|
||
|
|
mu sync.Mutex
|
||
|
|
limit int
|
||
|
|
stopped error
|
||
|
|
cancel context.CancelCauseFunc
|
||
|
|
}
|
||
|
|
|
||
|
|
func bindProbeExecution(ctx context.Context, key, model string, limit int) (context.Context, *probeExecutionBudget, context.CancelFunc, error) {
|
||
|
|
if _, err := probeSendLimit(&limit, "", false); err != nil {
|
||
|
|
return ctx, nil, func() {}, err
|
||
|
|
}
|
||
|
|
ctx, base, baseCancel, err := requestbudget.BindGraphExecution(ctx, key, model)
|
||
|
|
if err != nil {
|
||
|
|
return ctx, nil, baseCancel, err
|
||
|
|
}
|
||
|
|
ctx, cancel := context.WithCancelCause(ctx)
|
||
|
|
guard := &probeExecutionBudget{GraphExecution: base, limit: limit, cancel: cancel}
|
||
|
|
return requestbudget.WithGuard(ctx, guard), guard, func() { cancel(nil); baseCancel() }, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *probeExecutionBudget) Check(ctx context.Context) error {
|
||
|
|
g.mu.Lock()
|
||
|
|
defer g.mu.Unlock()
|
||
|
|
if g.stopped != nil {
|
||
|
|
return g.stopped
|
||
|
|
}
|
||
|
|
return g.GraphExecution.Check(ctx)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *probeExecutionBudget) Reserve(ctx context.Context, model, phase string) error {
|
||
|
|
g.mu.Lock()
|
||
|
|
defer g.mu.Unlock()
|
||
|
|
if g.stopped != nil {
|
||
|
|
return g.stopped
|
||
|
|
}
|
||
|
|
if err := g.GraphExecution.Check(ctx); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if g.Reservations() >= g.limit {
|
||
|
|
g.stopped = fmt.Errorf("%w: probe HTTP send cap reached", requestbudget.ErrStopped)
|
||
|
|
g.cancel(g.stopped)
|
||
|
|
return g.stopped
|
||
|
|
}
|
||
|
|
return g.GraphExecution.Reserve(ctx, model, phase)
|
||
|
|
}
|