Fix the heap escape on cost-tracking with a fixed sized, static buffer (#1600)

Author: TristonianJones

common/cost/tracker.go

@@ -15,8 +15,9 @@
 package cost
 
 import (
-	"cel.dev/cel-go/common/types/ref"
 	"fmt"
+
+	"cel.dev/cel-go/common/types/ref"
 )
 
 // WARNING: Any changes to cost calculations in this file require a corresponding change in estimator.go
@@ -129,6 +130,7 @@ type Tracker struct {
 	cost                     uint64
 	presenceTestHasCost      bool
 	modelVersion             uint32
+	argsBuf                  [2]ref.Val
 }
 
 // NewTracker creates a new Tracker with a given estimator and a set of functional TrackerOption values.
@@ -231,17 +233,19 @@ func (c *Tracker) EvalZeroArity(vars any, id int64, call Call, result ref.Val) {
 
 // EvalUnary records the cost for a unary call expression.
 func (c *Tracker) EvalUnary(vars any, id int64, call Call, arg ref.Val, result ref.Val) {
-	var buf [1]ref.Val
-	buf[0] = arg
-	c.recordCallCost(call, buf[:], result)
+	c.argsBuf[0] = arg
+	c.argsBuf[1] = nil
+	c.recordCallCost(call, c.argsBuf[:1], result)
+	c.argsBuf[0] = nil
 }
 
 // EvalBinary records the cost for a binary call expression.
 func (c *Tracker) EvalBinary(vars any, id int64, call Call, lhs, rhs ref.Val, result ref.Val) {
-	var buf [2]ref.Val
-	buf[0] = lhs
-	buf[1] = rhs
-	c.recordCallCost(call, buf[:], result)
+	c.argsBuf[0] = lhs
+	c.argsBuf[1] = rhs
+	c.recordCallCost(call, c.argsBuf[:2], result)
+	c.argsBuf[0] = nil
+	c.argsBuf[1] = nil
 }
 
 // EvalVarArgs records the cost for a variadic call expression.

common/cost/tracker_test.go

@@ -404,3 +404,54 @@ func TestTracker_ConcurrentCloneRace(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+type zeroAllocEstimator struct {
+	cost uint64
+}
+
+func (e *zeroAllocEstimator) CallCost(function, overloadID string, args []ref.Val, result ref.Val) *uint64 {
+	return &e.cost
+}
+
+func TestTracker_EvalUnaryBinaryAllocs(t *testing.T) {
+	estimator := &zeroAllocEstimator{cost: 3}
+	tracker, err := cost.NewTracker(estimator)
+	if err != nil {
+		t.Fatalf("NewTracker() failed: %v", err)
+	}
+	unmodeledTracker, err := cost.NewTracker(nil)
+	if err != nil {
+		t.Fatalf("NewTracker(nil) failed: %v", err)
+	}
+
+	unaryCall := &testCall{function: "!", overloadID: overloads.LogicalNot}
+	binaryCall := &testCall{function: "_+_", overloadID: overloads.AddInt64}
+	arg1 := types.Int(1)
+	arg2 := types.Int(2)
+	res := types.Int(3)
+
+	if allocs := testing.AllocsPerRun(100, func() {
+		unmodeledTracker.EvalUnary(nil, 1, unaryCall, types.True, types.False)
+	}); allocs != 0 {
+		t.Errorf("EvalUnary (unmodeled) allocated %v times, want 0", allocs)
+	}
+
+	if allocs := testing.AllocsPerRun(100, func() {
+		unmodeledTracker.EvalBinary(nil, 2, binaryCall, arg1, arg2, res)
+	}); allocs != 0 {
+		t.Errorf("EvalBinary (unmodeled) allocated %v times, want 0", allocs)
+	}
+
+	if allocs := testing.AllocsPerRun(100, func() {
+		tracker.EvalUnary(nil, 1, unaryCall, types.True, types.False)
+	}); allocs != 0 {
+		t.Errorf("EvalUnary (with estimator) allocated %v times, want 0", allocs)
+	}
+
+	if allocs := testing.AllocsPerRun(100, func() {
+		tracker.EvalBinary(nil, 2, binaryCall, arg1, arg2, res)
+	}); allocs != 0 {
+		t.Errorf("EvalBinary (with estimator) allocated %v times, want 0", allocs)
+	}
+}
+