sql: add and adopt fast memo staleness checks

Previously, determining if a memo was stale involved re-resolving
descriptor references to ensure that the versions and descriptor IDs
were the same. This could be fairly expensive and hit various points of
contention. To address this, this patch takes advantage of the
lease manager and stats cache generation to determine if descriptors
have and statistics have to be re-resolved. It does so by introducing a
fingerprint that will contain this generation information, current
database and search_path. If none of these things change, then we know
that a memo can be re-used without any extra checks.

Release note: None
Fixes: #105867

pkg/bench/rttanalysis/testdata/benchmark_expectations

@@ -27,7 +27,7 @@ exp,benchmark
 8,AlterTableUnsplit/alter_table_unsplit_at_1_value
 11,AlterTableUnsplit/alter_table_unsplit_at_2_values
 14,AlterTableUnsplit/alter_table_unsplit_at_3_values
-5,Audit/select_from_an_audit_table
+2-4,Audit/select_from_an_audit_table
 26,CreateRole/create_role_with_1_option
 29,CreateRole/create_role_with_2_options
 32,CreateRole/create_role_with_3_options
@@ -52,7 +52,7 @@ exp,benchmark
 17,DropView/drop_2_views
 17,DropView/drop_3_views
 5,GenerateObjects/generate_1000_tables_-_this_test_should_use_the_same_number_of_RTTs_as_for_10_tables
-11,GenerateObjects/generate_100_tables_from_template
+8-10,GenerateObjects/generate_100_tables_from_template
 5,GenerateObjects/generate_10_tables
 16,GenerateObjects/generate_10x10_schemas_and_tables_in_existing_db
 5,GenerateObjects/generate_50000_tables
@@ -125,9 +125,9 @@ exp,benchmark
 20,Truncate/truncate_2_column_0_rows
 20,Truncate/truncate_2_column_1_rows
 20,Truncate/truncate_2_column_2_rows
-3,UDFResolution/select_from_udf
-5,UseManyRoles/use_2_roles
-3,UseManyRoles/use_50_roles
+0,UDFResolution/select_from_udf
+2,UseManyRoles/use_2_roles
+0,UseManyRoles/use_50_roles
 1,VirtualTableQueries/select_crdb_internal.invalid_objects_with_1_fk
 1,VirtualTableQueries/select_crdb_internal.tables_with_1_fk
 5,VirtualTableQueries/virtual_table_cache_with_point_lookups

pkg/sql/opt/cat/catalog.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/roleoption"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/catid"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/lib/pq/oid"
 )
@@ -77,6 +78,24 @@ type Flags struct {
 	IncludeNonActiveIndexes bool
 }
 
+// DataSourcesFingerprint can be stored and confirm that all data sources resolved,
+// at the current point in time will have the same version, stats, and name
+// resolution.
+type DataSourcesFingerprint struct {
+	// leaseGeneration tracks if any new version of descriptors has been published
+	// by the lease manager.
+	LeaseGeneration int64
+
+	// statsGeneration tracks if any new statistics have been published.
+	StatsGeneration int64
+
+	// currentDatabase tracks the currently set database
+	CurrentDatabase string
+
+	// searchPath tracks the current search path
+	SearchPath sessiondata.SearchPath
+}
+
 // Catalog is an interface to a database catalog, exposing only the information
 // needed by the query optimizer.
 //
