opt: fix race caused by shared table annotations

When we assign placeholders on a detached memo, we copy the metadata.
However, this does a shallow copy of table annotations; this is
problematic for the statistics annotation which is a mutable object.

The fix is to register a copy function for each type of table
annotation as part of `NewTableAnnID`.

Fixes cockroachdb#34904.

Release note (bug fix): Fixed a crash related to cached plans.

pkg/sql/opt/memo/logical_props_builder.go

@@ -26,8 +26,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util"
 )
 
-var fdAnnID = opt.NewTableAnnID()
-
 // logicalPropsBuilder is a helper class that consolidates the code that derives
 // a parent expression's logical properties from those of its children.
 //
@@ -1670,3 +1668,9 @@ func (h *joinPropsHelper) cardinality() props.Cardinality {
 		return innerJoinCard
 	}
 }
+
+var fdAnnID = opt.NewTableAnnID(func(from interface{}) interface{} {
+	// The *FuncDepSet we cache in the table annotation is immutable; it's safe to
+	// share it.
+	return from
+})

pkg/sql/opt/memo/statistics_builder.go

@@ -28,8 +28,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/json"
 )
 
-var statsAnnID = opt.NewTableAnnID()
-
 // statisticsBuilder is responsible for building the statistics that are
 // used by the coster to estimate the cost of expressions.
 //
@@ -2705,3 +2703,11 @@ func RequestColStat(evalCtx *tree.EvalContext, e RelExpr, cols opt.ColSet) {
 	sb.init(evalCtx, e.Memo().Metadata())
 	sb.colStat(cols, e)
 }
+
+var statsAnnID = opt.NewTableAnnID(func(from interface{}) interface{} {
+	// In principle, we could return nil. But in practice many of the same
+	// statistics would have to be re-derived.
+	var res props.Statistics
+	res.CopyFrom(from.(*props.Statistics))
+	return &res
+})

pkg/sql/opt/metadata.go

