xinmi/scripts/graph_extraction_probe/budget.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)
}