@@ -210,6 +229,13 @@ type Catalog interface {
 	// GetCurrentUser returns the username.SQLUsername of the current session.
 	GetCurrentUser() username.SQLUsername
 
+	// CompareDataSourcesFingerprint compares a data source fingerprint to see if
+	// any data sources could have changed. A new fingerprint is returned if they
+	// have changed
+	CompareDataSourcesFingerprint(
+		f *DataSourcesFingerprint, new *DataSourcesFingerprint,
+	) bool
+
 	// GetRoutineOwner returns the username.SQLUsername of the routine's
 	// (specified by routineOid) owner.
 	GetRoutineOwner(ctx context.Context, routineOid oid.Oid) (username.SQLUsername, error)

pkg/sql/opt/memo/BUILD.bazel

@@ -46,6 +46,7 @@ go_library(
         "//pkg/sql/sem/tree",
         "//pkg/sql/sem/tree/treewindow",
         "//pkg/sql/sem/volatility",
+        "//pkg/sql/sessiondata",
         "//pkg/sql/stats",
         "//pkg/sql/types",
         "//pkg/util/buildutil",

pkg/sql/opt/memo/memo.go

@@ -396,7 +396,10 @@ func (m *Memo) HasPlaceholders() bool {
 // perform KV operations on behalf of the transaction associated with the
 // provided catalog, and those errors are required to be propagated.
 func (m *Memo) IsStale(
-	ctx context.Context, evalCtx *eval.Context, catalog cat.Catalog,
+	ctx context.Context,
+	evalCtx *eval.Context,
+	catalog cat.Catalog,
+	dataSourcesFingerprint *cat.DataSourcesFingerprint,
 ) (bool, error) {
 	// Memo is stale if fields from SessionData that can affect planning have
 	// changed.
@@ -457,6 +460,17 @@ func (m *Memo) IsStale(
 		m.txnIsoLevel != evalCtx.TxnIsoLevel {
 		return true, nil
 	}
+
+	// If the query is AOST we must check all the dependencies, since the descriptors
+	// may have been different in the past. Otherwise, the data sources fingerprint
+	// is sufficient.
+	var newFP cat.DataSourcesFingerprint
+	if evalCtx.AsOfSystemTime == nil &&
+		dataSourcesFingerprint != nil &&
+		catalog.CompareDataSourcesFingerprint(dataSourcesFingerprint, &newFP) {
+		return false, nil
+	}
+
 	// Memo is stale if the fingerprint of any object in the memo's metadata has
 	// changed, or if the current user no longer has sufficient privilege to
 	// access the object.
@@ -465,6 +479,10 @@ func (m *Memo) IsStale(
 	} else if !depsUpToDate {
 		return true, nil
 	}
+
+	if evalCtx.AsOfSystemTime == nil && dataSourcesFingerprint != nil {
+		*dataSourcesFingerprint = newFP
+	}
 	return false, nil
 }
 

pkg/sql/opt/memo/memo_test.go

@@ -180,7 +180,7 @@ func TestMemoIsStale(t *testing.T) {
 	ctx := context.Background()
 	stale := func() {
 		t.Helper()
-		if isStale, err := o.Memo().IsStale(ctx, &evalCtx, catalog); err != nil {
+		if isStale, err := o.Memo().IsStale(ctx, &evalCtx, catalog, nil); err != nil {
 			t.Fatal(err)
 		} else if !isStale {
 			t.Errorf("memo should be stale")
@@ -192,7 +192,7 @@ func TestMemoIsStale(t *testing.T) {
 		var o2 xform.Optimizer
 		opttestutils.BuildQuery(t, &o2, catalog, &evalCtx, query)
 
-		if isStale, err := o2.Memo().IsStale(ctx, &evalCtx, catalog); err != nil {
+		if isStale, err := o2.Memo().IsStale(ctx, &evalCtx, catalog, nil); err != nil {
 			t.Fatal(err)
 		} else if isStale {
 			t.Errorf("memo should not be stale")
@@ -201,7 +201,7 @@ func TestMemoIsStale(t *testing.T) {
 
 	notStale := func() {
 		t.Helper()
-		if isStale, err := o.Memo().IsStale(ctx, &evalCtx, catalog); err != nil {
+		if isStale, err := o.Memo().IsStale(ctx, &evalCtx, catalog, nil); err != nil {
 			t.Fatal(err)
 		} else if isStale {
 			t.Errorf("memo should not be stale")
@@ -529,7 +529,7 @@ func TestMemoIsStale(t *testing.T) {
 
 	// User no longer has access to view.
 	catalog.View(tree.NewTableNameWithSchema("t", catconstants.PublicSchemaName, "abcview")).Revoked = true
-	_, err = o.Memo().IsStale(ctx, &evalCtx, catalog)
+	_, err = o.Memo().IsStale(ctx, &evalCtx, catalog, nil)
 	if exp := "user does not have privilege"; !testutils.IsError(err, exp) {
 		t.Fatalf("expected %q error, but got %+v", exp, err)
 	}
@@ -538,7 +538,7 @@ func TestMemoIsStale(t *testing.T) {
 
 	// User no longer has execution privilege on a UDF.
 	catalog.RevokeExecution(catalog.Function("one").Oid)
-	_, err = o.Memo().IsStale(ctx, &evalCtx, catalog)
+	_, err = o.Memo().IsStale(ctx, &evalCtx, catalog, nil)
 	if exp := "user does not have privilege to execute function"; !testutils.IsError(err, exp) {
 		t.Fatalf("expected %q error, but got %+v", exp, err)
 	}

pkg/sql/opt/testutils/testcat/test_catalog.go

@@ -51,8 +51,9 @@ type Catalog struct {
 	counter    int
 	enumTypes  map[string]*types.T
 
-	udfs           map[string]*tree.ResolvedFunctionDefinition
-	revokedUDFOids intsets.Fast
+	udfs                  map[string]*tree.ResolvedFunctionDefinition
+	revokedUDFOids        intsets.Fast
+	uniqueLeaseGeneration int64
 }
 
 type dataSource interface {
@@ -453,6 +454,14 @@ func (tc *Catalog) AddSequence(seq *Sequence) {
 	tc.testSchema.dataSources[fq] = seq
 }
 
+// CompareDataSourcesFingerprint always assume that the fingerprints are changing
+// on us.
+func (tc *Catalog) CompareDataSourcesFingerprint(
+	a *cat.DataSourcesFingerprint, b *cat.DataSourcesFingerprint,
+) bool {
+	return false
+}
+
 // ExecuteMultipleDDL parses the given semicolon-separated DDL SQL statements
 // and applies each of them to the test catalog.
 func (tc *Catalog) ExecuteMultipleDDL(sql string) error {

pkg/sql/opt_catalog.go

@@ -508,6 +508,31 @@ func (oc *optCatalog) GetCurrentUser() username.SQLUsername {
 	return oc.planner.User()
 }
 
+// CompareDataSourcesFingerprint is part of the cat.Catalog interface.
+func (oc *optCatalog) CompareDataSourcesFingerprint(
+	f *cat.DataSourcesFingerprint, new *cat.DataSourcesFingerprint,
+) (matches bool) {
+	// Ensure that neither the lease manager or stats collection has seen
+	// any updates, which will allow us to skip re-resolution. Confirm that
+	// the current databse and search path are the same as the previous usage
+	// of this memo.
+	lg := oc.planner.Descriptors().GetLeaseGeneration()
+	sg := oc.planner.execCfg.TableStatsCache.GetGeneration()
+	if f.LeaseGeneration == lg &&
+		f.StatsGeneration == sg &&
+		f.CurrentDatabase == oc.planner.CurrentDatabase() &&
+		f.SearchPath.Equals(&oc.planner.SessionData().SearchPath) {
+		return true
+	}
+	*new = cat.DataSourcesFingerprint{
+		LeaseGeneration: lg,
+		StatsGeneration: sg,
+		CurrentDatabase: oc.planner.CurrentDatabase(),
+		SearchPath:      oc.planner.SessionData().SearchPath,
+	}
+	return false
+}
+
 // GetRoutineOwner is part of the cat.Catalog interface.
 func (oc *optCatalog) GetRoutineOwner(
 	ctx context.Context, routineOid oid.Oid,

pkg/sql/plan_opt.go

@@ -139,7 +139,7 @@ func (p *planner) prepareUsingOptimizer(
 			if !pm.TypeHints.Identical(p.semaCtx.Placeholders.TypeHints) {
 				opc.log(ctx, "query cache hit but type hints don't match")
 			} else {
-				isStale, err := cachedData.Memo.IsStale(ctx, p.EvalContext(), opc.catalog)
+				isStale, err := cachedData.Memo.IsStale(ctx, p.EvalContext(), opc.catalog, nil)
 				if err != nil {
 					return 0, err
 				}
@@ -652,7 +652,7 @@ func (opc *optPlanningCtx) chooseValidPreparedMemo(ctx context.Context) (*memo.M
 			// A generic plan does not yet exist.
 			return nil, nil
 		}
-		isStale, err := prep.GenericMemo.IsStale(ctx, opc.p.EvalContext(), opc.catalog)
+		isStale, err := prep.GenericMemo.IsStale(ctx, opc.p.EvalContext(), opc.catalog, &prep.GenericFingerPrint)
 		if err != nil {
 			return nil, err
 		} else if !isStale {
@@ -666,7 +666,7 @@ func (opc *optPlanningCtx) chooseValidPreparedMemo(ctx context.Context) (*memo.M
 	}
 
 	if prep.BaseMemo != nil {
-		isStale, err := prep.BaseMemo.IsStale(ctx, opc.p.EvalContext(), opc.catalog)
+		isStale, err := prep.BaseMemo.IsStale(ctx, opc.p.EvalContext(), opc.catalog, &prep.BaseFingerPrint)
 		if err != nil {
 			return nil, err
 		} else if !isStale {
@@ -791,7 +791,7 @@ func (opc *optPlanningCtx) buildExecMemo(ctx context.Context) (_ *memo.Memo, _ e
 		// Consult the query cache.
 		cachedData, ok := p.execCfg.QueryCache.Find(&p.queryCacheSession, opc.p.stmt.SQL)
 		if ok {
-			if isStale, err := cachedData.Memo.IsStale(ctx, p.EvalContext(), opc.catalog); err != nil {
+			if isStale, err := cachedData.Memo.IsStale(ctx, p.EvalContext(), opc.catalog, nil); err != nil {
 				return nil, err
 			} else if isStale {
 				opc.log(ctx, "query cache hit but needed update")

pkg/sql/prepared_stmt.go

@@ -14,6 +14,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/appstatspb"
 	"github.com/cockroachdb/cockroach/pkg/sql/clusterunique"
 	"github.com/cockroachdb/cockroach/pkg/sql/flowinfra"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgwirebase"
 	"github.com/cockroachdb/cockroach/pkg/sql/querycache"
@@ -54,10 +55,14 @@ type PreparedStatement struct {
 	// optimizer during prepare of a SQL statement.
 	BaseMemo *memo.Memo
 
+	BaseFingerPrint cat.DataSourcesFingerprint
+
 	// GenericMemo, if present, is a fully-optimized memo that can be executed
 	// as-is.
 	GenericMemo *memo.Memo
 
+	GenericFingerPrint cat.DataSourcesFingerprint
+
 	// IdealGenericPlan is true if GenericMemo is guaranteed to be optimal
 	// across all executions of the prepared statement. Ideal generic plans are
 	// generated when the statement has no placeholders nor fold-able stable