Skip to content

Commit 9329452

Browse files
TristonianJonesTristonianJones
committed
Scale sizes for strings and bytes, add memory tracker
1 parent f513cf8 commit 9329452

4 files changed

Lines changed: 334 additions & 3 deletions

File tree

‎common/types/memory_tracker.go‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package types
16+
17+
const (
18+
defaultMemoryTrackerSampleInterval = 1
19+
)
20+
21+
// MemoryTrackerOption configures a MemoryTracker instance.
22+
type MemoryTrackerOption func(*MemoryTracker)
23+
24+
// MemoryTrackerLimit sets a limit on the peak aggregate memory observed during tracking.
25+
//
26+
// The tracker does not enforce the limit itself; callers should consult ExceedsLimit after
27+
// tracking observations and terminate evaluation as appropriate.
28+
func MemoryTrackerLimit(limit uint32) MemoryTrackerOption {
29+
return func(t *MemoryTracker) {
30+
t.limit = &limit
31+
}
32+
}
33+
34+
// MemoryTrackerSampleInterval configures how frequently Sample observations compute a size.
35+
//
36+
// An interval of N means every Nth call to Sample performs a size computation; intervening
37+
// calls are skipped. Values less than 1 are treated as 1, meaning every sample is computed.
38+
// Sampling bounds the tracking overhead for high-frequency observation points such as
39+
// comprehension loops and bind initializers.
40+
func MemoryTrackerSampleInterval(interval uint32) MemoryTrackerOption {
41+
return func(t *MemoryTracker) {
42+
if interval < 1 {
43+
interval = 1
44+
}
45+
t.sampleInterval = interval
46+
}
47+
}
48+
49+
// MemoryTrackerSizeCalculator overrides the SizeCalculator used to compute value sizes.
50+
func MemoryTrackerSizeCalculator(calc *SizeCalculator) MemoryTrackerOption {
51+
return func(t *MemoryTracker) {
52+
t.calc = calc
53+
}
54+
}
55+
56+
// MemoryTracker records the peak aggregate memory observed during an evaluation.
57+
//
58+
// Memory is measured in aggregate element counts as computed by a SizeCalculator, with all
59+
// arithmetic saturating at math.MaxUint32. The tracker is independent of any interpreter
60+
// implementation; evaluators feed it observations at points where values materialize:
61+
//
62+
// - inputs to a call (e.g. the target and arguments of a.join(', '))
63+
// - the output of a call (e.g. the result of a + a)
64+
// - resolved attribute values (e.g. the value of a.b.c)
65+
// - sampled values built up within comprehensions or bind initializers
66+
//
67+
// The peak is the largest single observation, where one Track call observes a set of
68+
// coexistent values as a single watermark.
69+
//
70+
// A MemoryTracker is stateful and intended for use by a single evaluation at a time; it is
71+
// not safe for concurrent use.
72+
type MemoryTracker struct {
73+
version int
74+
calc *SizeCalculator
75+
limit *uint32
76+
sampleInterval uint32
77+
78+
sampleCount uint32
79+
peak uint32
80+
calcLimitExceeded bool
81+
}
82+
83+
// NewMemoryTracker returns a new MemoryTracker configured with optional MemoryTrackerOption
84+
// settings, using a default SizeCalculator when one is not provided.
85+
func NewMemoryTracker(opts ...MemoryTrackerOption) *MemoryTracker {
86+
t := &MemoryTracker{
87+
version: 0,
88+
sampleInterval: defaultMemoryTrackerSampleInterval,
89+
}
90+
for _, opt := range opts {
91+
opt(t)
92+
}
93+
if t.calc == nil {
94+
t.calc = NewSizeCalculator()
95+
}
96+
return t
97+
}
98+
99+
// Version returns the tracking version.
100+
func (t *MemoryTracker) Version() int {
101+
return t.version
102+
}
103+
104+
// Track observes a set of coexistent values as a single watermark, returning their combined
105+
// aggregate size with saturation at math.MaxUint32.
106+
//
107+
// Call sites with multiple live values, such as the input arguments to a function call,
108+
// should be tracked in a single call so the watermark reflects their combined footprint.
109+
func (t *MemoryTracker) Track(vals ...any) uint32 {
110+
total := uint32(0)
111+
for _, val := range vals {
112+
est := t.calc.EstimateAggregateSize(val)
113+
if est.LimitExceeded {
114+
t.calcLimitExceeded = true
115+
}
116+
total = safeAddUint32(total, est.Size)
117+
}
118+
if total > t.peak {
119+
t.peak = total
120+
}
121+
return total
122+
}
123+
124+
// Sample observes a value subject to the tracker's sample interval, returning the value's
125+
// aggregate size when computed, or zero when the observation is skipped.
126+
//
127+
// Sample is intended for high-frequency observation points, such as accumulator values built
128+
// up by comprehension loops or bind initializers, where sizing every iteration would be
129+
// prohibitively expensive.
130+
func (t *MemoryTracker) Sample(val any) uint32 {
131+
t.sampleCount++
132+
if t.sampleInterval > 1 && t.sampleCount%t.sampleInterval != 0 {
133+
return 0
134+
}
135+
return t.Track(val)
136+
}
137+
138+
// Peak returns the largest single watermark observed, saturating at math.MaxUint32.
139+
func (t *MemoryTracker) Peak() uint32 {
140+
return t.peak
141+
}
142+
143+
// ExceedsLimit indicates whether the peak observed memory exceeds the configured limit.
144+
// When no limit is configured, ExceedsLimit always returns false.
145+
func (t *MemoryTracker) ExceedsLimit() bool {
146+
return t.limit != nil && t.peak > *t.limit
147+
}
148+
149+
// CalculationLimitExceeded indicates whether any tracked value was too expensive to size,
150+
// causing the size computation to abort at the SizeCalculator's depth or traversal limits.
151+
//
152+
// Such observations saturate to math.MaxUint32; this signal distinguishes values which were
153+
// too costly to measure from values whose measured size genuinely saturated uint32.
154+
func (t *MemoryTracker) CalculationLimitExceeded() bool {
155+
return t.calcLimitExceeded
156+
}

