Improve size calculation and memoization (#1431)

Author: TristonianJones

common/types/BUILD.bazel

@@ -22,6 +22,7 @@ go_library(
         "format.go",
         "list.go",
         "map.go",
+        "memory_tracker.go",
         "native.go",
         "null.go",
         "object.go",
@@ -72,6 +73,7 @@ go_test(
         "json_struct_test.go",
         "list_test.go",
         "map_test.go",
+        "memory_tracker_test.go",
         "native_test.go",
         "null_test.go",
         "object_test.go",

common/types/list.go

@@ -274,9 +274,17 @@ func (l *baseList) AggregateSize(sizer AggregateSizer) uint32 {
 	if sz := atomic.LoadUint32(&l.aggSize); sz != 0 {
 		return sz
 	}
-	total := uint32(1)
-	for i := range l.size {
-		total = safeAddUint32(total, sizer.AggregateSize(l.get(i)))
+	var total uint32
+	if l.value != nil {
+		if t, ok := getSliceElementsAggregateSize(sizer, l.value); ok {
+			total = t
+		}
+	}
+	if total == 0 {
+		total = uint32(1)
+		for i := range l.size {
+			total = safeAddUint32(total, sizer.AggregateSize(l.get(i)))
+		}
 	}
 	if cacheableAggregateSize(sizer) {
 		atomic.StoreUint32(&l.aggSize, total)
@@ -365,6 +373,7 @@ type concatList struct {
 	prevList   traits.Lister
 	nextList   traits.Lister
 	cachedSize ref.Val
+	aggSize    uint32
 }
 
 func newConcatList(adapter Adapter, prevList, nextList traits.Lister) ref.Val {
@@ -499,7 +508,14 @@ func (l *concatList) Size() ref.Val {
 
 // AggregateSize implements the AggregateSizeVisitor interface method.
 func (l *concatList) AggregateSize(sizer AggregateSizer) uint32 {
-	return safeAddUint32(sizer.AggregateSize(l.prevList), sizer.AggregateSize(l.nextList))
+	if sz := atomic.LoadUint32(&l.aggSize); sz != 0 {
+		return sz
+	}
+	total := safeAddUint32(sizer.AggregateSize(l.prevList), sizer.AggregateSize(l.nextList))
+	if cacheableAggregateSize(sizer) {
+		atomic.StoreUint32(&l.aggSize, total)
+	}
+	return total
 }
 
 // String converts the concatenated list to a human-readable string.

common/types/map.go

@@ -314,12 +314,21 @@ func (m *baseMap) AggregateSize(sizer AggregateSizer) uint32 {
 	if sz := atomic.LoadUint32(&m.aggSize); sz != 0 {
 		return sz
 	}
-	f := foldableAggregateSizer{sizer: sizer, total: 1}
-	m.Fold(&f)
+	var total uint32
+	if m.value != nil {
+		if t, ok := getMapElementsAggregateSize(sizer, m.value); ok {
+			total = t
+		}
+	}
+	if total == 0 {
+		f := foldableAggregateSizer{sizer: sizer, total: 1}
+		m.Fold(&f)
+		total = f.total
+	}
 	if cacheableAggregateSize(sizer) {
-		atomic.StoreUint32(&m.aggSize, f.total)
+		atomic.StoreUint32(&m.aggSize, total)
 	}
-	return f.total
+	return total
 }
 
 // String converts the map into a human-readable string.
@@ -331,7 +340,7 @@ func (m *baseMap) String() string {
 	for it.HasNext() == True {
 		k := it.Next()
 		v, _ := m.Find(k)
-		sb.WriteString(fmt.Sprintf("%v: %v", k, v))
+		fmt.Fprintf(&sb, "%v: %v", k, v)
 		if i != m.size-1 {
 			sb.WriteString(", ")
 		}
@@ -710,7 +719,8 @@ func (a *stringIfaceMapAccessor) Fold(f traits.Folder) {
 // accessing protoreflect.Map values.
 type protoMap struct {
 	Adapter
-	value *pb.Map
+	value   *pb.Map
+	aggSize uint32
 }
 
 // Contains returns whether the map contains the given key.
@@ -935,12 +945,18 @@ func (m *protoMap) AggregateSize(sizer AggregateSizer) uint32 {
 	if m.value == nil {
 		return 0
 	}
+	if sz := atomic.LoadUint32(&m.aggSize); sz != 0 {
+		return sz
+	}
 	total := uint32(1)
 	m.value.Range(func(k protoreflect.MapKey, v protoreflect.Value) bool {
 		total = safeAddUint32(total, sizer.AggregateSize(k))
 		total = safeAddUint32(total, sizer.AggregateSize(v))
 		return true
 	})
+	if cacheableAggregateSize(sizer) {
+		atomic.StoreUint32(&m.aggSize, total)
+	}
 	return total
 }
 

common/types/memory_tracker.go

@@ -14,6 +14,10 @@
 
 package types
 
+import (
+	"cel.dev/cel-go/common/types/ref"
+)
+
 const (
 	defaultMemoryTrackerSampleInterval = 1
 )
@@ -27,7 +31,8 @@ type MemoryTrackerOption func(*MemoryTracker)
 // tracking observations and terminate evaluation as appropriate.
 func MemoryTrackerLimit(limit uint32) MemoryTrackerOption {
 	return func(t *MemoryTracker) {
-		t.limit = &limit
+		t.limit = limit
+		t.hasLimit = true
 	}
 }
 
@@ -57,24 +62,23 @@ func MemoryTrackerSizeCalculator(calc *SizeCalculator) MemoryTrackerOption {
 //
 // Memory is measured in aggregate element counts as computed by a SizeCalculator, with all
 // arithmetic saturating at math.MaxUint32. The tracker is independent of any interpreter
-// implementation; evaluators feed it observations at points where values materialize:
+// implementation; evaluators feed it observations at the points where values materialize
+// during evaluation, such as resolved attributes, call results, constructed aggregates, and
+// values built up within comprehensions or bind initializers.
 //
-//   - inputs and outputs of calls (e.g. the target and arguments of a.join(', '))
-//   - resolved attribute values (e.g. the value of a.b.c)
-//   - values built up within literal blocks, comprehensions, and bind initializers
-//
-// The peak is the largest single observation, where one Track call observes a set of
-// coexistent values as a single watermark.
+// The peak is the largest single observation, where one Track call observes a value
+// as a watermark.
 //
 // A MemoryTracker is stateful and intended for use by a single evaluation at a time; it is
 // not safe for concurrent use.
 type MemoryTracker struct {
 	version        int
 	calc           *SizeCalculator
-	limit          *uint32
+	limit          uint32
+	hasLimit       bool
 	sampleInterval uint32
 
-	sampleCount       uint32
+	sampleCounts      map[int64]uint32
 	peak              uint32
 	calcLimitExceeded bool
 }
@@ -100,41 +104,60 @@ func (t *MemoryTracker) Version() int {
 	return t.version
 }
 
-// Track observes a set of coexistent values as a single watermark, returning their combined
-// aggregate size with saturation at math.MaxUint32.
-//
-// Call sites with multiple live values, such as the input arguments to a function call,
-// should be tracked in a single call so the watermark reflects their combined footprint.
-//
-// Within observers, consider whether to choose the aggregated argument sizes, the result size,
-// or both when working with allocating operations.
-func (t *MemoryTracker) Track(vals ...any) uint32 {
-	total := uint32(0)
-	for _, val := range vals {
-		est := t.calc.EstimateAggregateSize(val)
-		if est.LimitExceeded {
-			t.calcLimitExceeded = true
+// Track observes a value as a watermark, returning its aggregate size with saturation at math.MaxUint32.
+func (t *MemoryTracker) Track(val ref.Val) uint32 {
+	if val == nil {
+		return 0
+	}
+	switch v := val.(type) {
+	case Bool, Int, Uint, Double, Duration, Timestamp, Null, *Type, *Err, *Unknown:
+		if 1 > t.peak {
+			t.peak = 1
 		}
-		total = safeAddUint32(total, est.Size)
+		return 1
+	case String:
+		sz := t.calc.stringSize(len(v))
+		if sz > t.peak {
+			t.peak = sz
+		}
+		return sz
+	case Bytes:
+		sz := t.calc.stringSize(len(v))
+		if sz > t.peak {
+			t.peak = sz
+		}
+		return sz
+	}
+
+	est := t.calc.ApproximateAggregateSize(val)
+	if est.LimitExceeded {
+		t.calcLimitExceeded = true
 	}
-	if total > t.peak {
-		t.peak = total
+	if est.Size > t.peak {
+		t.peak = est.Size
 	}
-	return total
+	return est.Size
 }
 
-// Sample observes a value subject to the tracker's sample interval, returning the value's
-// aggregate size when computed, or zero when the observation is skipped.
+// Sample observes a value for a specific node ID subject to the tracker's sample interval,
+// returning the value's aggregate size when computed, or zero when the observation is skipped.
+//
+// The first observation of any node ID is always tracked; subsequent observations are sampled
+// at multiples of the sample interval.
 //
 // Sample is intended for high-frequency observation points, such as accumulator values built
 // up by comprehension loops or bind initializers, where sizing every iteration would be
 // prohibitively expensive.
-func (t *MemoryTracker) Sample(val any) uint32 {
-	t.sampleCount++
-	if t.sampleInterval > 1 && t.sampleCount%t.sampleInterval != 0 {
-		return 0
+func (t *MemoryTracker) Sample(id int64, val ref.Val) uint32 {
+	if t.sampleCounts == nil {
+		t.sampleCounts = make(map[int64]uint32)
+	}
+	t.sampleCounts[id]++
+	count := t.sampleCounts[id]
+	if count == 1 || t.sampleInterval <= 1 || count%t.sampleInterval == 0 {
+		return t.Track(val)
 	}
-	return t.Track(val)
+	return 0
 }
 
 // Peak returns the largest single watermark observed, saturating at math.MaxUint32.
@@ -145,7 +168,7 @@ func (t *MemoryTracker) Peak() uint32 {
 // ExceedsLimit indicates whether the peak observed memory exceeds the configured limit.
 // When no limit is configured, ExceedsLimit always returns false.
 func (t *MemoryTracker) ExceedsLimit() bool {
-	return t.limit != nil && t.peak > *t.limit
+	return t.hasLimit && t.peak > t.limit
 }
 
 // CalculationLimitExceeded indicates whether any tracked value was too expensive to size,

common/types/memory_tracker_test.go

@@ -37,16 +37,18 @@ func TestMemoryTrackerTrack(t *testing.T) {
 		}
 	})
 
-	t.Run("call_inputs_combined_watermark", func(t *testing.T) {
-		// Models the inputs to a.join(', ') where the list and separator coexist.
+	t.Run("consecutive_watermarks", func(t *testing.T) {
 		tracker := NewMemoryTracker()
-		list := NewRefValList(adapter, []ref.Val{String("a"), String("b")})
-		sep := String(", ")
-		if got := tracker.Track(list, sep); got != 4 {
-			t.Errorf("Track(list, sep) got %d, want 4 (3 list + 1 separator)", got)
+		list := NewRefValList(adapter, []ref.Val{Int(1), Int(2)})
+		if got := tracker.Track(list); got != 3 {
+			t.Errorf("Track(list) got %d, want 3", got)
 		}
-		if got := tracker.Peak(); got != 4 {
-			t.Errorf("Peak() got %d, want 4", got)
+		arg := Int(0)
+		if got := tracker.Track(arg); got != 1 {
+			t.Errorf("Track(arg) got %d, want 1", got)
+		}
+		if got := tracker.Peak(); got != 3 {
+			t.Errorf("Peak() got %d, want 3", got)
 		}
 	})
 
@@ -63,18 +65,18 @@ func TestMemoryTrackerTrack(t *testing.T) {
 		// Models the output of a + a, a string twice the size of its inputs.
 		tracker := NewMemoryTracker()
 		in := String(strings.Repeat("a", 50))
-		tracker.Track(in, in)
+		tracker.Track(in)
 		out := String(strings.Repeat("a", 100))
 		tracker.Track(out)
 		if got := tracker.Peak(); got != 10 {
 			t.Errorf("Peak() got %d, want 10 (100-char output at 10 chars per unit)", got)
 		}
 	})
 
-	t.Run("saturating_sum", func(t *testing.T) {
+	t.Run("saturating_value", func(t *testing.T) {
 		tracker := NewMemoryTracker()
 		big := customSizerVal(math.MaxUint32)
-		if got := tracker.Track(big, big); got != math.MaxUint32 {
+		if got := tracker.Track(big); got != math.MaxUint32 {
 			t.Errorf("Track() got %d, want MaxUint32", got)
 		}
 		if got := tracker.Peak(); got != math.MaxUint32 {
@@ -104,7 +106,7 @@ func TestMemoryTrackerSample(t *testing.T) {
 	t.Run("default_interval_samples_every_value", func(t *testing.T) {
 		tracker := NewMemoryTracker()
 		for i := 0; i < 3; i++ {
-			if got := tracker.Sample(Int(i)); got != 1 {
+			if got := tracker.Sample(1, Int(i)); got != 1 {
 				t.Errorf("Sample() got %d, want 1", got)
 			}
 		}
@@ -116,23 +118,48 @@ func TestMemoryTrackerSample(t *testing.T) {
 	t.Run("interval_skips_intermediate_samples", func(t *testing.T) {
 		tracker := NewMemoryTracker(MemoryTrackerSampleInterval(3))
 		list := NewRefValList(adapter, []ref.Val{Int(1), Int(2)})
-		if got := tracker.Sample(list); got != 0 {
-			t.Errorf("Sample() #1 got %d, want 0 (skipped)", got)
+		if got := tracker.Sample(1, list); got != 3 {
+			t.Errorf("Sample() #1 got %d, want 3 (computed on first observation)", got)
 		}
-		if got := tracker.Sample(list); got != 0 {
+		if got := tracker.Sample(1, list); got != 0 {
 			t.Errorf("Sample() #2 got %d, want 0 (skipped)", got)
 		}
-		if got := tracker.Sample(list); got != 3 {
+		if got := tracker.Sample(1, list); got != 3 {
 			t.Errorf("Sample() #3 got %d, want 3 (computed)", got)
 		}
 		if got := tracker.Peak(); got != 3 {
 			t.Errorf("Peak() got %d, want 3", got)
 		}
 	})
 
+	t.Run("per_id_tracking", func(t *testing.T) {
+		tracker := NewMemoryTracker(MemoryTrackerSampleInterval(2))
+		list := NewRefValList(adapter, []ref.Val{Int(1), Int(2)})
+		// id 1: sample 1 (computed on first observation)
+		if got := tracker.Sample(1, list); got != 3 {
+			t.Errorf("Sample(1) #1 got %d, want 3", got)
+		}
+		// id 2: sample 1 (computed on first observation)
+		if got := tracker.Sample(2, list); got != 3 {
+			t.Errorf("Sample(2) #1 got %d, want 3", got)
+		}
+		// id 1: sample 2 (computed, multiple of 2)
+		if got := tracker.Sample(1, list); got != 3 {
+			t.Errorf("Sample(1) #2 got %d, want 3", got)
+		}
+		// id 2: sample 2 (computed, multiple of 2)
+		if got := tracker.Sample(2, list); got != 3 {
+			t.Errorf("Sample(2) #2 got %d, want 3", got)
+		}
+		// id 1: sample 3 (skipped)
+		if got := tracker.Sample(1, list); got != 0 {
+			t.Errorf("Sample(1) #3 got %d, want 0", got)
+		}
+	})
+
 	t.Run("zero_interval_clamped_to_one", func(t *testing.T) {
 		tracker := NewMemoryTracker(MemoryTrackerSampleInterval(0))
-		if got := tracker.Sample(Int(1)); got != 1 {
+		if got := tracker.Sample(1, Int(1)); got != 1 {
 			t.Errorf("Sample() got %d, want 1", got)
 		}
 	})
@@ -172,4 +199,5 @@ func TestMemoryTrackerVersion(t *testing.T) {
 
 // Interface conformance check: the tracker's calculator remains usable as an AggregateSizer
 // by visitor implementations.
+var _ ref.Val = customSizerVal(0)
 var _ traits.Sizer = customSizerVal(0)