From 0d2989d4d6af494e8ab534922d6b197dcca526af Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Thu, 20 Aug 2026 17:32:57 -0700 Subject: [PATCH] Switch to direct observation of costs on a per-operation basis --- cel/program.go | 73 +++++++---- ext/encoders_test.go | 2 +- interpreter/async.go | 6 +- interpreter/attributes.go | 10 ++ interpreter/frame.go | 16 +++ interpreter/interpretable.go | 148 +++++++++++++++++---- interpreter/interpreter_test.go | 2 +- interpreter/planner.go | 22 ++-- interpreter/runtimecost.go | 226 +++++++++++++------------------- interpreter/runtimecost_test.go | 82 +++++++++++- 10 files changed, 381 insertions(+), 206 deletions(-) diff --git a/cel/program.go b/cel/program.go index b6df752ee..740181803 100644 --- a/cel/program.go +++ b/cel/program.go @@ -341,7 +341,7 @@ func newProgram(e *Env, a *ast.AST, opts []ProgramOption) (Program, error) { observers = append(observers, interpreter.EvalStateObserver()) } if p.evalOpts&OptTrackCost == OptTrackCost { - observers = append(observers, interpreter.CostObserver(interpreter.CostTrackerFactory(trackerFactory))) + plannerOptions = append(plannerOptions, interpreter.CostObserver(interpreter.CostTrackerFactory(trackerFactory))) } // Enable exhaustive eval over a basic observer since it offers a superset of features. if p.evalOpts&OptExhaustiveEval == OptExhaustiveEval { @@ -369,19 +369,6 @@ func (p *prog) initInterpretable(a *ast.AST, plannerOptions []interpreter.Planne // Eval implements the Program interface method. func (p *prog) Eval(input any) (out ref.Val, det *EvalDetails, err error) { - // Configure error recovery for unexpected panics during evaluation. Note, the use of named - // return values makes it possible to modify the error response during the recovery - // function. - defer func() { - if r := recover(); r != nil { - switch t := r.(type) { - case interpreter.EvalCancelledError: - err = t - default: - err = fmt.Errorf("internal error: %v", r) - } - } - }() // Asynchronous calls cannot be resolved by a single-pass evaluation. Reject before doing any // work (this also covers ContextEval, which delegates here); ConcurrentEval does not call Eval. if p.hasAsync { @@ -398,19 +385,36 @@ func (p *prog) Eval(input any) (out ref.Val, det *EvalDetails, err error) { } defer frame.Close() } + // Configure error recovery and details capture for evaluation. + defer func() { + if tracker := frame.CostTracker(); tracker != nil { + if det == nil { + det = &EvalDetails{} + } + det.costTracker = tracker + } + if r := recover(); r != nil { + switch t := r.(type) { + case interpreter.EvalCancelledError: + err = t + default: + err = fmt.Errorf("internal error: %v", r) + } + } + }() + if p.observable != nil { det = &EvalDetails{} out = p.observable.ObserveExec(frame, func(observed any) { switch o := observed.(type) { case interpreter.EvalState: det.state = o - case *interpreter.CostTracker: - det.costTracker = o } }) } else { out = p.interpretable.Exec(frame) } + // The output of an internal Eval may have a value (`v`) that is a types.Err. This step // translates the CEL value to a Go error response. This interface does not quite match the // RPC signature which allows for multiple errors to be returned, but should be sufficient. @@ -511,25 +515,32 @@ func (p *prog) ConcurrentEval(ctx context.Context, input any) <-chan EvalResult go func() { defer close(resCh) + frame, err := p.newAsyncFrame(ctx, input) + if err != nil { + resCh <- EvalResult{Err: err} + return + } + defer frame.Close() + + var det *EvalDetails // Ensure concurrent eval handles panic / recovery properly defer func() { + if tracker := frame.CostTracker(); tracker != nil { + if det == nil { + det = &EvalDetails{} + } + det.costTracker = tracker + } if r := recover(); r != nil { switch t := r.(type) { case interpreter.EvalCancelledError: - resCh <- EvalResult{Err: t} + resCh <- EvalResult{EvalDetails: det, Err: t} default: - resCh <- EvalResult{Err: fmt.Errorf("internal error: %v", r)} + resCh <- EvalResult{EvalDetails: det, Err: fmt.Errorf("internal error: %v", r)} } } }() - frame, err := p.newAsyncFrame(ctx, input) - if err != nil { - resCh <- EvalResult{Err: err} - return - } - defer frame.Close() - // Completions are signaled to this channel as async calls finish. The asyncCallState // fan-in also selects on ctx.Done(), so the sender will not leak if this loop returns early. completions := make(chan int64, p.resolveCompletionBufferSize()) @@ -537,7 +548,6 @@ func (p *prog) ConcurrentEval(ctx context.Context, input any) <-chan EvalResult for { var out ref.Val - var det *EvalDetails if p.observable != nil { det = &EvalDetails{} @@ -545,13 +555,20 @@ func (p *prog) ConcurrentEval(ctx context.Context, input any) <-chan EvalResult switch o := observed.(type) { case interpreter.EvalState: det.state = o - case *interpreter.CostTracker: - det.costTracker = o } }) } else { out = p.interpretable.Exec(frame) } + // This ensures that cost tracking is present on the result passed through the channel + // in the positive outcome case, as opposed to the defer which captures these details + // in the evaluation error scenarios. + if tracker := frame.CostTracker(); tracker != nil { + if det == nil { + det = &EvalDetails{} + } + det.costTracker = tracker + } // Communicate errors quickly. if types.IsError(out) { diff --git a/ext/encoders_test.go b/ext/encoders_test.go index 1d526a15d..05da80126 100644 --- a/ext/encoders_test.go +++ b/ext/encoders_test.go @@ -249,7 +249,7 @@ func TestEncodersCosts(t *testing.T) { "x": 100, }, estimatedCost: checker.CostEstimate{Min: 2, Max: math.MaxUint64}, - actualCost: 1, + actualCost: math.MaxUint64, version: 1, }, } diff --git a/interpreter/async.go b/interpreter/async.go index a678705a7..61a43e3cf 100644 --- a/interpreter/async.go +++ b/interpreter/async.go @@ -111,7 +111,11 @@ func (fn *evalAsyncFunc) Exec(frame *ExecutionFrame) ref.Val { return unk } result := frame.ComputeResult(fn.ID(), fn.Function(), fn.OverloadID(), fn.impl, argVals) - return types.LabelErrNode(fn.id, result) + val := types.LabelErrNode(fn.id, result) + if costs := frame.CostTracker(); costs != nil { + costs.EvalVarArgs(frame, fn.id, fn, argVals, val) + } + return val } // asyncCallStateTracker manages async call states across re-evaluations of a single program. diff --git a/interpreter/attributes.go b/interpreter/attributes.go index ce344eb62..edfd6e7ca 100644 --- a/interpreter/attributes.go +++ b/interpreter/attributes.go @@ -1318,6 +1318,11 @@ func applyQualifiers(vars Activation, obj any, qualifiers []Qualifier) (any, boo if err != nil { return nil, false, err } + if frame := AsFrame(vars); frame != nil { + if costs := frame.CostTracker(); costs != nil { + costs.Qualify(qual.ID()) + } + } if !present { // We return optional none here with a presence of 'false' as the layers // above will attempt to call types.OptionalOf() on a present value if any @@ -1329,6 +1334,11 @@ func applyQualifiers(vars Activation, obj any, qualifiers []Qualifier) (any, boo if err != nil { return nil, false, err } + if frame := AsFrame(vars); frame != nil { + if costs := frame.CostTracker(); costs != nil { + costs.Qualify(qual.ID()) + } + } } obj = qualObj } diff --git a/interpreter/frame.go b/interpreter/frame.go index 2fd93052d..bc2b990f1 100644 --- a/interpreter/frame.go +++ b/interpreter/frame.go @@ -231,6 +231,22 @@ func (f *ExecutionFrame) CheckInterrupt() bool { return false } +// CostTracker returns the active CostTracker for this evaluation pass, or nil. +func (f *ExecutionFrame) CostTracker() *CostTracker { + if f.ctx == nil { + return nil + } + return f.ctx.costs +} + +// SetCostTracker sets the active CostTracker for this evaluation pass. +func (f *ExecutionFrame) SetCostTracker(tracker *CostTracker) { + if f.ctx == nil { + f.ctx = evalContextPool.Get().(*evalContext) + } + f.ctx.costs = tracker +} + // ComputeResult tracks and computes the result of the given asynchronous function. // // The first invocation for a given (node id, args) tuple registers the call state and returns an diff --git a/interpreter/interpretable.go b/interpreter/interpretable.go index d17e51d10..c13b76631 100644 --- a/interpreter/interpretable.go +++ b/interpreter/interpretable.go @@ -247,10 +247,16 @@ func (test *evalTestOnly) Exec(frame *ExecutionFrame) ref.Val { if err != nil { return types.LabelErrNode(test.id, types.WrapErr(err)) } + var res ref.Val if optVal, isOpt := val.(*types.Optional); isOpt { - return types.Bool(optVal.HasValue()) + res = types.Bool(optVal.HasValue()) + } else { + res = test.Adapter().NativeToValue(val) + } + if costs := frame.CostTracker(); costs != nil { + costs.EvalAttribute(test.id, true, res) } - return test.Adapter().NativeToValue(val) + return res } // Eval implements the Interpretable interface method. @@ -449,9 +455,16 @@ func (eq *evalEq) Exec(frame *ExecutionFrame) ref.Val { unk, _ = types.MaybeMergeUnknowns(lVal, unk) unk, _ = types.MaybeMergeUnknowns(rVal, unk) if unk != nil { + if costs := frame.CostTracker(); costs != nil { + costs.EvalBinary(frame, eq.id, eq, lVal, rVal, unk) + } return unk } - return types.Equal(lVal, rVal) + res := types.Equal(lVal, rVal) + if costs := frame.CostTracker(); costs != nil { + costs.EvalBinary(frame, eq.id, eq, lVal, rVal, res) + } + return res } // Eval implements the Interpretable interface method. @@ -499,9 +512,16 @@ func (ne *evalNe) Exec(frame *ExecutionFrame) ref.Val { unk, _ = types.MaybeMergeUnknowns(lVal, unk) unk, _ = types.MaybeMergeUnknowns(rVal, unk) if unk != nil { + if costs := frame.CostTracker(); costs != nil { + costs.EvalBinary(frame, ne.id, ne, lVal, rVal, unk) + } return unk } - return types.Bool(types.Equal(lVal, rVal) != types.True) + res := types.Bool(types.Equal(lVal, rVal) != types.True) + if costs := frame.CostTracker(); costs != nil { + costs.EvalBinary(frame, ne.id, ne, lVal, rVal, res) + } + return res } // Eval implements the Interpretable interface method. @@ -538,7 +558,11 @@ func (zero *evalZeroArity) ID() int64 { // Exec implements the InterpretableV2 interface method. func (zero *evalZeroArity) Exec(frame *ExecutionFrame) ref.Val { - return types.LabelErrNode(zero.id, zero.impl()) + res := types.LabelErrNode(zero.id, zero.impl()) + if costs := frame.CostTracker(); costs != nil { + costs.EvalZeroArity(frame, zero.id, zero, res) + } + return res } // Eval implements the Interpretable interface method. @@ -579,22 +603,33 @@ func (un *evalUnary) ID() int64 { // Exec implements the InterpretableV2 interface method. func (un *evalUnary) Exec(frame *ExecutionFrame) ref.Val { argVal := un.arg.Exec(frame) - // Early return if the argument to the function is unknown or error. + // Early return if the argument to the function is error in strict mode. strict := !un.nonStrict - if strict && types.IsUnknownOrError(argVal) { + if strict && types.IsError(argVal) { return argVal } + if strict && types.IsUnknown(argVal) { + if costs := frame.CostTracker(); costs != nil { + costs.EvalUnary(frame, un.id, un, argVal, argVal) + } + return argVal + } + var res ref.Val // If the implementation is bound and the argument value has the right traits required to // invoke it, then call the implementation. if un.impl != nil && (un.trait == 0 || (!strict && types.IsUnknownOrError(argVal)) || argVal.Type().HasTrait(un.trait)) { - return types.LabelErrNode(un.id, un.impl(argVal)) + res = types.LabelErrNode(un.id, un.impl(argVal)) + } else if argVal.Type().HasTrait(traits.ReceiverType) { + // Otherwise, if the argument is a ReceiverType attempt to invoke the receiver method on the + // operand (arg0). + res = types.LabelErrNode(un.id, argVal.(traits.Receiver).Receive(un.function, un.overload, []ref.Val{})) + } else { + res = types.NewErrWithNodeID(un.id, "no such overload: %s", un.function) } - // Otherwise, if the argument is a ReceiverType attempt to invoke the receiver method on the - // operand (arg0). - if argVal.Type().HasTrait(traits.ReceiverType) { - return types.LabelErrNode(un.id, argVal.(traits.Receiver).Receive(un.function, un.overload, []ref.Val{})) + if costs := frame.CostTracker(); costs != nil { + costs.EvalUnary(frame, un.id, un, argVal, res) } - return types.NewErrWithNodeID(un.id, "no such overload: %s", un.function) + return res } // Eval implements the Interpretable interface method. @@ -649,20 +684,28 @@ func (bin *evalBinary) Exec(frame *ExecutionFrame) ref.Val { unk, _ = types.MaybeMergeUnknowns(lVal, unk) unk, _ = types.MaybeMergeUnknowns(rVal, unk) if unk != nil { + if costs := frame.CostTracker(); costs != nil { + costs.EvalBinary(frame, bin.id, bin, lVal, rVal, unk) + } return unk } } + var res ref.Val // If the implementation is bound and the argument value has the right traits required to // invoke it, then call the implementation. if bin.impl != nil && (bin.trait == 0 || (!strict && types.IsUnknownOrError(lVal)) || lVal.Type().HasTrait(bin.trait)) { - return types.LabelErrNode(bin.id, bin.impl(lVal, rVal)) + res = types.LabelErrNode(bin.id, bin.impl(lVal, rVal)) + } else if lVal.Type().HasTrait(traits.ReceiverType) { + // Otherwise, if the argument is a ReceiverType attempt to invoke the receiver method on the + // operand (arg0). + res = types.LabelErrNode(bin.id, lVal.(traits.Receiver).Receive(bin.function, bin.overload, []ref.Val{rVal})) + } else { + res = types.NewErrWithNodeID(bin.id, "no such overload: %s", bin.function) } - // Otherwise, if the argument is a ReceiverType attempt to invoke the receiver method on the - // operand (arg0). - if lVal.Type().HasTrait(traits.ReceiverType) { - return types.LabelErrNode(bin.id, lVal.(traits.Receiver).Receive(bin.function, bin.overload, []ref.Val{rVal})) + if costs := frame.CostTracker(); costs != nil { + costs.EvalBinary(frame, bin.id, bin, lVal, rVal, res) } - return types.NewErrWithNodeID(bin.id, "no such overload: %s", bin.function) + return res } // Eval implements the Interpretable interface method. @@ -726,20 +769,40 @@ func (fn *evalVarArgs) Exec(frame *ExecutionFrame) ref.Val { } } if strict && unk != nil { + if costs := frame.CostTracker(); costs != nil { + costs.EvalVarArgs(frame, fn.id, fn, argVals, unk) + } return unk } + if len(argVals) == 0 { + var res ref.Val + if fn.impl != nil { + res = types.LabelErrNode(fn.id, fn.impl()) + } else { + res = types.NewErrWithNodeID(fn.id, "no such overload: %s %d", fn.function, fn.id) + } + if costs := frame.CostTracker(); costs != nil { + costs.EvalZeroArity(frame, fn.id, fn, res) + } + return res + } + var res ref.Val // If the implementation is bound and the argument value has the right traits required to // invoke it, then call the implementation. arg0 := argVals[0] if fn.impl != nil && (fn.trait == 0 || (!strict && types.IsUnknownOrError(arg0)) || arg0.Type().HasTrait(fn.trait)) { - return types.LabelErrNode(fn.id, fn.impl(argVals...)) + res = types.LabelErrNode(fn.id, fn.impl(argVals...)) + } else if arg0.Type().HasTrait(traits.ReceiverType) { + // Otherwise, if the argument is a ReceiverType attempt to invoke the receiver method on the + // operand (arg0). + res = types.LabelErrNode(fn.id, arg0.(traits.Receiver).Receive(fn.function, fn.overload, argVals[1:])) + } else { + res = types.NewErrWithNodeID(fn.id, "no such overload: %s %d", fn.function, fn.id) } - // Otherwise, if the argument is a ReceiverType attempt to invoke the receiver method on the - // operand (arg0). - if arg0.Type().HasTrait(traits.ReceiverType) { - return types.LabelErrNode(fn.id, arg0.(traits.Receiver).Receive(fn.function, fn.overload, argVals[1:])) + if costs := frame.CostTracker(); costs != nil { + costs.EvalVarArgs(frame, fn.id, fn, argVals, res) } - return types.NewErrWithNodeID(fn.id, "no such overload: %s %d", fn.function, fn.id) + return res } // Eval implements the Interpretable interface method. @@ -802,9 +865,16 @@ func (l *evalList) Exec(frame *ExecutionFrame) ref.Val { elemVals = append(elemVals, elemVal) } if unk != nil { + if costs := frame.CostTracker(); costs != nil { + costs.CreateList(l.id, unk) + } return unk } - return types.NewRefValList(l.adapter, elemVals) + res := types.NewRefValList(l.adapter, elemVals) + if costs := frame.CostTracker(); costs != nil { + costs.CreateList(l.id, res) + } + return res } // Eval implements the Interpretable interface method. @@ -865,9 +935,16 @@ func (m *evalMap) Exec(frame *ExecutionFrame) ref.Val { entries[keyVal] = valVal } if unk != nil { + if costs := frame.CostTracker(); costs != nil { + costs.CreateMap(m.id, unk) + } return unk } - return types.NewRefValMap(m.adapter, entries) + res := types.NewRefValMap(m.adapter, entries) + if costs := frame.CostTracker(); costs != nil { + costs.CreateMap(m.id, res) + } + return res } // Eval implements the Interpretable interface method. @@ -934,9 +1011,16 @@ func (o *evalObj) Exec(frame *ExecutionFrame) ref.Val { fieldVals[field] = val } if unk != nil { + if costs := frame.CostTracker(); costs != nil { + costs.CreateStruct(o.id, unk) + } return unk } - return types.LabelErrNode(o.id, o.provider.NewValue(o.typeName, fieldVals)) + res := types.LabelErrNode(o.id, o.provider.NewValue(o.typeName, fieldVals)) + if costs := frame.CostTracker(); costs != nil { + costs.CreateStruct(o.id, res) + } + return res } // Eval implements the Interpretable interface method. @@ -1446,7 +1530,13 @@ func (a *evalAttr) Exec(frame *ExecutionFrame) ref.Val { if err != nil { return types.LabelErrNode(a.ID(), types.WrapErr(err)) } - return a.adapter.NativeToValue(v) + res := a.adapter.NativeToValue(v) + if costs := frame.CostTracker(); costs != nil { + if _, isCond := a.attr.(*conditionalAttribute); !isCond { + costs.EvalAttribute(a.ID(), false, res) + } + } + return res } // Eval implements the Interpretable interface method. diff --git a/interpreter/interpreter_test.go b/interpreter/interpreter_test.go index 76e9e1b41..bc8df4f50 100644 --- a/interpreter/interpreter_test.go +++ b/interpreter/interpreter_test.go @@ -2661,7 +2661,7 @@ func newTestPartialActivation(t testing.TB, in any, unknowns ...*AttributePatter // newStandardInterpreter builds a Dispatcher and TypeProvider with support for all of the CEL // builtins defined in the language definition. -func newStandardInterpreter(t *testing.T, +func newStandardInterpreter(t testing.TB, container *containers.Container, provider types.Provider, adapter types.Adapter, diff --git a/interpreter/planner.go b/interpreter/planner.go index bdf183be7..fdd57fbb7 100644 --- a/interpreter/planner.go +++ b/interpreter/planner.go @@ -51,15 +51,16 @@ func newPlanner(disp Dispatcher, // planner is an implementation of the interpretablePlanner interface. type planner struct { - disp Dispatcher - provider types.Provider - adapter types.Adapter - attrFactory AttributeFactory - container *containers.Container - refMap map[int64]*ast.ReferenceInfo - typeMap map[int64]*types.Type - decorators []InterpretableDecoratorV2 - observers []StatefulObserver + disp Dispatcher + provider types.Provider + adapter types.Adapter + attrFactory AttributeFactory + container *containers.Container + refMap map[int64]*ast.ReferenceInfo + typeMap map[int64]*types.Type + decorators []InterpretableDecoratorV2 + observers []StatefulObserver + costTrackerFactory func() (*CostTracker, error) } type planBuilder struct { @@ -85,6 +86,9 @@ func (p *planner) Plan(expr ast.Expr) (InterpretableV2, error) { if err != nil { return nil, err } + if p.costTrackerFactory != nil { + i = &costTrackingInterpretable{InterpretableV2: i, factory: p.costTrackerFactory} + } if len(p.observers) == 0 { return i, nil } diff --git a/interpreter/runtimecost.go b/interpreter/runtimecost.go index 5558cf6fd..e92acf3ab 100644 --- a/interpreter/runtimecost.go +++ b/interpreter/runtimecost.go @@ -56,7 +56,7 @@ func CostObserver(opts ...costTrackPlanOption) PlannerOption { if ct.factory == nil { return nil, errors.New("cost tracker factory not configured") } - p.observers = append(p.observers, ct) + p.costTrackerFactory = ct.factory return p, nil } } @@ -91,86 +91,8 @@ func (ct *costTrackerFactory) GetState(frame *ExecutionFrame) any { return frame.ctx.costs } -// Observe computes the incremental cost of each step and records it into the CostTracker associated -// with the evaluation. +// Observe implements the StatefulObserver interface. func (ct *costTrackerFactory) Observe(vars Activation, id int64, programStep any, val ref.Val) { - frame := AsFrame(vars) - state := ct.GetState(frame) - if state == nil { - return - } - tracker, ok := state.(*CostTracker) - if !ok { - // The state is configured with CostTrackFactory so this shouldn't happen. - return - } - switch t := programStep.(type) { - case ConstantQualifier: - // TODO: Push identifiers on to the stack before observing constant qualifiers that apply to them - // and enable the below pop. Once enabled this can case can be collapsed into the Qualifier case. - tracker.cost++ - case InterpretableConst: - // zero cost - case InterpretableAttribute: - switch a := t.Attr().(type) { - case *conditionalAttribute: - // Ternary has no direct cost. All cost is from the conditional and the true/false branch expressions. - tracker.stack.drop(a.falsy.ID(), a.truthy.ID(), a.expr.ID()) - default: - tracker.stack.drop(t.Attr().ID()) - tracker.cost += common.SelectAndIdentCost - } - if !tracker.presenceTestHasCost { - if _, isTestOnly := programStep.(*evalTestOnly); isTestOnly { - tracker.cost -= common.SelectAndIdentCost - } - } - case *evalExhaustiveConditional: - // Ternary has no direct cost. All cost is from the conditional and the true/false branch expressions. - tracker.stack.drop(t.attr.falsy.ID(), t.attr.truthy.ID(), t.attr.expr.ID()) - - // While the field names are identical, the boolean operation eval structs do not share an interface and so - // must be handled individually. - case *evalOr: - for _, term := range t.terms { - tracker.stack.drop(term.ID()) - } - case *evalAnd: - for _, term := range t.terms { - tracker.stack.drop(term.ID()) - } - case *evalExhaustiveOr: - for _, term := range t.terms { - tracker.stack.drop(term.ID()) - } - case *evalExhaustiveAnd: - for _, term := range t.terms { - tracker.stack.drop(term.ID()) - } - case *evalFold: - tracker.stack.drop(t.iterRange.ID()) - case Qualifier: - tracker.cost++ - case InterpretableCall: - if argVals, ok := tracker.stack.dropArgs(t.Args()); ok { - tracker.cost += tracker.costCall(t, argVals, val) - } - case InterpretableConstructor: - tracker.stack.dropArgs(t.InitVals()) - switch t.Type() { - case types.ListType: - tracker.cost += common.ListCreateBaseCost - case types.MapType: - tracker.cost += common.MapCreateBaseCost - default: - tracker.cost += common.StructCreateBaseCost - } - } - tracker.stack.push(val, id) - - if tracker.Limit != nil && tracker.cost > *tracker.Limit { - panic(EvalCancelledError{Cause: CostLimitExceeded, Message: "operation cancelled: actual cost limit exceeded"}) - } } // CostTrackerOption configures the behavior of CostTracker objects. @@ -231,8 +153,7 @@ type CostTracker struct { Limit *uint64 presenceTestHasCost bool - cost uint64 - stack refValStack + cost uint64 } // Clone makes a shallow copy of the tracker. @@ -253,6 +174,93 @@ func (c *CostTracker) ActualCost() uint64 { return c.cost } +// CreateList records list literal construction cost. +func (c *CostTracker) CreateList(id int64, res ref.Val) { + c.cost = cost.SafeAdd(c.cost, common.ListCreateBaseCost) + c.checkLimit() +} + +// CreateMap records map literal construction cost. +func (c *CostTracker) CreateMap(id int64, res ref.Val) { + c.cost = cost.SafeAdd(c.cost, common.MapCreateBaseCost) + c.checkLimit() +} + +// CreateStruct records struct/object construction cost. +func (c *CostTracker) CreateStruct(id int64, res ref.Val) { + c.cost = cost.SafeAdd(c.cost, common.StructCreateBaseCost) + c.checkLimit() +} + +// EvalAttribute records attribute resolution cost (ident / select). +func (c *CostTracker) EvalAttribute(id int64, isTestOnly bool, res ref.Val) { + if !isTestOnly || c.presenceTestHasCost { + c.cost = cost.SafeAdd(c.cost, common.SelectAndIdentCost) + c.checkLimit() + } +} + +// Qualify records qualifier cost. +func (c *CostTracker) Qualify(id int64) { + c.cost = cost.SafeAdd(c.cost, 1) + c.checkLimit() +} + +type costTrackingInterpretable struct { + InterpretableV2 + factory func() (*CostTracker, error) +} + +func (c *costTrackingInterpretable) Exec(frame *ExecutionFrame) ref.Val { + if frame.CostTracker() == nil { + tracker, err := c.factory() + if err != nil { + return types.NewErr("cost tracker factory: %v", err) + } + frame.SetCostTracker(tracker) + } + return c.InterpretableV2.Exec(frame) +} + +func (c *costTrackingInterpretable) Eval(ctx Activation) ref.Val { + return c.Exec(AsFrame(ctx)) +} + +// EvalZeroArity records the cost for a 0-arity call expression. +func (c *CostTracker) EvalZeroArity(vars Activation, id int64, call InterpretableCall, result ref.Val) { + c.cost = cost.SafeAdd(c.cost, c.costCall(call, nil, result)) + c.checkLimit() +} + +// EvalUnary records the cost for a unary call expression. +func (c *CostTracker) EvalUnary(vars Activation, id int64, call InterpretableCall, arg ref.Val, result ref.Val) { + var buf [1]ref.Val + buf[0] = arg + c.cost = cost.SafeAdd(c.cost, c.costCall(call, buf[:], result)) + c.checkLimit() +} + +// EvalBinary records the cost for a binary call expression. +func (c *CostTracker) EvalBinary(vars Activation, id int64, call InterpretableCall, lhs, rhs ref.Val, result ref.Val) { + var buf [2]ref.Val + buf[0] = lhs + buf[1] = rhs + c.cost = cost.SafeAdd(c.cost, c.costCall(call, buf[:], result)) + c.checkLimit() +} + +// EvalVarArgs records the cost for a variadic call expression. +func (c *CostTracker) EvalVarArgs(vars Activation, id int64, call InterpretableCall, args []ref.Val, result ref.Val) { + c.cost = cost.SafeAdd(c.cost, c.costCall(call, args, result)) + c.checkLimit() +} + +func (c *CostTracker) checkLimit() { + if c.Limit != nil && c.cost > *c.Limit { + panic(EvalCancelledError{Cause: CostLimitExceeded, Message: "operation cancelled: actual cost limit exceeded"}) + } +} + func (c *CostTracker) costCall(call InterpretableCall, args []ref.Val, result ref.Val) uint64 { var total uint64 if len(c.overloadTrackers) != 0 { @@ -342,57 +350,3 @@ func actualSize(value ref.Val) uint64 { } return 1 } - -type stackVal struct { - Val ref.Val - ID int64 -} - -// refValStack keeps track of values of the stack for cost calculation purposes -type refValStack []stackVal - -func (s *refValStack) push(val ref.Val, id int64) { - value := stackVal{Val: val, ID: id} - *s = append(*s, value) -} - -// TODO: Allowing drop and dropArgs to remove stack items above the IDs they are provided is a workaround. drop and dropArgs -// should find and remove only the stack items matching the provided IDs once all attributes are properly pushed and popped from stack. - -// drop searches the stack for each ID and removes the ID and all stack items above it. -// If none of the IDs are found, the stack is not modified. -// WARNING: It is possible for multiple expressions with the same ID to exist (due to how macros are implemented) so it's -// possible that a dropped ID will remain on the stack. They should be removed when IDs on the stack are popped. -func (s *refValStack) drop(ids ...int64) { - for _, id := range ids { - for idx := len(*s) - 1; idx >= 0; idx-- { - if (*s)[idx].ID == id { - *s = (*s)[:idx] - break - } - } - } -} - -// dropArgs searches the stack for all the args by their IDs, accumulates their associated ref.Vals and drops any -// stack items above any of the arg IDs. If any of the IDs are not found the stack, false is returned. -// Args are assumed to be found in the stack in reverse order, i.e. the last arg is expected to be found highest in -// the stack. -// WARNING: It is possible for multiple expressions with the same ID to exist (due to how macros are implemented) so it's -// possible that a dropped ID will remain on the stack. They should be removed when IDs on the stack are popped. -func (s *refValStack) dropArgs(args []InterpretableV2) ([]ref.Val, bool) { - result := make([]ref.Val, len(args)) -argloop: - for nIdx := len(args) - 1; nIdx >= 0; nIdx-- { - for idx := len(*s) - 1; idx >= 0; idx-- { - if (*s)[idx].ID == args[nIdx].ID() { - el := (*s)[idx] - *s = (*s)[:idx] - result[nIdx] = el.Val - continue argloop - } - } - return nil, false - } - return result, true -} diff --git a/interpreter/runtimecost_test.go b/interpreter/runtimecost_test.go index b160318b7..347d3ba6b 100644 --- a/interpreter/runtimecost_test.go +++ b/interpreter/runtimecost_test.go @@ -174,7 +174,7 @@ func computeCost(t *testing.T, expr string, vars []*decls.VariableDecl, ctx Acti return costTracker.cost, est, err } -func constructActivation(t *testing.T, in any) Activation { +func constructActivation(t testing.TB, in any) Activation { t.Helper() if in == nil { return EmptyActivation() @@ -904,3 +904,83 @@ func TestRuntimeCost(t *testing.T) { }) } } + +func BenchmarkCostTracking(b *testing.B) { + benchmarks := []struct { + name string + expr string + vars []*decls.VariableDecl + in map[string]any + }{ + { + name: "simple_comparison", + expr: "x > 10", + vars: []*decls.VariableDecl{decls.NewVariable("x", types.IntType)}, + in: map[string]any{"x": 15}, + }, + { + name: "function_calls", + expr: "str.startsWith('hello') && str.endsWith('world')", + vars: []*decls.VariableDecl{decls.NewVariable("str", types.StringType)}, + in: map[string]any{"str": "hello beautiful world"}, + }, + { + name: "comprehension", + expr: "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10].map(x, x * 2).filter(x, x > 10)", + }, + { + name: "nested_comprehensions", + expr: "[1, 2, 3, 4, 5].all(i, [1, 2, 3, 4, 5].exists(j, i + j == 6))", + }, + } + + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + s := common.NewTextSource(bm.expr) + p, err := parser.NewParser(parser.Macros(parser.AllMacros...)) + if err != nil { + b.Fatalf("Failed to initialize parser: %v", err) + } + parsed, errs := p.Parse(s) + if len(errs.GetErrors()) != 0 { + b.Fatalf("Parse(%s) failed: %v", bm.expr, errs.GetErrors()) + } + + cont := containers.DefaultContainer + reg := newTestRegistry(b, types.ProtoTypeDefs(&proto3pb.TestAllTypes{})) + attrs := NewAttributeFactory(cont, reg, reg) + env := newTestEnv(b, cont, reg) + if len(bm.vars) > 0 { + err = env.AddIdents(bm.vars...) + if err != nil { + b.Fatalf("Failed to add idents: %v", err) + } + } + checked, errs := checker.Check(parsed, s, env) + if len(errs.GetErrors()) != 0 { + b.Fatalf("Check(%s) failed: %v", bm.expr, errs.GetErrors()) + } + + evalCostTracker, err := NewCostTracker(nil) + if err != nil { + b.Fatalf("NewCostTracker() failed: %v", err) + } + trackerFactory := func() (*CostTracker, error) { + return evalCostTracker.Clone() + } + interp := newStandardInterpreter(b, cont, reg, reg, attrs) + prg, err := interp.NewInterpretable(checked, CostObserver(CostTrackerFactory(trackerFactory))) + if err != nil { + b.Fatalf("NewInterpretable(%s) failed: %v", bm.expr, err) + } + + ctx := constructActivation(b, bm.in) + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + prg.Eval(ctx) + } + }) + } +} +