‎common/types/memory_tracker_test.go‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package types
16+
17+
import (
18+
"math"
19+
"strings"
20+
"testing"
21+
22+
"github.com/google/cel-go/common/types/ref"
23+
"github.com/google/cel-go/common/types/traits"
24+
)
25+
26+
func TestMemoryTrackerTrack(t *testing.T) {
27+
adapter := DefaultTypeAdapter
28+
29+
t.Run("single_value_watermark", func(t *testing.T) {
30+
tracker := NewMemoryTracker()
31+
list := NewRefValList(adapter, []ref.Val{Int(1), Int(2)})
32+
if got := tracker.Track(list); got != 3 {
33+
t.Errorf("Track() got %d, want 3", got)
34+
}
35+
if got := tracker.Peak(); got != 3 {
36+
t.Errorf("Peak() got %d, want 3", got)
37+
}
38+
})
39+
40+
t.Run("call_inputs_combined_watermark", func(t *testing.T) {
41+
// Models the inputs to a.join(', ') where the list and separator coexist.
42+
tracker := NewMemoryTracker()
43+
list := NewRefValList(adapter, []ref.Val{String("a"), String("b")})
44+
sep := String(", ")
45+
if got := tracker.Track(list, sep); got != 4 {
46+
t.Errorf("Track(list, sep) got %d, want 4 (3 list + 1 separator)", got)
47+
}
48+
if got := tracker.Peak(); got != 4 {
49+
t.Errorf("Peak() got %d, want 4", got)
50+
}
51+
})
52+
53+
t.Run("peak_retains_max", func(t *testing.T) {
54+
tracker := NewMemoryTracker()
55+
tracker.Track(NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3)}))
56+
tracker.Track(Int(1))
57+
if got := tracker.Peak(); got != 4 {
58+
t.Errorf("Peak() got %d, want 4", got)
59+
}
60+
})
61+
62+
t.Run("call_output_watermark", func(t *testing.T) {
63+
// Models the output of a + a, a string twice the size of its inputs.
64+
tracker := NewMemoryTracker()
65+
in := String(strings.Repeat("a", 50))
66+
tracker.Track(in, in)
67+
out := String(strings.Repeat("a", 100))
68+
tracker.Track(out)
69+
if got := tracker.Peak(); got != 10 {
70+
t.Errorf("Peak() got %d, want 10 (100-char output at 10 chars per unit)", got)
71+
}
72+
})
73+
74+
t.Run("saturating_sum", func(t *testing.T) {
75+
tracker := NewMemoryTracker()
76+
big := customSizerVal(math.MaxUint32)
77+
if got := tracker.Track(big, big); got != math.MaxUint32 {
78+
t.Errorf("Track() got %d, want MaxUint32", got)
79+
}
80+
if got := tracker.Peak(); got != math.MaxUint32 {
81+
t.Errorf("Peak() got %d, want MaxUint32", got)
82+
}
83+
if tracker.CalculationLimitExceeded() {
84+
t.Error("CalculationLimitExceeded() got true, want false for pure saturation")
85+
}
86+
})
87+
88+
t.Run("calculation_limit_exceeded", func(t *testing.T) {
89+
tracker := NewMemoryTracker(
90+
MemoryTrackerSizeCalculator(NewSizeCalculator(SizeCalculatorMaxTraversal(2))))
91+
list := NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3)})
92+
if got := tracker.Track(list); got != math.MaxUint32 {
93+
t.Errorf("Track() got %d, want MaxUint32", got)
94+
}
95+
if !tracker.CalculationLimitExceeded() {
96+
t.Error("CalculationLimitExceeded() got false, want true")
97+
}
98+
})
99+
}
100+
101+
func TestMemoryTrackerSample(t *testing.T) {
102+
adapter := DefaultTypeAdapter
103+
104+
t.Run("default_interval_samples_every_value", func(t *testing.T) {
105+
tracker := NewMemoryTracker()
106+
for i := 0; i < 3; i++ {
107+
if got := tracker.Sample(Int(i)); got != 1 {
108+
t.Errorf("Sample() got %d, want 1", got)
109+
}
110+
}
111+
if got := tracker.Peak(); got != 1 {
112+
t.Errorf("Peak() got %d, want 1", got)
113+
}
114+
})
115+
116+
t.Run("interval_skips_intermediate_samples", func(t *testing.T) {
117+
tracker := NewMemoryTracker(MemoryTrackerSampleInterval(3))
118+
list := NewRefValList(adapter, []ref.Val{Int(1), Int(2)})
119+
if got := tracker.Sample(list); got != 0 {
120+
t.Errorf("Sample() #1 got %d, want 0 (skipped)", got)
121+
}
122+
if got := tracker.Sample(list); got != 0 {
123+
t.Errorf("Sample() #2 got %d, want 0 (skipped)", got)
124+
}
125+
if got := tracker.Sample(list); got != 3 {
126+
t.Errorf("Sample() #3 got %d, want 3 (computed)", got)
127+
}
128+
if got := tracker.Peak(); got != 3 {
129+
t.Errorf("Peak() got %d, want 3", got)
130+
}
131+
})
132+
133+
t.Run("zero_interval_clamped_to_one", func(t *testing.T) {
134+
tracker := NewMemoryTracker(MemoryTrackerSampleInterval(0))
135+
if got := tracker.Sample(Int(1)); got != 1 {
136+
t.Errorf("Sample() got %d, want 1", got)
137+
}
138+
})
139+
}
140+
141+
func TestMemoryTrackerLimit(t *testing.T) {
142+
t.Run("no_limit", func(t *testing.T) {
143+
tracker := NewMemoryTracker()
144+
tracker.Track(customSizerVal(math.MaxUint32))
145+
if tracker.ExceedsLimit() {
146+
t.Error("ExceedsLimit() got true, want false when no limit configured")
147+
}
148+
})
149+
150+
t.Run("under_limit", func(t *testing.T) {
151+
tracker := NewMemoryTracker(MemoryTrackerLimit(10))
152+
tracker.Track(customSizerVal(10))
153+
if tracker.ExceedsLimit() {
154+
t.Error("ExceedsLimit() got true, want false at exactly the limit")
155+
}
156+
})
157+
158+
t.Run("over_limit", func(t *testing.T) {
159+
tracker := NewMemoryTracker(MemoryTrackerLimit(10))
160+
tracker.Track(customSizerVal(11))
161+
if !tracker.ExceedsLimit() {
162+
t.Error("ExceedsLimit() got false, want true")
163+
}
164+
})
165+
}
166+
167+
func TestMemoryTrackerVersion(t *testing.T) {
168+
if got := NewMemoryTracker().Version(); got != 0 {
169+
t.Errorf("Version() got %d, want 0", got)
170+
}
171+
}
172+
173+
// Interface conformance check: the tracker's calculator remains usable as an AggregateSizer
174+
// by visitor implementations.
175+
var _ traits.Sizer = customSizerVal(0)

‎common/types/size_calc.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -179,8 +179,8 @@ func (s *SizeCalculator) EstimateAggregateSize(val any) AggregateSizeEstimate {
179179
return AggregateSizeEstimate{Size: size, LimitExceeded: exceeded}
180180
}
181181

182-
// stringSize converts a byte length to an element count where stringUnitLength bytes count
183-
// as a single element, rounding up with a minimum size of 1.
182+
// stringSize converts a character or byte length to an element count where stringUnitLength
183+
// characters count as a single element, rounding up with a minimum size of 1.
184184
func (s *SizeCalculator) stringSize(length int) uint32 {
185185
if length <= 0 {
186186
return 1

‎common/types/size_calc_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ func TestCalculateSize(t *testing.T) {
4646
{
4747
name: "sizer_string",
4848
val: String("hello"),
49-
want: 1, // 5 bytes round up to a single 10-byte element unit
49+
want: 1, // 5 chars round up to a single 10-char element unit
5050
},
5151
{
5252
name: "sizer_bytes",

0 commit comments

Comments
 (0)