Backport vitessio#19872: cache proto3 row encoding to reduce query consolidation memory (#846)

* vttablet: cache proto3 row encoding to reduce query consolidation memory utilization (vitessio#19872)

Signed-off-by: Brett Wines <bwines@slack-corp.com>
Co-authored-by: Claude <svc-devxp-claude@slack-corp.com>

* fix: resolve CI failures from backport cherry-pick

Two issues from the cherry-pick:
1. sizegen formatting: remove blank lines between generated functions
   in cached_size.go to match sizegen's output format
2. test panic: set FakePendingResult.Consolidator in the waiter cap
   reject test to prevent nil pointer dereference in AddWaiterCounter

Co-Authored-By: Claude <svc-devxp-claude@slack-corp.com>

AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :)

* Fix TestQueryExecutorConsolidatorWaiterCapReject: set consolidator total waiter count

The test was setting per-query WaiterCount but not the consolidator's
TotalWaiterCount, which is what the actual code checks to decide whether
the cap is exceeded. Also corrected assertions to match the actual code
path (one AddWaiterCounter(-1) call, not two).

Co-Authored-By: Claude <svc-devxp-claude@slack-corp.com>

AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :)

---------

Signed-off-by: Brett Wines <bwines@slack-corp.com>
Co-authored-by: Claude <svc-devxp-claude@slack-corp.com>

Author: bgwines

go/flags/endtoend/vtcombo.txt

@@ -44,6 +44,7 @@ Flags:
       --config-path strings                                              Paths to search for config files in. (default [{{ .Workdir }}])
       --config-persistence-min-interval duration                         minimum interval between persisting dynamic config changes back to disk (if no change has occurred, nothing is done). (default 1s)
       --config-type string                                               Config file type (omit to infer config type from file extension).
+      --consolidator-cache-proto3-rows                                   If true, the consolidation leader pre-caches proto3-encoded rows so that waiters avoid redundant encoding work.
       --consolidator-query-waiter-cap int                                Configure the maximum number of clients allowed to wait on the consolidator.
       --consolidator-query-waiter-cap-method string                      Configure the method when consolidator waiter cap is exceeded. Options: fallthrough, reject. (default "fallthrough")
       --consolidator-stream-query-size int                               Configure the stream consolidator query size in bytes. Setting to 0 disables the stream consolidator. (default 2097152)

go/flags/endtoend/vttablet.txt

@@ -78,6 +78,7 @@ Flags:
       --config-path strings                                              Paths to search for config files in. (default [{{ .Workdir }}])
       --config-persistence-min-interval duration                         minimum interval between persisting dynamic config changes back to disk (if no change has occurred, nothing is done). (default 1s)
       --config-type string                                               Config file type (omit to infer config type from file extension).
+      --consolidator-cache-proto3-rows                                   If true, the consolidation leader pre-caches proto3-encoded rows so that waiters avoid redundant encoding work.
       --consolidator-query-waiter-cap int                                Configure the maximum number of clients allowed to wait on the consolidator.
       --consolidator-query-waiter-cap-method string                      Configure the method when consolidator waiter cap is exceeded. Options: fallthrough, reject. (default "fallthrough")
       --consolidator-stream-query-size int                               Configure the stream consolidator query size in bytes. Setting to 0 disables the stream consolidator. (default 2097152)

go/sqltypes/cached_size.go

@@ -94,17 +94,25 @@ func proto3ToRows(fields []*querypb.Field, rows []*querypb.Row) [][]Value {
 	return result
 }
 
