diff --git a/cel/memory_test.go b/cel/memory_test.go new file mode 100644 index 000000000..03514af0c --- /dev/null +++ b/cel/memory_test.go @@ -0,0 +1,223 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cel + +import ( + "context" + "math" + "strings" + "testing" + + "cel.dev/cel-go/common/types" +) + +func TestMemoryTracking(t *testing.T) { + tests := []struct { + name string + expr string + decls []EnvOption + memOpts []types.MemoryTrackerOption + in any + wantPeak uint32 + }{ + { + name: "attribute_resolution", + expr: `a`, + decls: []EnvOption{Variable("a", ListType(IntType))}, + in: map[string]any{"a": []int64{1, 2, 3}}, + // 1 (list container) + 3 elements, then the call-free program peaks at the attribute. + wantPeak: 4, + }, + { + name: "call_output", + expr: `a + a`, + decls: []EnvOption{Variable("a", ListType(IntType))}, + in: map[string]any{"a": []int64{1, 2}}, + // The peak is the call output: a lazy concat list backed by both inputs, sizing + // as their sum (3 + 3 = 6), which exceeds either input observed on its own. + wantPeak: 6, + }, + { + name: "attribute_field_selection", + expr: `m.vals`, + decls: []EnvOption{Variable("m", MapType(StringType, ListType(IntType)))}, + in: map[string]any{"m": map[string][]int64{"vals": {1, 2, 3}}}, + // The resolved attribute value is the inner list: 1 (container) + 3 elements. + wantPeak: 4, + }, + } + + for _, tst := range tests { + tc := tst + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + env := testEnv(t, tc.decls...) + ast, iss := env.Compile(tc.expr) + if iss.Err() != nil { + t.Fatalf("env.Compile(%v) failed: %v", tc.expr, iss.Err()) + } + program, err := env.Program(ast, MemoryTracking(tc.memOpts...)) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + _, details, err := program.Eval(tc.in) + if err != nil { + t.Fatalf("program.Eval() failed: %v", err) + } + peak := details.PeakMemory() + if peak == nil { + t.Fatalf("EvalDetails.PeakMemory() got nil, wanted %d", tc.wantPeak) + } + if *peak != tc.wantPeak { + t.Errorf("EvalDetails.PeakMemory() got %d, wanted %d", *peak, tc.wantPeak) + } + }) + } +} + +func TestMemoryTrackingComprehension(t *testing.T) { + env := testEnv(t, Variable("a", ListType(IntType))) + ast, iss := env.Compile(`a.map(x, x * 2)`) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + program, err := env.Program(ast, MemoryTracking()) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + in := map[string]any{"a": []int64{1, 2, 3, 4, 5}} + _, details, err := program.Eval(in) + if err != nil { + t.Fatalf("program.Eval() failed: %v", err) + } + peak := details.PeakMemory() + if peak == nil { + t.Fatal("EvalDetails.PeakMemory() got nil, wanted non-nil peak") + } + // The comprehension result is 1 (container) + 5 elements; the peak must be at least as + // large since the final accumulation observes the input alongside the built list. + if *peak < 6 { + t.Errorf("EvalDetails.PeakMemory() got %d, wanted at least 6", *peak) + } +} + +func TestMemoryTrackingConcurrentEval(t *testing.T) { + env := testEnv(t, Variable("a", ListType(IntType))) + ast, iss := env.Compile(`a + a`) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + program, err := env.Program(ast, MemoryTracking()) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + res := <-program.ConcurrentEval(ctx, map[string]any{"a": []int64{1, 2}}) + if res.Err != nil { + t.Fatalf("program.ConcurrentEval() failed: %v", res.Err) + } + peak := res.EvalDetails.PeakMemory() + if peak == nil { + t.Fatal("EvalDetails.PeakMemory() got nil, wanted non-nil peak") + } + if *peak != 6 { + t.Errorf("EvalDetails.PeakMemory() got %d, wanted 6", *peak) + } +} + +func TestMemoryTrackingDisabled(t *testing.T) { + env := testEnv(t, Variable("a", ListType(IntType))) + ast, iss := env.Compile(`a + a`) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + program, err := env.Program(ast, EvalOptions(OptTrackState)) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + _, details, err := program.Eval(map[string]any{"a": []int64{1, 2}}) + if err != nil { + t.Fatalf("program.Eval() failed: %v", err) + } + if peak := details.PeakMemory(); peak != nil { + t.Errorf("EvalDetails.PeakMemory() got %d, wanted nil when tracking disabled", *peak) + } +} + +func TestMemoryLimit(t *testing.T) { + tests := []struct { + name string + memLimit uint32 + wantErr string + }{ + { + name: "under_limit", + memLimit: 1000, + }, + { + name: "over_limit", + memLimit: 5, + wantErr: "memory limit exceeded", + }, + } + for _, tst := range tests { + tc := tst + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + env := testEnv(t, Variable("a", ListType(IntType))) + ast, iss := env.Compile(`a + a`) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + program, err := env.Program(ast, MemoryLimit(tc.memLimit)) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + _, _, err = program.Eval(map[string]any{"a": []int64{1, 2, 3}}) + if tc.wantErr == "" { + if err != nil { + t.Fatalf("program.Eval() failed: %v", err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("program.Eval() got error %v, wanted error containing %q", err, tc.wantErr) + } + }) + } +} + +func TestMemoryTrackingCalculationLimitExceeded(t *testing.T) { + env := testEnv(t, Variable("a", ListType(StringType))) + ast, iss := env.Compile(`a`) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + program, err := env.Program(ast, + MemoryTracking( + types.MemoryTrackerSizeCalculator( + types.NewSizeCalculator(types.SizeCalculatorMaxTraversal(2))))) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + _, details, err := program.Eval(map[string]any{"a": []string{"a", "b", "c", "d", "e"}}) + if err != nil { + t.Fatalf("program.Eval() failed: %v", err) + } + if peak := details.PeakMemory(); peak == nil || *peak != math.MaxUint32 { + t.Errorf("EvalDetails.PeakMemory() got %v, wanted MaxUint32", peak) + } +} diff --git a/cel/options.go b/cel/options.go index 21cdb3615..ea3c3f949 100644 --- a/cel/options.go +++ b/cel/options.go @@ -729,6 +729,10 @@ const ( // // Deprecated: use ext.StringsValidateFormatCalls() as this option is now a no-op. OptCheckStringFormat EvalOption = 1 << iota + + // OptTrackMemory enables the runtime peak memory tracking and returns the peak watermark within + // evalDetails via func PeakMemory() + OptTrackMemory EvalOption = 1 << iota ) // EvalOptions sets one or more evaluation options which may affect the evaluation or Result. @@ -818,6 +822,36 @@ func CostTracking(costEstimator interpreter.ActualCostEstimator) ProgramOption { } } +// MemoryTracking enables peak memory tracking during evaluation with an optional set of +// types.MemoryTrackerOption values to configure the tracker's size calculator, sample +// interval, and limit behaviors. +// +// Peak memory is measured in aggregate element counts as computed by a types.SizeCalculator +// and is observed at the points where values materialize during evaluation: resolved +// attributes, call results, constructed aggregates, and comprehension results. The peak +// watermark is available via the EvalDetails.PeakMemory() method. +func MemoryTracking(memOpts ...types.MemoryTrackerOption) ProgramOption { + return func(p *prog) (*prog, error) { + p.memoryOptions = append(p.memoryOptions, memOpts...) + p.evalOpts |= OptTrackMemory + return p, nil + } +} + +// MemoryLimit enables memory tracking and configures program evaluation to exit early with a +// "memory limit exceeded" error if the peak tracked memory exceeds the limit. +// +// The MemoryLimit is a metric that corresponds to the aggregate element counts of the values +// observed during evaluation. It is indicative of memory usage, not CPU usage; see CostLimit +// for bounding compute. +func MemoryLimit(memLimit uint32) ProgramOption { + return func(p *prog) (*prog, error) { + p.memoryLimit = &memLimit + p.evalOpts |= OptTrackMemory + return p, nil + } +} + // CostLimit enables cost tracking and sets configures program evaluation to exit early with a // "runtime cost limit exceeded" error if the runtime cost exceeds the costLimit. // The CostLimit is a metric that corresponds to the number and estimated expense of operations diff --git a/cel/program.go b/cel/program.go index 740181803..798f7eef5 100644 --- a/cel/program.go +++ b/cel/program.go @@ -141,6 +141,7 @@ type AttributePatternType = interpreter.AttributePattern type EvalDetails struct { state interpreter.EvalState costTracker *interpreter.CostTracker + memTracker *types.MemoryTracker } // State of the evaluation, non-nil if the OptTrackState or OptExhaustiveEval is specified @@ -162,6 +163,16 @@ func (ed *EvalDetails) ActualCost() *uint64 { return &cost } +// PeakMemory returns the peak memory watermark observed through the course of execution when +// `MemoryTracking` is enabled. Otherwise, returns nil if memory tracking was not enabled. +func (ed *EvalDetails) PeakMemory() *uint32 { + if ed == nil || ed.memTracker == nil { + return nil + } + peak := ed.memTracker.Peak() + return &peak +} + // EvalResult encapsulates the response from a ConcurrentEval call. type EvalResult struct { Val ref.Val @@ -189,6 +200,8 @@ type prog struct { callCostEstimator interpreter.ActualCostEstimator costOptions []interpreter.CostTrackerOption costLimit *uint64 + memoryOptions []types.MemoryTrackerOption + memoryLimit *uint32 // hasAsync indicates the planned expression contains an asynchronous function call, which can // only be resolved by ConcurrentEval. @@ -313,36 +326,52 @@ func newProgram(e *Env, a *ast.AST, opts []ProgramOption) (Program, error) { plannerOptions = append(plannerOptions, interpreter.RegexProgramSizeLimit(limit)) } - // Enable exhaustive eval, state tracking and cost tracking last since they require a factory. - if p.evalOpts&(OptExhaustiveEval|OptTrackState|OptTrackCost) != 0 { - costOptCount := len(p.costOptions) - if p.costLimit != nil { - costOptCount++ - } - costOpts := make([]interpreter.CostTrackerOption, 0, costOptCount) - costOpts = append(costOpts, p.costOptions...) - if p.costLimit != nil { - costOpts = append(costOpts, interpreter.CostTrackerLimit(*p.costLimit)) - } - // Creating a new cost tracker for each evaluation causes significant work that - // needs to be repeated for each evaluation even though the cost tracker is - // mostly read-only once constructed. Therefore it gets constructed - // once now and later a cheap clone is used for each evaluation. - tracker, err := interpreter.NewCostTracker(p.callCostEstimator, costOpts...) - if err != nil { - return nil, fmt.Errorf("construct cost tracker: %w", err) - } - trackerFactory := func() (*interpreter.CostTracker, error) { - return tracker.Clone() - } + // Enable exhaustive eval, state tracking, cost tracking, and memory tracking last since they + // require a factory. + if p.evalOpts&(OptExhaustiveEval|OptTrackState|OptTrackCost|OptTrackMemory) != 0 { var observers []interpreter.PlannerOption if p.evalOpts&(OptExhaustiveEval|OptTrackState) != 0 { // EvalStateObserver is required for OptExhaustiveEval. observers = append(observers, interpreter.EvalStateObserver()) } if p.evalOpts&OptTrackCost == OptTrackCost { + costOptCount := len(p.costOptions) + if p.costLimit != nil { + costOptCount++ + } + costOpts := make([]interpreter.CostTrackerOption, 0, costOptCount) + costOpts = append(costOpts, p.costOptions...) + if p.costLimit != nil { + costOpts = append(costOpts, interpreter.CostTrackerLimit(*p.costLimit)) + } + // Creating a new cost tracker for each evaluation causes significant work that + // needs to be repeated for each evaluation even though the cost tracker is + // mostly read-only once constructed. Therefore it gets constructed + // once now and later a cheap clone is used for each evaluation. + tracker, err := interpreter.NewCostTracker(p.callCostEstimator, costOpts...) + if err != nil { + return nil, fmt.Errorf("construct cost tracker: %w", err) + } + trackerFactory := func() (*interpreter.CostTracker, error) { + return tracker.Clone() + } plannerOptions = append(plannerOptions, interpreter.CostObserver(interpreter.CostTrackerFactory(trackerFactory))) } + if p.evalOpts&OptTrackMemory == OptTrackMemory { + memOptCount := len(p.memoryOptions) + if p.memoryLimit != nil { + memOptCount++ + } + memOpts := make([]types.MemoryTrackerOption, 0, memOptCount) + memOpts = append(memOpts, p.memoryOptions...) + if p.memoryLimit != nil { + memOpts = append(memOpts, types.MemoryTrackerLimit(*p.memoryLimit)) + } + memTrackerFactory := func() (*types.MemoryTracker, error) { + return types.NewMemoryTracker(memOpts...), nil + } + observers = append(observers, interpreter.MemoryObserver(interpreter.MemoryTrackerFactory(memTrackerFactory))) + } // Enable exhaustive eval over a basic observer since it offers a superset of features. if p.evalOpts&OptExhaustiveEval == OptExhaustiveEval { plannerOptions = append(plannerOptions, @@ -409,6 +438,8 @@ func (p *prog) Eval(input any) (out ref.Val, det *EvalDetails, err error) { switch o := observed.(type) { case interpreter.EvalState: det.state = o + case *types.MemoryTracker: + det.memTracker = o } }) } else { @@ -555,6 +586,8 @@ func (p *prog) ConcurrentEval(ctx context.Context, input any) <-chan EvalResult switch o := observed.(type) { case interpreter.EvalState: det.state = o + case *types.MemoryTracker: + det.memTracker = o } }) } else { diff --git a/ext/bindings_test.go b/ext/bindings_test.go index ced229716..981b786fa 100644 --- a/ext/bindings_test.go +++ b/ext/bindings_test.go @@ -16,9 +16,11 @@ package ext import ( "fmt" + "runtime" "strings" "sync" "testing" + "time" "cel.dev/cel-go/cel" "cel.dev/cel-go/checker" @@ -189,6 +191,289 @@ func TestBindings(t *testing.T) { } } +// nestedBindExpr produces an exponentially amplified value through nested cel.bind calls: +// +// cel.bind(a0, [0,1,2,3,4,5,6,7,8,9], +// cel.bind(a1, [a0,a0,a0,a0,a0,a0,a0,a0,a0,a0], +// ... cel.bind(aN, [aN-1 x10], aN == aN))) +// +// Each level is physically a 10-element list of references to the level below, so the +// logical element count grows by 10x per level while the physical allocation stays small. +func nestedBindExpr(levels int) string { + expr := fmt.Sprintf("a%d == a%d", levels, levels) + for i := levels; i >= 1; i-- { + refs := strings.TrimSuffix(strings.Repeat(fmt.Sprintf("a%d,", i-1), 10), ",") + expr = fmt.Sprintf("cel.bind(a%d, [%s], %s)", i, refs, expr) + } + return fmt.Sprintf("cel.bind(a0, [0,1,2,3,4,5,6,7,8,9], %s)", expr) +} + +func TestBindingsMemoryPeakAmplification(t *testing.T) { + // The tracked peak reflects the full logical element count of the largest bound value: + // size(a0) = 11, and size(aN) = 1 + 10*size(aN-1). Each bound level is a list of + // references to a single shared instance whose aggregate size is memoized on first + // computation, so sizing each level costs ~11 traversals while remaining exact — the + // calculator's traversal and depth budgets never bind on the shared structure. + tests := []struct { + levels int + wantPeak uint32 + }{ + {levels: 2, wantPeak: 1111}, + {levels: 4, wantPeak: 111111}, + } + env, err := cel.NewEnv(Bindings()) + if err != nil { + t.Fatalf("cel.NewEnv(Bindings()) failed: %v", err) + } + for _, tc := range tests { + t.Run(fmt.Sprintf("levels_%d", tc.levels), func(t *testing.T) { + expr := nestedBindExpr(tc.levels) + ast, iss := env.Compile(expr) + if iss.Err() != nil { + t.Fatalf("env.Compile(%v) failed: %v", expr, iss.Err()) + } + prg, err := env.Program(ast, cel.MemoryTracking()) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + out, details, err := prg.Eval(cel.NoVars()) + if err != nil { + t.Fatalf("prg.Eval() failed: %v", err) + } + if out != types.True { + t.Errorf("prg.Eval() got %v, wanted true", out) + } + peak := details.PeakMemory() + if peak == nil { + t.Fatal("EvalDetails.PeakMemory() got nil, wanted non-nil peak") + } + if *peak != tc.wantPeak { + t.Errorf("EvalDetails.PeakMemory() got %d, wanted %d", *peak, tc.wantPeak) + } + }) + } +} + +func TestBindingsMemoryLimitAmplification(t *testing.T) { + // With eight amplification levels the logical size is ~10^9 elements while the physical + // structure is ~90 small lists sharing backing storage. The memory limit must trip while + // the bound values are being constructed, long before the a8 == a8 comparison would + // attempt to traverse the logical structure. + // + // Because each level's aggregate size is memoized on the shared list instances, sizing + // stays exact (size(aN) = 1 + 10*size(aN-1)) at ~11 traversals per level, and the + // calculator's traversal and depth budgets never bind under either configuration below. + // The limit trips at the a5 binding, whose exact logical size of 1,111,111 elements + // exceeds the 1M limit. Since sizing walks shared references without materializing the + // logical value, the Go-native allocations stay flat (~20KB) regardless of the + // calculator budgets — the flat allocation bound below asserts exactly that. + env, err := cel.NewEnv(Bindings()) + if err != nil { + t.Fatalf("cel.NewEnv(Bindings()) failed: %v", err) + } + expr := nestedBindExpr(8) + ast, iss := env.Compile(expr) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + + tests := []struct { + name string + opts []cel.ProgramOption + }{ + { + name: "default_calculator_limits", + opts: []cel.ProgramOption{cel.MemoryLimit(1_000_000)}, + }, + { + name: "custom_calculator_limits_100k_traversals_10_deep", + opts: []cel.ProgramOption{ + cel.MemoryTracking( + types.MemoryTrackerSizeCalculator(types.NewSizeCalculator( + types.SizeCalculatorMaxTraversal(100_000), + types.SizeCalculatorMaxDepth(10)))), + cel.MemoryLimit(1_000_000), + }, + }, + } + + // Subtests share process-global MemStats and must not run in parallel. + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + prg, err := env.Program(ast, tc.opts...) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + _, _, err = prg.Eval(cel.NoVars()) + runtime.ReadMemStats(&after) + + if err == nil || !strings.Contains(err.Error(), "memory limit exceeded") { + t.Fatalf("prg.Eval() got error %v, wanted error containing 'memory limit exceeded'", err) + } + allocated := after.TotalAlloc - before.TotalAlloc + t.Logf("evaluation allocated %d bytes", allocated) + // Measured allocation is ~20KB in either configuration; 1MiB provides ample + // headroom for incidental runtime allocations while still proving the sizing + // pass allocates nothing proportional to the traversal budget or the logical + // value size. + const maxAllocBytes = 1 << 20 // 1MiB + if allocated > maxAllocBytes { + t.Errorf("evaluation allocated %d bytes, wanted less than %d", allocated, maxAllocBytes) + } + }) + } +} + +// doublingStringBindExpr doubles the input string s through nested cel.bind calls: +// +// cel.bind(a0, s, cel.bind(a1, a0 + a0, ... cel.bind(aN, aN-1 + aN-1, aN == aN))) +// +// Unlike list concatenation, string concatenation materializes a new backing string at each +// level, so the physical allocation doubles alongside the logical size. +func doublingStringBindExpr(levels int) string { + expr := fmt.Sprintf("a%d == a%d", levels, levels) + for i := levels; i >= 1; i-- { + expr = fmt.Sprintf("cel.bind(a%d, a%d + a%d, %s)", i, i-1, i-1, expr) + } + return fmt.Sprintf("cel.bind(a0, s, %s)", expr) +} + +func TestBindingsMemoryLimitStringConcat(t *testing.T) { + // A 1MiB input string doubled per bind level. The doublings rapidly exceed the 1M unit + // memory limit, so evaluation must terminate while the bound values are still being + // constructed; unchecked, the doublings would allocate ~510MiB cumulative through a8. + env, err := cel.NewEnv(Bindings(), cel.Variable("s", cel.StringType)) + if err != nil { + t.Fatalf("cel.NewEnv(Bindings()) failed: %v", err) + } + in := map[string]any{"s": strings.Repeat("a", 1<<20)} + compile := func(levels int) *cel.Ast { + ast, iss := env.Compile(doublingStringBindExpr(levels)) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + return ast + } + measure := func(t *testing.T, prg cel.Program) (time.Duration, uint64, error) { + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + start := time.Now() + _, _, err := prg.Eval(in) + elapsed := time.Since(start) + runtime.ReadMemStats(&after) + return elapsed, after.TotalAlloc - before.TotalAlloc, err + } + + // Subtests share process-global MemStats and must not run in parallel. + t.Run("limit_trips_mid_amplification", func(t *testing.T) { + prg, err := env.Program(compile(8), cel.MemoryLimit(1_000_000)) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + elapsed, allocated, err := measure(t, prg) + if err == nil || !strings.Contains(err.Error(), "memory limit exceeded") { + t.Fatalf("prg.Eval() got error %v, wanted error containing 'memory limit exceeded'", err) + } + t.Logf("limit tripped in %v after allocating %d bytes", elapsed, allocated) + // The enforcement point follows the materialization of the offending value, so a + // few doublings of intermediates is the expected floor; the bound proves the + // remaining exponential trajectory (~510MiB unchecked) was halted. + const maxAllocBytes = 64 << 20 // 64MiB + if allocated > maxAllocBytes { + t.Errorf("evaluation allocated %d bytes, wanted less than %d", allocated, maxAllocBytes) + } + if elapsed > 5*time.Second { + t.Errorf("evaluation took %v, wanted under 5s", elapsed) + } + }) + + t.Run("unchecked_go_layer_baseline", func(t *testing.T) { + // The same workload the tripped case performed (doubling through 16MiB) without + // memory tracking: quantifies the Go-layer time and allocation the limit raced + // against, and the tracking overhead paid by the tripped case. + prg, err := env.Program(compile(4)) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + elapsed, allocated, err := measure(t, prg) + if err != nil { + t.Fatalf("prg.Eval() failed: %v", err) + } + t.Logf("unchecked baseline completed in %v allocating %d bytes", elapsed, allocated) + const maxAllocBytes = 64 << 20 // 64MiB + if allocated > maxAllocBytes { + t.Errorf("evaluation allocated %d bytes, wanted less than %d", allocated, maxAllocBytes) + } + }) +} + +func BenchmarkBindingsMemoryAmplification(b *testing.B) { + env, err := cel.NewEnv(Bindings()) + if err != nil { + b.Fatalf("cel.NewEnv(Bindings()) failed: %v", err) + } + compile := func(levels int) *cel.Ast { + ast, iss := env.Compile(nestedBindExpr(levels)) + if iss.Err() != nil { + b.Fatalf("env.Compile() failed: %v", iss.Err()) + } + return ast + } + ast2 := compile(2) + ast8 := compile(8) + customCalc := cel.MemoryTracking( + types.MemoryTrackerSizeCalculator(types.NewSizeCalculator( + types.SizeCalculatorMaxTraversal(100_000), + types.SizeCalculatorMaxDepth(10)))) + + benchmarks := []struct { + name string + ast *cel.Ast + opts []cel.ProgramOption + wantErr bool + }{ + // Successful evaluation of the 2-level expression, without and with tracking, to + // isolate the per-eval overhead of watermark observation and sizing. + {name: "levels_2_no_tracking", ast: ast2}, + {name: "levels_2_tracking", ast: ast2, opts: []cel.ProgramOption{cel.MemoryTracking()}}, + // Enforcement trip on the 8-level expression: sizing runs until the calculator's + // traversal budget aborts, the saturated watermark trips the limit, and evaluation + // unwinds through the cancellation panic. + { + name: "levels_8_limit_default_calculator", + ast: ast8, + opts: []cel.ProgramOption{cel.MemoryLimit(1_000_000)}, + wantErr: true, + }, + { + name: "levels_8_limit_custom_calculator_100k_10", + ast: ast8, + opts: []cel.ProgramOption{customCalc, cel.MemoryLimit(1_000_000)}, + wantErr: true, + }, + } + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + prg, err := env.Program(bm.ast, bm.opts...) + if err != nil { + b.Fatalf("env.Program() failed: %v", err) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _, err := prg.Eval(cel.NoVars()) + if bm.wantErr != (err != nil) { + b.Fatalf("prg.Eval() got err=%v, wantErr=%v", err, bm.wantErr) + } + } + }) + } +} + func TestBindingsNonMatch(t *testing.T) { env, err := cel.NewEnv(Bindings(), Strings()) if err != nil { diff --git a/interpreter/BUILD.bazel b/interpreter/BUILD.bazel index ee1f2ebda..6613d94e8 100644 --- a/interpreter/BUILD.bazel +++ b/interpreter/BUILD.bazel @@ -22,6 +22,7 @@ go_library( "planner.go", "prune.go", "runtimecost.go", + "runtimememory.go", ], importpath = "cel.dev/cel-go/interpreter", deps = [ @@ -55,6 +56,7 @@ go_test( "interpreter_test.go", "prune_test.go", "runtimecost_test.go", + "runtimememory_test.go", ], embed = [ ":go_default_library", diff --git a/interpreter/frame.go b/interpreter/frame.go index bc2b990f1..ff8191154 100644 --- a/interpreter/frame.go +++ b/interpreter/frame.go @@ -49,6 +49,9 @@ type evalContext struct { // costs provides the context for tracking the evaluation costs. costs *CostTracker + // memory provides the context for tracking peak memory during evaluation. + memory *types.MemoryTracker + // ctx is the context for async call implementations to use. ctx context.Context @@ -128,6 +131,7 @@ func (f *ExecutionFrame) Close() { f.ctx.interrupt = nil f.ctx.state = nil f.ctx.costs = nil + f.ctx.memory = nil f.ctx.interrupted.Store(false) f.ctx.interruptCheckCount.Store(0) f.ctx.interruptCheckFrequency = 0 diff --git a/interpreter/interpreter.go b/interpreter/interpreter.go index e5a583f6d..07297478d 100644 --- a/interpreter/interpreter.go +++ b/interpreter/interpreter.go @@ -76,6 +76,10 @@ const ( // CostLimitExceeded indicates that the operation was cancelled in response to the actual cost limit being // exceeded. CostLimitExceeded + + // MemoryLimitExceeded indicates that the operation was cancelled in response to the peak memory limit + // being exceeded. + MemoryLimitExceeded ) // evalStateOption configures the evalStateFactory behavior. diff --git a/interpreter/runtimememory.go b/interpreter/runtimememory.go new file mode 100644 index 000000000..4175fb6c9 --- /dev/null +++ b/interpreter/runtimememory.go @@ -0,0 +1,105 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package interpreter + +import ( + "errors" + + "cel.dev/cel-go/common/types" + "cel.dev/cel-go/common/types/ref" +) + +// memoryTrackPlanOption modifies the memory tracking factory associated with the MemoryObserver. +type memoryTrackPlanOption func(*memoryTrackerFactory) *memoryTrackerFactory + +// MemoryTrackerFactory configures the factory method to generate a new memory-tracker +// per-evaluation. +func MemoryTrackerFactory(factory func() (*types.MemoryTracker, error)) memoryTrackPlanOption { + return func(fac *memoryTrackerFactory) *memoryTrackerFactory { + fac.factory = factory + return fac + } +} + +// MemoryObserver provides an observer that tracks runtime peak memory. +func MemoryObserver(opts ...memoryTrackPlanOption) PlannerOption { + mt := &memoryTrackerFactory{} + for _, o := range opts { + mt = o(mt) + } + return func(p *planner) (*planner, error) { + if mt.factory == nil { + return nil, errors.New("memory tracker factory not configured") + } + p.observers = append(p.observers, mt) + return p, nil + } +} + +// memoryTrackerFactory holds a factory for producing new MemoryTracker instances on each Eval call. +type memoryTrackerFactory struct { + factory func() (*types.MemoryTracker, error) +} + +// InitState produces a MemoryTracker and bundles it into the ExecutionFrame in a way which is +// not visible to expression evaluation. +func (mt *memoryTrackerFactory) InitState(frame *ExecutionFrame) (any, error) { + if frame.ctx != nil && frame.ctx.memory != nil { + return frame.ctx.memory, nil + } + tracker, err := mt.factory() + if err != nil { + return nil, err + } + if frame.ctx == nil { + frame.ctx = evalContextPool.Get().(*evalContext) + } + frame.ctx.memory = tracker + return tracker, nil +} + +// GetState extracts the MemoryTracker from the ExecutionFrame. +func (mt *memoryTrackerFactory) GetState(frame *ExecutionFrame) any { + if frame == nil || frame.ctx == nil { + return nil + } + return frame.ctx.memory +} + +// Observe records the peak memory watermarks associated with each evaluation step. +// +// Watermarks are observed at the points where values materialize during evaluation: resolved +// attributes, function call results, constructed aggregate literals, and comprehension results. +// Since every intermediate Interpretable is observed, the inputs to a call contribute to the +// peak at the expression nodes which produced them; constants are part of the program image +// rather than runtime-materialized memory and are not observed. +func (mt *memoryTrackerFactory) Observe(vars Activation, id int64, programStep any, val ref.Val) { + frame := AsFrame(vars) + if frame == nil || frame.ctx == nil || frame.ctx.memory == nil { + return + } + tracker := frame.ctx.memory + switch programStep.(type) { + case InterpretableAttribute, InterpretableConst, InterpretableCall, InterpretableConstructor, *evalFold: + if frame.parent != nil { + tracker.Sample(id, val) + } else { + tracker.Track(val) + } + if tracker.ExceedsLimit() { + panic(EvalCancelledError{Cause: MemoryLimitExceeded, Message: "operation cancelled: memory limit exceeded"}) + } + } +} diff --git a/interpreter/runtimememory_test.go b/interpreter/runtimememory_test.go new file mode 100644 index 000000000..cf3d4b028 --- /dev/null +++ b/interpreter/runtimememory_test.go @@ -0,0 +1,570 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package interpreter + +import ( + "errors" + "strings" + "testing" + + "cel.dev/cel-go/checker" + "cel.dev/cel-go/common" + "cel.dev/cel-go/common/containers" + "cel.dev/cel-go/common/decls" + "cel.dev/cel-go/common/types" + "cel.dev/cel-go/common/types/ref" + "cel.dev/cel-go/parser" + proto3pb "cel.dev/cel-go/test/proto3pb" +) + +func TestMemoryObserver_FactoryNotConfigured(t *testing.T) { + cont := containers.DefaultContainer + reg := newTestRegistry(t) + attrs := NewAttributeFactory(cont, reg, reg) + interp := newStandardInterpreter(t, cont, reg, reg, attrs) + + s := common.NewTextSource(`1 + 1`) + p, err := parser.NewParser() + if err != nil { + t.Fatalf("parser.NewParser() failed: %v", err) + } + parsed, errs := p.Parse(s) + if len(errs.GetErrors()) != 0 { + t.Fatalf("Parse() failed: %v", errs.GetErrors()) + } + env := newTestEnv(t, cont, reg) + checked, errs := checker.Check(parsed, s, env) + if len(errs.GetErrors()) != 0 { + t.Fatalf("Check() failed: %v", errs.GetErrors()) + } + + _, err = interp.NewInterpretable(checked, MemoryObserver()) + if err == nil || !strings.Contains(err.Error(), "memory tracker factory not configured") { + t.Fatalf("NewInterpretable() got error %v, wanted 'memory tracker factory not configured'", err) + } +} + +func TestMemoryObserver_StateLifecycle(t *testing.T) { + fac := &memoryTrackerFactory{ + factory: func() (*types.MemoryTracker, error) { + return types.NewMemoryTracker(), nil + }, + } + + frame, err := NewExecutionFrame(EmptyActivation()) + if err != nil { + t.Fatalf("NewExecutionFrame() failed: %v", err) + } + defer frame.Close() + + if got := fac.GetState(nil); got != nil { + t.Errorf("GetState(nil) got %v, wanted nil", got) + } + if got := fac.GetState(frame); got != nil { + t.Errorf("GetState(uninitialized) got %v, wanted nil", got) + } + + state, err := fac.InitState(frame) + if err != nil { + t.Fatalf("InitState() failed: %v", err) + } + if state == nil { + t.Fatal("InitState() returned nil state") + } + + gotState := fac.GetState(frame) + if gotState != state { + t.Errorf("GetState() got %v, wanted %v", gotState, state) + } + + // Calling InitState a second time returns the existing state + state2, err := fac.InitState(frame) + if err != nil { + t.Fatalf("InitState() second call failed: %v", err) + } + if state2 != state { + t.Errorf("InitState() second call got %v, wanted %v", state2, state) + } + + // Test factory error propagation + errFac := &memoryTrackerFactory{ + factory: func() (*types.MemoryTracker, error) { + return nil, errors.New("factory failure") + }, + } + frame2, err := NewExecutionFrame(EmptyActivation()) + if err != nil { + t.Fatalf("NewExecutionFrame() failed: %v", err) + } + defer frame2.Close() + + _, err = errFac.InitState(frame2) + if err == nil || !strings.Contains(err.Error(), "factory failure") { + t.Fatalf("InitState() got error %v, wanted 'factory failure'", err) + } +} + +func TestMemoryObserver_Eval(t *testing.T) { + tests := []struct { + name string + expr string + vars []*decls.VariableDecl + in any + memOpts []types.MemoryTrackerOption + wantPeak uint32 + }{ + { + name: "ident_attribute", + expr: `a`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3}}, + wantPeak: 4, + }, + { + name: "field_selection", + expr: `m.vals`, + vars: []*decls.VariableDecl{decls.NewVariable("m", types.NewMapType(types.StringType, types.NewListType(types.IntType)))}, + in: map[string]any{"m": map[string][]int64{"vals": {1, 2, 3}}}, + wantPeak: 4, + }, + { + name: "call_list_concat", + expr: `a + a`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2}}, + wantPeak: 6, + }, + { + name: "call_string_concat", + expr: `s + s`, + vars: []*decls.VariableDecl{decls.NewVariable("s", types.StringType)}, + in: map[string]any{"s": strings.Repeat("a", 50)}, + wantPeak: 10, + }, + { + name: "list_literal", + expr: `[1, 2, 3, 4]`, + wantPeak: 5, + }, + { + name: "map_literal", + expr: `{'k1': 1, 'k2': 2}`, + wantPeak: 5, + }, + { + name: "comprehension_map", + expr: `a.map(x, x * 2)`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5}}, + wantPeak: 6, + }, + { + name: "comprehension_filter", + expr: `a.filter(x, x % 2 == 0)`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5}}, + wantPeak: 3, + }, + { + name: "comprehension_exists", + expr: `a.exists(x, x == 3)`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5}}, + wantPeak: 1, + }, + { + name: "comprehension_all", + expr: `a.all(x, x > 0)`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5}}, + wantPeak: 1, + }, + { + name: "comprehension_sampled", + expr: `a.map(x, x * 2)`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5}}, + memOpts: []types.MemoryTrackerOption{types.MemoryTrackerSampleInterval(3)}, + wantPeak: 6, + }, + { + name: "comprehension_single_element", + expr: `[1].map(x, [x, x])`, + memOpts: []types.MemoryTrackerOption{types.MemoryTrackerSampleInterval(10)}, + wantPeak: 3, + }, + { + name: "comprehension_nested_sampled", + expr: `a.map(x, [x, x]).filter(z, z.size() == 2)`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5}}, + memOpts: []types.MemoryTrackerOption{types.MemoryTrackerSampleInterval(2)}, + wantPeak: 6, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tracker, res, err := evalTestMemoryTracker(t, tc.expr, tc.vars, tc.in, tc.memOpts...) + if err != nil { + t.Fatalf("evalTestMemoryTracker() failed: %v", err) + } + if types.IsError(res) { + t.Fatalf("eval result error: %v", res) + } + if tracker == nil { + t.Fatal("expected non-nil MemoryTracker") + } + if tracker.Peak() < tc.wantPeak { + t.Errorf("tracker.Peak() got %d, want at least %d", tracker.Peak(), tc.wantPeak) + } + }) + } +} + +func TestMemoryObserver_LimitExceeded(t *testing.T) { + t.Run("over_limit_panics", func(t *testing.T) { + defer func() { + r := recover() + if r == nil { + t.Fatal("expected evaluation to panic on memory limit exceeded") + } + evalErr, ok := r.(EvalCancelledError) + if !ok { + t.Fatalf("expected EvalCancelledError, got %T: %v", r, r) + } + if evalErr.Cause != MemoryLimitExceeded { + t.Errorf("EvalCancelledError.Cause got %v, want %v", evalErr.Cause, MemoryLimitExceeded) + } + }() + + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2, 3}} + _, _, _ = evalTestMemoryTracker(t, `a + a`, vars, in, types.MemoryTrackerLimit(5)) + }) + + t.Run("under_limit_succeeds", func(t *testing.T) { + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2}} + tracker, res, err := evalTestMemoryTracker(t, `a + a`, vars, in, types.MemoryTrackerLimit(100)) + if err != nil { + t.Fatalf("eval failed: %v", err) + } + if types.IsError(res) { + t.Fatalf("eval result error: %v", res) + } + if tracker.ExceedsLimit() { + t.Error("tracker.ExceedsLimit() got true, want false") + } + }) +} + +func evalTestMemoryTracker(t testing.TB, expr string, vars []*decls.VariableDecl, in any, memOpts ...types.MemoryTrackerOption) (*types.MemoryTracker, ref.Val, error) { + t.Helper() + s := common.NewTextSource(expr) + p, err := parser.NewParser(parser.Macros(parser.AllMacros...)) + if err != nil { + t.Fatalf("Failed to initialize parser: %v", err) + } + parsed, errs := p.Parse(s) + if len(errs.GetErrors()) != 0 { + t.Fatalf("Failed to parse expression %q: %v", expr, errs.GetErrors()) + } + + cont := containers.DefaultContainer + reg := newTestRegistry(t) + attrs := NewAttributeFactory(cont, reg, reg) + env := newTestEnv(t, cont, reg) + if len(vars) > 0 { + if err := env.AddIdents(vars...); err != nil { + t.Fatalf("Failed to add variables: %v", err) + } + } + checked, errs := checker.Check(parsed, s, env) + if len(errs.GetErrors()) != 0 { + t.Fatalf("Failed to check expression %q: %v", expr, errs.GetErrors()) + } + + var tracker *types.MemoryTracker + factory := func() (*types.MemoryTracker, error) { + tracker = types.NewMemoryTracker(memOpts...) + return tracker, nil + } + + disp := NewDispatcher() + addFunctionBindings(t, disp) + interp := NewInterpreter(disp, cont, reg, reg, attrs) + prg, err := interp.NewInterpretable(checked, Optimize(), MemoryObserver(MemoryTrackerFactory(factory))) + if err != nil { + return nil, nil, err + } + + act := constructTestActivation(t, in) + res := prg.Eval(act) + return tracker, res, nil +} + +func constructTestActivation(t testing.TB, in any) Activation { + t.Helper() + if in == nil { + return EmptyActivation() + } + a, err := NewActivation(in) + if err != nil { + t.Fatalf("NewActivation(%v) failed: %v", in, err) + } + return a +} + +func benchmarkMemoryTracker(b *testing.B, expr string, vars []*decls.VariableDecl, in any, enableTracking bool, memOpts ...types.MemoryTrackerOption) { + b.Helper() + s := common.NewTextSource(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("Failed to parse expression %q: %v", expr, errs.GetErrors()) + } + + cont := containers.DefaultContainer + reg := newTestRegistry(b) + attrs := NewAttributeFactory(cont, reg, reg) + env := newTestEnv(b, cont, reg) + if len(vars) > 0 { + if err := env.AddIdents(vars...); err != nil { + b.Fatalf("Failed to add variables: %v", err) + } + } + checked, errs := checker.Check(parsed, s, env) + if len(errs.GetErrors()) != 0 { + b.Fatalf("Failed to check expression %q: %v", expr, errs.GetErrors()) + } + + disp := NewDispatcher() + addFunctionBindings(b, disp) + interp := NewInterpreter(disp, cont, reg, reg, attrs) + + planOpts := []PlannerOption{Optimize()} + if enableTracking { + factory := func() (*types.MemoryTracker, error) { + return types.NewMemoryTracker(memOpts...), nil + } + planOpts = append(planOpts, MemoryObserver(MemoryTrackerFactory(factory))) + } + + prg, err := interp.NewInterpretable(checked, planOpts...) + if err != nil { + b.Fatalf("NewInterpretable() failed: %v", err) + } + + frame := AsFrame(constructTestActivation(b, in)) + + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + prg.Eval(frame) + } +} + +func BenchmarkMemoryTracker_SimpleCall_TrackingDisabled(b *testing.B) { + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2, 3, 4, 5}} + benchmarkMemoryTracker(b, `a + a`, vars, in, false) +} + +func BenchmarkMemoryTracker_SimpleCall_TrackingEnabled(b *testing.B) { + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2, 3, 4, 5}} + benchmarkMemoryTracker(b, `a + a`, vars, in, true) +} + +func BenchmarkMemoryTracker_Constructor_TrackingDisabled(b *testing.B) { + benchmarkMemoryTracker(b, `[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]`, nil, nil, false) +} + +func BenchmarkMemoryTracker_Constructor_TrackingEnabled(b *testing.B) { + benchmarkMemoryTracker(b, `[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]`, nil, nil, true) +} + +func BenchmarkMemoryTracker_Comprehension_TrackingDisabled(b *testing.B) { + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}} + benchmarkMemoryTracker(b, `a.map(x, x * 2)`, vars, in, false) +} + +func BenchmarkMemoryTracker_Comprehension_SampleInterval1(b *testing.B) { + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}} + benchmarkMemoryTracker(b, `a.map(x, x * 2)`, vars, in, true, types.MemoryTrackerSampleInterval(1)) +} + +func BenchmarkMemoryTracker_Comprehension_SampleInterval5(b *testing.B) { + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}} + benchmarkMemoryTracker(b, `a.map(x, x * 2)`, vars, in, true, types.MemoryTrackerSampleInterval(5)) +} + +func BenchmarkMemoryTracker_NestedComprehension_TrackingDisabled(b *testing.B) { + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2, 3, 4, 5}} + benchmarkMemoryTracker(b, `a.map(x, [x, x]).filter(z, z.size() == 2)`, vars, in, false) +} + +func BenchmarkMemoryTracker_NestedComprehension_TrackingEnabled(b *testing.B) { + vars := []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))} + in := map[string]any{"a": []int64{1, 2, 3, 4, 5}} + benchmarkMemoryTracker(b, `a.map(x, [x, x]).filter(z, z.size() == 2)`, vars, in, true) +} + +type optionsScenario struct { + name string + trackMem bool + memSample uint32 + trackCost bool + trackTrac bool +} + +func runOptionsBenchmark(b *testing.B, expr string, vars []*decls.VariableDecl, in any, sc optionsScenario) { + b.Helper() + s := common.NewTextSource(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("Failed to parse expression %q: %v", 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(vars) > 0 { + if err := env.AddIdents(vars...); err != nil { + b.Fatalf("Failed to add variables: %v", err) + } + } + checked, errs := checker.Check(parsed, s, env) + if len(errs.GetErrors()) != 0 { + b.Fatalf("Failed to check expression %q: %v", expr, errs.GetErrors()) + } + + disp := NewDispatcher() + addFunctionBindings(b, disp) + interp := NewInterpreter(disp, cont, reg, reg, attrs) + + planOpts := []PlannerOption{Optimize()} + if sc.trackMem { + factory := func() (*types.MemoryTracker, error) { + var memOpts []types.MemoryTrackerOption + if sc.memSample > 0 { + memOpts = append(memOpts, types.MemoryTrackerSampleInterval(sc.memSample)) + } + return types.NewMemoryTracker(memOpts...), nil + } + planOpts = append(planOpts, MemoryObserver(MemoryTrackerFactory(factory))) + } + if sc.trackCost { + costFac := func() (*CostTracker, error) { + return NewCostTracker(nil) + } + planOpts = append(planOpts, CostObserver(CostTrackerFactory(costFac))) + } + if sc.trackTrac { + planOpts = append(planOpts, EvalStateObserver()) + } + + prg, err := interp.NewInterpretable(checked, planOpts...) + if err != nil { + b.Fatalf("NewInterpretable() failed: %v", err) + } + + frame := AsFrame(constructTestActivation(b, in)) + + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + prg.Eval(frame) + } +} + +func BenchmarkOptionsComparison(b *testing.B) { + scenarios := []optionsScenario{ + {name: "None"}, + {name: "CostTracking", trackCost: true}, + {name: "MemoryTracking_Sample1", trackMem: true, memSample: 1}, + {name: "MemoryTracking_Sample5", trackMem: true, memSample: 5}, + {name: "Tracing_State", trackTrac: true}, + {name: "Cost_and_Memory", trackCost: true, trackMem: true, memSample: 1}, + {name: "Cost_and_Tracing", trackCost: true, trackTrac: true}, + {name: "Memory_and_Tracing", trackMem: true, memSample: 1, trackTrac: true}, + {name: "All_Options", trackCost: true, trackMem: true, memSample: 1, trackTrac: true}, + } + + workloads := []struct { + name string + expr string + vars []*decls.VariableDecl + in any + }{ + { + name: "SimpleCall", + expr: `a + a`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5}}, + }, + { + name: "LiteralConstructor", + expr: `[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]`, + }, + { + name: "MapComprehension_10", + expr: `a.map(x, x * 2)`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}, + }, + { + name: "NestedComprehension_10", + expr: `a.map(x, [x, x]).filter(z, z.size() == 2)`, + vars: []*decls.VariableDecl{decls.NewVariable("a", types.NewListType(types.IntType))}, + in: map[string]any{"a": []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}, + }, + { + name: "ProtoFieldAccess", + expr: `msg.single_int64 + msg.single_int32`, + vars: []*decls.VariableDecl{decls.NewVariable("msg", types.NewObjectType("google.expr.proto3.test.TestAllTypes"))}, + in: map[string]any{"msg": &proto3pb.TestAllTypes{SingleInt64: 42, SingleInt32: 10, SingleString: "hello world"}}, + }, + { + name: "ProtoRepeatedComprehension_10", + expr: `msg.repeated_int64.map(x, x * 2)`, + vars: []*decls.VariableDecl{decls.NewVariable("msg", types.NewObjectType("google.expr.proto3.test.TestAllTypes"))}, + in: map[string]any{"msg": &proto3pb.TestAllTypes{RepeatedInt64: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}}, + }, + } + + for _, w := range workloads { + b.Run(w.name, func(b *testing.B) { + for _, sc := range scenarios { + b.Run(sc.name, func(b *testing.B) { + runOptionsBenchmark(b, w.expr, w.vars, w.in, sc) + }) + } + }) + } +}