Backport proto3 row caching for query consolidator to slack-19.0

Backport of vitessio#19872, which was already cherry-picked to
slack-22.0 in #846. When multiple queries are consolidated,
the original query now pre-computes and caches the proto3 row encoding so
that N waiters share the encoded rows instead of each redundantly calling
RowsToProto3.

Gated behind --consolidator-cache-proto3-rows (default false).

Key v19-specific adaptations:
- No InsertIDChanged field (v22-only)
- AppendResult uses src.InsertID != 0 (not InsertIDUpdated())
- Benchmarks use for/b.N loops (Go 1.22, no b.Loop())
- Sizegen format: no blank lines between functions

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 :)

Author: bgwines

go/flags/endtoend/vtcombo.txt

@@ -51,6 +51,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                                   Cache proto3-encoded rows in consolidated results to avoid redundant encoding by waiting queries.
       --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

@@ -85,6 +85,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                                   Cache proto3-encoded rows in consolidated results to avoid redundant encoding by waiting queries.
       --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,16 +94,24 @@ 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,
-		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"
@@ -320,3 +324,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 i := 0; i < b.N; i++ {
+				ResultToProto3(result)
+			}
+		})
+
+		b.Run(fmt.Sprintf("rows=%d/cached", numRows), func(b *testing.B) {
+			result.CacheProto3Rows()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				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 i := 0; i < b.N; i++ {
+				var wg sync.WaitGroup
+				wg.Add(waiters)
+				for j := 0; j < waiters; j++ {
+					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 i := 0; i < b.N; i++ {
+				var wg sync.WaitGroup
+				wg.Add(waiters)
+				for j := 0; j < waiters; j++ {
+					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 i := 0; i < b.N; i++ {
+						runtime.GC()
+						var before runtime.MemStats
+						runtime.ReadMemStats(&before)
+
+						var wg sync.WaitGroup
+						wg.Add(waiters)
+						for j := 0; j < waiters; j++ {
+							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

@@ -35,6 +35,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
@@ -66,6 +72,7 @@ type ResultStream 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 {
@@ -108,6 +115,7 @@ func (result *Result) Copy() *Result {
 			out.Rows = append(out.Rows, CopyRow(r))
 		}
 	}
+	// proto3Rows is intentionally not propagated: callers may modify Rows
 	return out
 }
 
@@ -120,6 +128,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)
 	}
 }
 
@@ -324,6 +341,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
 	if src.RowsAffected == 0 && len(src.Rows) == 0 && len(src.Fields) == 0 {
 		return
 	}