mem-track: Peak memory tracking options with per-id sampling

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)
+	}
+}

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

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 {