@@ -89,6 +89,8 @@ type Metadata struct {
 	// might resolve to the same object now but to different objects later; we
 	// want to verify the resolution of both names.
 	deps []mdDep
+
+	// NOTE! When adding fields here, update CopyFrom.
 }
 
 type mdDep struct {
@@ -133,12 +135,23 @@ func (md *Metadata) Init() {
 // CopyFrom initializes the metadata with a copy of the provided metadata.
 // This metadata can then be modified independent of the copied metadata.
 func (md *Metadata) CopyFrom(from *Metadata) {
-	if len(md.cols) != 0 || len(md.tables) != 0 || len(md.deps) != 0 {
+	if len(md.schemas) != 0 || len(md.cols) != 0 || len(md.tables) != 0 ||
+		len(md.sequences) != 0 || len(md.deps) != 0 {
 		panic("CopyFrom requires empty destination")
 	}
 	md.schemas = append(md.schemas, from.schemas...)
 	md.cols = append(md.cols, from.cols...)
-	md.tables = append(md.tables, from.tables...)
+
+	if cap(md.tables) >= len(from.tables) {
+		md.tables = md.tables[:len(from.tables)]
+	} else {
+		md.tables = make([]TableMeta, len(from.tables))
+	}
+	for i := range md.tables {
+		md.tables[i].CopyFrom(&from.tables[i])
+	}
+
+	md.sequences = append(md.sequences, from.sequences...)
 	md.deps = append(md.deps, from.deps...)
 }
 

pkg/sql/opt/norm/factory.go

@@ -169,7 +169,7 @@ func (f *Factory) CustomFuncs() *CustomFuncs {
 // CopyAndRelaceDefault, which constructs a copy of the source operator using
 // children returned by recursive calls to the replace callback. Note that if a
 // non-leaf replacement node is constructed, its inputs must be copied using
-// CopyAndReplaceDefault
+// CopyAndReplaceDefault.
 //
 // Sample usage:
 //

pkg/sql/opt/props/col_stats_map.go

@@ -216,6 +216,27 @@ func (m *ColStatsMap) RemoveIntersecting(cols opt.ColSet) {
 	}
 }
 
+// CopyFrom initializes m with a copy of the given ColStatsMap.
+func (m *ColStatsMap) CopyFrom(from *ColStatsMap) {
+	for i := 0; i < m.count && i < initialColStatsCap; i++ {
+		m.initial[i].CopyFrom(&from.initial[i])
+	}
+	if len(from.other) > 0 {
+		m.other = make([]ColumnStatistic, len(from.other))
+		for i := range from.other {
+			m.other[i].CopyFrom(&from.other[i])
+		}
+	}
+	if len(from.index) > 0 {
+		m.index = make(map[colStatKey]colStatVal, len(from.index))
+		for k, v := range from.index {
+			m.index[k] = v
+		}
+	}
+	m.count = from.count
+	m.unique = from.unique
+}
+
 // Clear empties the map of all column statistics.
 func (m *ColStatsMap) Clear() {
 	m.count = 0

pkg/sql/opt/props/statistics.go

@@ -73,6 +73,14 @@ func (s *Statistics) Init(relProps *Relational) (zeroCardinality bool) {
 	return false
 }
 
+// CopyFrom initializes s with the values from another object; s then can be
+// independently modified.
+func (s *Statistics) CopyFrom(from *Statistics) {
+	s.RowCount = from.RowCount
+	s.ColStats.CopyFrom(&from.ColStats)
+	s.Selectivity = from.Selectivity
+}
+
 // ApplySelectivity applies a given selectivity to the statistics. RowCount and
 // Selectivity are updated. Note that DistinctCounts are not updated, other than
 // limiting them to the new RowCount. See ColumnStatistic.ApplySelectivity for
@@ -144,6 +152,13 @@ type ColumnStatistic struct {
 	NullCount float64
 }
 
+// CopyFrom initializes c with the values from another column statistic.
+func (c *ColumnStatistic) CopyFrom(from *ColumnStatistic) {
+	c.Cols = c.Cols.Copy()
+	c.DistinctCount = from.DistinctCount
+	c.NullCount = from.NullCount
+}
+
 // ApplySelectivity updates the distinct count and null count according to a
 // given selectivity.
 func (c *ColumnStatistic) ApplySelectivity(selectivity, inputRows float64) {

pkg/sql/opt/table_meta.go

@@ -109,6 +109,14 @@ var tableAnnIDCount TableAnnID
 // table struct.
 const maxTableAnnIDCount = 2
 
+// CopyTableAnnFn is a function that must be provided for each TableAnnID.
+// It is used when the metadata is copied; it must return an instance of the
+// annotation that can be used in the new copy. For immutable types, this can be
+// the same instance. For mutable types, it should be a copy.
+type CopyTableAnnFn func(from interface{}) interface{}
+
+var copyTableAnnFns [maxTableAnnIDCount]CopyTableAnnFn
+
 // TableMeta stores information about one of the tables stored in the metadata.
 type TableMeta struct {
 	// MetaID is the identifier for this table that is unique within the query
@@ -155,6 +163,18 @@ func (tm *TableMeta) IndexKeyColumns(indexOrd int) ColSet {
 	return indexCols
 }
 
+// CopyFrom copies the given TableMeta into this TableMeta, replacing any
+// existing data.
+func (tm *TableMeta) CopyFrom(from *TableMeta) {
+	*tm = *from
+	// Make copies of the objects stored as table annotations.
+	for i, val := range tm.anns {
+		if val != nil {
+			tm.anns[i] = copyTableAnnFns[i](val)
+		}
+	}
+}
+
 // TableAnnotation returns the given annotation that is associated with the
 // given table. If the table has no such annotation, TableAnnotation returns
 // nil.
@@ -182,11 +202,12 @@ func (md *Metadata) SetTableAnnotation(tabID TableID, tabAnnID TableAnnID, ann i
 // variables).
 //
 // See the TableAnnID comment for more details and a usage example.
-func NewTableAnnID() TableAnnID {
+func NewTableAnnID(copyFn CopyTableAnnFn) TableAnnID {
 	if tableAnnIDCount == maxTableAnnIDCount {
 		panic("can't allocate table annotation id; increase maxTableAnnIDCount to allow")
 	}
-	cnt := tableAnnIDCount
+	id := tableAnnIDCount
+	copyTableAnnFns[id] = copyFn
 	tableAnnIDCount++
-	return cnt
+	return id
 }