-// ResultToProto3 converts Result to proto3.
 func ResultToProto3(qr *Result) *querypb.QueryResult {
 	if qr == nil {
 		return nil
 	}
+	// This read is susceptible to TOCTOU if proto3Rows is populated
+	// concurrently, but the worst case is a redundant RowsToProto3 call.
+	// In the consolidation path this can't happen today (the leader sets
+	// the cached value before releasing the RWMutex), but the fallback is
+	// harmless.
+	rows := qr.proto3Rows
+	if rows == nil {
+		rows = RowsToProto3(qr.Rows)
+	}
 	return &querypb.QueryResult{
 		Fields:              qr.Fields,
 		RowsAffected:        qr.RowsAffected,
 		InsertId:            qr.InsertID,
 		InsertIdChanged:     qr.InsertIDChanged,
-		Rows:                RowsToProto3(qr.Rows),
+		Rows:                rows,
 		Info:                qr.Info,
 		SessionStateChanges: qr.SessionStateChanges,
 	}

go/sqltypes/proto3.go

@@ -17,6 +17,10 @@ limitations under the License.
 package sqltypes
 
 import (
+	"fmt"
+	"runtime"
+	"strconv"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -330,3 +334,226 @@ func TestProto3ValuesEqual(t *testing.T) {
 		require.Equal(t, tc.expected, Proto3ValuesEqual(tc.v1, tc.v2))
 	}
 }
+
+func TestResultToProto3_CachedRows(t *testing.T) {
+	fields := []*querypb.Field{{
+		Name: "col1",
+		Type: VarChar,
+	}, {
+		Name: "col2",
+		Type: Int64,
+	}}
+	result := &Result{
+		Fields:       fields,
+		RowsAffected: 2,
+		Rows: [][]Value{{
+			TestValue(VarChar, "aa"),
+			TestValue(Int64, "1"),
+		}, {
+			TestValue(VarChar, "bb"),
+			TestValue(Int64, "2"),
+		}},
+	}
+
+	uncached := ResultToProto3(result)
+	require.Len(t, uncached.Rows, 2)
+
+	result.CacheProto3Rows()
+
+	cached := ResultToProto3(result)
+	require.True(t, proto.Equal(uncached, cached), "cached and uncached proto3 results differ")
+
+	require.Same(t, cached.Rows[0], result.proto3Rows[0])
+	require.Same(t, cached.Rows[1], result.proto3Rows[1])
+}
+
+func TestResultToProto3_NilAndEmptyCache(t *testing.T) {
+	require.Nil(t, ResultToProto3(nil))
+	var nilResult *Result
+	nilResult.CacheProto3Rows() // does not panic
+
+	result := &Result{
+		Fields: []*querypb.Field{{Name: "col1", Type: VarChar}},
+	}
+	result.CacheProto3Rows()
+	require.Nil(t, result.proto3Rows)
+
+	p3 := ResultToProto3(result)
+	require.NotNil(t, p3)
+	require.Empty(t, p3.Rows)
+}
+
+func TestCopy_DoesNotPropagateProto3RowCache(t *testing.T) {
+	result := &Result{
+		Fields: []*querypb.Field{{Name: "col1", Type: Int64}},
+		Rows: [][]Value{{
+			TestValue(Int64, "42"),
+		}},
+	}
+	result.CacheProto3Rows()
+	require.NotNil(t, result.proto3Rows)
+
+	require.Nil(t, result.ShallowCopy().proto3Rows)
+	require.Nil(t, result.Copy().proto3Rows)
+}
+
+func TestMutations_InvalidateCachedProto3Rows(t *testing.T) {
+	resultWithProto3FieldPopulated := func(t *testing.T) *Result {
+		t.Helper()
+		result := &Result{
+			Fields: []*querypb.Field{{Name: "col1", Type: VarChar}},
+			Rows: [][]Value{{
+				TestValue(VarChar, "hello"),
+			}},
+		}
+		result.CacheProto3Rows()
+		require.NotNil(t, result.proto3Rows)
+		return result
+	}
+
+	t.Run("AppendResult", func(t *testing.T) {
+		result := resultWithProto3FieldPopulated(t)
+		result.AppendResult(&Result{
+			Rows: [][]Value{{
+				TestValue(VarChar, "world"),
+			}},
+		})
+		require.Nil(t, result.proto3Rows, "AppendResult must invalidate proto3Rows cache")
+
+		p3 := ResultToProto3(result)
+		require.Len(t, p3.Rows, 2)
+	})
+
+	t.Run("Repair", func(t *testing.T) {
+		result := resultWithProto3FieldPopulated(t)
+		result.Repair([]*querypb.Field{{Name: "col1", Type: VarBinary}})
+		require.Nil(t, result.proto3Rows, "Repair must invalidate proto3Rows cache")
+
+		p3 := ResultToProto3(result)
+		require.Len(t, p3.Rows, 1)
+	})
+}
+
+// makeTestResult builds a Result with numRows rows, each containing 5 columns
+// of mixed types. Row values are deterministic so benchmarks are reproducible.
+func makeTestResult(numRows int) *Result {
+	fields := MakeTestFields(
+		"col1|col2|col3|col4|col5",
+		"int64|varchar|varchar|float64|int64",
+	)
+	rows := make([][]Value, numRows)
+	for i := range rows {
+		rows[i] = []Value{
+			TestValue(Int64, strconv.Itoa(i)),
+			TestValue(VarChar, fmt.Sprintf("val-%06d", i)),
+			TestValue(VarChar, fmt.Sprintf("val-%06d-longervalue", i)),
+			TestValue(Float64, fmt.Sprintf("%d.%02d", i/100, i%100)),
+			TestValue(Int64, strconv.Itoa(i%2)),
+		}
+	}
+	return &Result{
+		Fields:       fields,
+		RowsAffected: uint64(numRows),
+		Rows:         rows,
+	}
+}
+
+// BenchmarkResultToProto3 measures per-call allocation cost of ResultToProto3
+// with and without CacheProto3Rows, across different result sizes.
+func BenchmarkResultToProto3(b *testing.B) {
+	for _, numRows := range []int{100, 1000, 10000} {
+		result := makeTestResult(numRows)
+
+		b.Run(fmt.Sprintf("rows=%d/uncached", numRows), func(b *testing.B) {
+			b.ReportAllocs()
+			for b.Loop() {
+				ResultToProto3(result)
+			}
+		})
+
+		b.Run(fmt.Sprintf("rows=%d/cached", numRows), func(b *testing.B) {
+			result.CacheProto3Rows()
+			b.ReportAllocs()
+			for b.Loop() {
+				ResultToProto3(result)
+			}
+		})
+	}
+}
+
+// BenchmarkConsolidationFanOut simulates the consolidator scenario: one shared
+// Result is read by N concurrent goroutines calling ResultToProto3. Reports
+// total heap bytes allocated across all waiters.
+func BenchmarkConsolidationFanOut(b *testing.B) {
+	const numRows = 10000
+
+	for _, waiters := range []int{1, 10, 50} {
+		result := makeTestResult(numRows)
+
+		b.Run(fmt.Sprintf("waiters=%d/uncached", waiters), func(b *testing.B) {
+			b.ReportAllocs()
+			for b.Loop() {
+				var wg sync.WaitGroup
+				wg.Add(waiters)
+				for range waiters {
+					go func() {
+						defer wg.Done()
+						ResultToProto3(result)
+					}()
+				}
+				wg.Wait()
+			}
+		})
+
+		b.Run(fmt.Sprintf("waiters=%d/cached", waiters), func(b *testing.B) {
+			result.CacheProto3Rows()
+			b.ReportAllocs()
+			for b.Loop() {
+				var wg sync.WaitGroup
+				wg.Add(waiters)
+				for range waiters {
+					go func() {
+						defer wg.Done()
+						ResultToProto3(result)
+					}()
+				}
+				wg.Wait()
+			}
+		})
+
+		// Report aggregate heap delta for a single iteration so the savings
+		// are visible in absolute terms, not just per-op.
+		b.Run(fmt.Sprintf("waiters=%d/heap-delta", waiters), func(b *testing.B) {
+			for _, cached := range []bool{false, true} {
+				label := "uncached"
+				r := makeTestResult(numRows)
+				if cached {
+					label = "cached"
+					r.CacheProto3Rows()
+				}
+				b.Run(label, func(b *testing.B) {
+					b.ReportAllocs()
+					for b.Loop() {
+						runtime.GC()
+						var before runtime.MemStats
+						runtime.ReadMemStats(&before)
+
+						var wg sync.WaitGroup
+						wg.Add(waiters)
+						for range waiters {
+							go func() {
+								defer wg.Done()
+								ResultToProto3(r)
+							}()
+						}
+						wg.Wait()
+
+						var after runtime.MemStats
+						runtime.ReadMemStats(&after)
+						b.ReportMetric(float64(after.TotalAlloc-before.TotalAlloc), "heap-bytes/op")
+					}
+				})
+			}
+		})
+	}
+}

go/sqltypes/proto3_test.go

@@ -36,6 +36,12 @@ type Result struct {
 	SessionStateChanges string           `json:"session_state_changes"`
 	StatusFlags         uint16           `json:"status_flags"`
 	Info                string           `json:"info"`
+
+	// proto3Rows caches the proto3-encoded representation of Rows, avoiding
+	// redundant encoding when multiple consumers share the same Result (i.e.
+	// query consolidation). Not populated for the non-consolidation flow;
+	// don't depend on it being present.
+	proto3Rows []*querypb.Row
 }
 
 //goland:noinspection GoUnusedConst
@@ -76,6 +82,7 @@ type MultiResultStream interface {
 // Repair fixes the type info in the rows
 // to conform to the supplied field types.
 func (result *Result) Repair(fields []*querypb.Field) {
+	result.proto3Rows = nil
 	// Usage of j is intentional.
 	for j, f := range fields {
 		for _, r := range result.Rows {
@@ -119,6 +126,7 @@ func (result *Result) Copy() *Result {
 			out.Rows = append(out.Rows, CopyRow(r))
 		}
 	}
+	// proto3Rows is intentionally not propagated: callers may modify Rows
 	return out
 }
 
@@ -132,6 +140,15 @@ func (result *Result) ShallowCopy() *Result {
 		Info:                result.Info,
 		SessionStateChanges: result.SessionStateChanges,
 		Rows:                result.Rows,
+		// proto3Rows is intentionally not propagated: callers may modify Rows
+	}
+}
+
+func (result *Result) CacheProto3Rows() {
+	// result can be nil when the query errors out (e.g. execDBConn returns
+	// a nil Result alongside an error).
+	if result != nil && len(result.Rows) > 0 {
+		result.proto3Rows = RowsToProto3(result.Rows)
 	}
 }
 
@@ -339,6 +356,7 @@ func (result *Result) StripMetadata(incl querypb.ExecuteOptions_IncludedFields)
 // to another result.Note currently it doesn't handle cases like
 // if two results have different fields.We will enhance this function.
 func (result *Result) AppendResult(src *Result) {
+	result.proto3Rows = nil
 	result.RowsAffected += src.RowsAffected
 	if src.InsertIDUpdated() {
 		result.InsertID = src.InsertID