sqlstats: fix race condition in ss_mem_iterator

StmtStatsIterator.Next() copies the next statement stats info into a
new CollectedStatementStatistics object, but the slice fields
of stmtStats.mu.data are not copied explicitly. This makes them
susecptible to race conditions, as seen in #138224.

To fix, all slice fields in the StatementStatistics are
explicitly copied, using a new util.CopySlice function.

Fixes: 138224
Epic: none
Release note: none

pkg/sql/sqlstats/ssmemstorage/ss_mem_iterator.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/appstatspb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
+	"github.com/cockroachdb/cockroach/pkg/util"
 )
 
 type baseIterator struct {
@@ -77,7 +78,7 @@ func (s *StmtStatsIterator) Next() bool {
 	}
 
 	statementStats.mu.Lock()
-	data := statementStats.mu.data
+	data := copyData(statementStats)
 	distSQLUsed := statementStats.mu.distSQLUsed
 	vectorized := statementStats.mu.vectorized
 	fullScan := statementStats.mu.fullScan
@@ -105,6 +106,20 @@ func (s *StmtStatsIterator) Next() bool {
 	return true
 }
 
+// copyData Copies the statement stat's data object under the mutex.
+// This function requires that there is a lock on the stats object.
+func copyData(stats *stmtStats) appstatspb.StatementStatistics {
+	stats.mu.AssertHeld()
+	data := stats.mu.data
+	data.Nodes = util.CopySlice(data.Nodes)
+	data.KVNodeIDs = util.CopySlice(data.KVNodeIDs)
+	data.Regions = util.CopySlice(data.Regions)
+	data.PlanGists = util.CopySlice(data.PlanGists)
+	data.IndexRecommendations = util.CopySlice(data.IndexRecommendations)
+	data.Indexes = util.CopySlice(data.Indexes)
+	return data
+}
+
 // Cur returns the appstatspb.CollectedStatementStatistics at the current internal
 // counter.
 func (s *StmtStatsIterator) Cur() *appstatspb.CollectedStatementStatistics {

pkg/util/slices.go

@@ -111,3 +111,11 @@ func InsertUnique[T cmp.Ordered](s []T, v T) []T {
 	}
 	return slices.Insert(s, idx, v)
 }
+
+// CopySlice creates a new slice and copies all elements from the provided
+// slice into it
+func CopySlice[T any, S []T](s S) S {
+	newSlice := make([]T, len(s))
+	copy(newSlice, s)
+	return newSlice
+}

pkg/util/slices_test.go

@@ -142,3 +142,25 @@ func TestMapFrom(t *testing.T) {
 		return i, i * i
 	}))
 }
+
+func TestCopySlice(t *testing.T) {
+	expected1 := []int{1, 2, 3}
+	expected2 := []string{"a", "b", "c"}
+	expected3 := []struct {
+		a int
+		b string
+	}{{1, "a"}, {2, "b"}}
+	actual1 := CopySlice(expected1)
+	actual2 := CopySlice(expected2)
+	actual3 := CopySlice(expected3)
+	require.Equal(t, expected1, actual1)
+	require.Equal(t, expected2, actual2)
+	require.Equal(t, expected3, actual3)
+
+	expected1[0] = 999
+	expected2[0] = "999"
+	expected3[0].a = 999
+	require.NotEqual(t, expected1, actual1)
+	require.NotEqual(t, expected2, actual2)
+	require.NotEqual(t, expected3, actual3)
+}