Merge cockroachdb#139279 cockroachdb#140035

139279: sqlstats: fix race condition in ss_mem_iterator r=kyle-a-wong a=kyle-a-wong

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 cockroachdb#138224.

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

Fixes: cockroachdb#138224
Epic: none
Release note: none

140035: opt: reduce `opt.ColSet` allocs when exploring partial index scans r=mgartner a=mgartner

When exploring a partial index scan, we previously built a set of
columns held constant by partial index predicate, and then built another
set of those columns that are not composite key encoded. Now, instead of
building the second set, we remove columns from the first. This will
reduce allocations in the case where the constant column IDs are larger
than 128.

Epic: None
Release note: None


Co-authored-by: Kyle Wong <[email protected]>
Co-authored-by: Marcus Gartner <[email protected]>

pkg/sql/opt/xform/scan_index_iter.go

@@ -364,15 +364,14 @@ func (it *scanIndexIter) filtersImplyPredicate(
 // the given filters and of types that do not have composite encodings.
 func (it *scanIndexIter) extractConstNonCompositeColumns(f memo.FiltersExpr) opt.ColSet {
 	constCols := memo.ExtractConstColumns(it.e.ctx, f, it.evalCtx)
-	var constNonCompositeCols opt.ColSet
 	for col, ok := constCols.Next(0); ok; col, ok = constCols.Next(col + 1) {
 		ord := it.tabMeta.MetaID.ColumnOrdinal(col)
 		typ := it.tabMeta.Table.Column(ord).DatumType()
-		if !colinfo.CanHaveCompositeKeyEncoding(typ) {
-			constNonCompositeCols.Add(col)
+		if colinfo.CanHaveCompositeKeyEncoding(typ) {
+			constCols.Remove(col)
 		}
 	}
-	return constNonCompositeCols
+	return constCols
 }
 
 // buildConstProjectionsFromPredicate builds a ProjectionsExpr that projects

pkg/sql/sqlstats/ssmemstorage/ss_mem_iterator.go

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