sql/hints, builtins: add two more statement_hints_cache builtins

Add new builtins `crdb_internal.clear_statement_hints_cache` and
`crdb_internal.await_statement_hints_cache` which will be useful for
testing.

The first builtin clears the hintCache, similar to
`clear_query_plan_cache` and `clear_table_stats_cache`. (I debated
whether this builtin should also set hintedHashes to nil and restart the
changefeed, but that seems too different from
`clear_table_stats_cache`.)

The second builtin blocks until the hints cache rangefeed watcher has
caught up with the present, which usually takes a couple of
seconds. This is handy when we want to test a statement using the hint
immediately after adding the hint to system.statement_hints.

Informs: #153633

Release note: None

Author: michae2

pkg/sql/faketreeeval/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "//pkg/jobs/jobspb",
         "//pkg/roachpb",
         "//pkg/security/username",
+        "//pkg/settings/cluster",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/hintpb",
         "//pkg/sql/pgwire/pgcode",

pkg/sql/faketreeeval/evalctx.go

@@ -14,6 +14,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
@@ -578,6 +579,12 @@ func (ep *DummyEvalPlanner) ClearQueryPlanCache() {}
 // ClearTableStatsCache is part of the eval.Planner interface.
 func (ep *DummyEvalPlanner) ClearTableStatsCache() {}
 
+// ClearStatementHintsCache is part of the eval.Planner interface.
+func (ep *DummyEvalPlanner) ClearStatementHintsCache() {}
+
+// AwaitStatementHintsCache is part of the eval.Planner interface.
+func (ep *DummyEvalPlanner) AwaitStatementHintsCache(ctx context.Context, st *cluster.Settings) {}
+
 // RetryCounter is part of the eval.Planner interface.
 func (ep *DummyEvalPlanner) RetryCounter() int {
 	return 0

pkg/sql/hints/BUILD.bazel

@@ -14,6 +14,8 @@ go_library(
         "//pkg/kv/kvclient/rangefeed/rangefeedbuffer",
         "//pkg/kv/kvclient/rangefeed/rangefeedcache",
         "//pkg/kv/kvpb",
+        "//pkg/kv/kvserver",
+        "//pkg/kv/kvserver/closedts",
         "//pkg/roachpb",
         "//pkg/settings",
         "//pkg/settings/cluster",

pkg/sql/hints/hint_cache.go

@@ -17,6 +17,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/rangefeed/rangefeedbuffer"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/rangefeed/rangefeedcache"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -92,6 +94,10 @@ type StatementHintsCache struct {
 	// generation is incremented any time the hint cache is updated by the
 	// rangefeed.
 	generation atomic.Int64
+
+	// frontier is the walltime of the rangefeed watcher frontier time, which is
+	// polled by goroutines in Await.
+	frontier atomic.Int64
 }
 
 // cacheSize is the size of the entries to store in the cache.
@@ -139,6 +145,16 @@ func NewStatementHintsCache(
 	return hintsCache
 }
 
+// Clear removes all entries from the StatementHintsCache.
+func (c *StatementHintsCache) Clear() {
+	defer c.generation.Add(1)
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.mu.hintCache.Clear()
+	// TODO(michae2): Consider whether we should also delete hintedHashes and
+	// restart the rangefeed.
+}
+
 // ============================================================================
 // Populating the cache.
 // ============================================================================
@@ -214,6 +230,7 @@ func (c *StatementHintsCache) onUpdate(
 		log.Dev.Info(ctx, "statement_hints rangefeed applying incremental update")
 		c.handleIncrementalUpdate(ctx, update)
 	}
+	c.frontier.Store(update.Timestamp.WallTime)
 }
 
 // handleInitialScan builds the hintedHashes map and adds it to
@@ -251,6 +268,7 @@ func (c *StatementHintsCache) handleIncrementalUpdate(
 	defer c.generation.Add(1)
 	c.mu.Lock()
 	defer c.mu.Unlock()
+
 	for _, ev := range update.Events {
 		// Drop outdated entries from hintCache.
 		c.mu.hintCache.Del(ev.hash)
@@ -358,6 +376,30 @@ var _ rangefeedbuffer.Event = &bufferEvent{}
 // Reading from the cache.
 // ============================================================================
 
+// Await blocks until the StatementHintsCache's rangefeed watcher catches up
+// with the present. After Await returns, MaybeGetStatementHints should
+// accurately reflect all hints that were modified before the call to Await
+// (assuming the ctx was not canceled).
+func (c *StatementHintsCache) Await(ctx context.Context, st *cluster.Settings) {
+	// The frontier timestamp comes from the rangefeed, and could be up to
+	// kv.closed_timestamp.target_duration +
+	// kv.rangefeed.closed_timestamp_refresh_interval behind the present.
+	targetDuration := closedts.TargetDuration.Get(&st.SV)
+	refreshInterval := kvserver.RangeFeedRefreshInterval.Get(&st.SV)
+	const fudge = 10 * time.Millisecond
+	waitUntil := c.clock.Now().AddDuration(targetDuration + refreshInterval + fudge).WallTime
+
+	// Await is only used for testing, so we don't need to wake up immediately. We
+	// can get away with polling the frontier time.
+	for frontier := c.frontier.Load(); frontier < waitUntil; frontier = c.frontier.Load() {
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(time.Duration(waitUntil-frontier) * time.Nanosecond):
+		}
+	}
+}
+
 // GetGeneration returns the current generation, which will change if any
 // modifications happen to the cache.
 func (c *StatementHintsCache) GetGeneration() int64 {

pkg/sql/hints/hint_table.go

@@ -102,5 +102,8 @@ func InsertHintIntoDB(
 	if err != nil {
 		return 0, err
 	}
+	// TODO(michae2,drewk): Consider calling
+	// StatementHintsCache.handleIncrementalUpdate here to eagerly update the
+	// local node's cache.
 	return int64(tree.MustBeDInt(row[0])), nil
 }

pkg/sql/planner.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/spanconfig"
 	"github.com/cockroachdb/cockroach/pkg/sql/auditlogging"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
@@ -1090,6 +1091,20 @@ func (p *planner) ClearTableStatsCache() {
 	}
 }
 
+// ClearStatementHintsCache is part of the eval.Planner interface.
+func (p *planner) ClearStatementHintsCache() {
+	if p.execCfg.StatementHintsCache != nil {
+		p.execCfg.StatementHintsCache.Clear()
+	}
+}
+
+// AwaitStatementHintsCache is part of the eval.Planner interface.
+func (p *planner) AwaitStatementHintsCache(ctx context.Context, st *cluster.Settings) {
+	if p.execCfg.StatementHintsCache != nil {
+		p.execCfg.StatementHintsCache.Await(ctx, st)
+	}
+}
+
 // innerPlansMustUseLeafTxn returns true if inner plans must use a leaf
 // transaction.
 func (p *planner) innerPlansMustUseLeafTxn() bool {

pkg/sql/sem/builtins/builtins.go

@@ -9714,6 +9714,40 @@ WHERE object_id = table_descriptor_id
 			Volatility: volatility.Volatile,
 		},
 	),
+	"crdb_internal.clear_statement_hints_cache": makeBuiltin(
+		tree.FunctionProperties{
+			Category:         builtinconstants.CategorySystemRepair,
+			Undocumented:     true,
+			DistsqlBlocklist: true, // applicable only on the gateway
+		},
+		tree.Overload{
+			Types:      tree.ParamTypes{},
+			ReturnType: tree.FixedReturnType(types.Void),
+			Fn: func(ctx context.Context, evalCtx *eval.Context, args tree.Datums) (tree.Datum, error) {
+				evalCtx.Planner.ClearStatementHintsCache()
+				return tree.DVoidDatum, nil
+			},
+			Info:       `This function is used to clear the statement hints cache on the gateway node`,
+			Volatility: volatility.Volatile,
+		},
+	),
+	"crdb_internal.await_statement_hints_cache": makeBuiltin(
+		tree.FunctionProperties{
+			Category:         builtinconstants.CategorySystemRepair,
+			Undocumented:     true,
+			DistsqlBlocklist: true, // applicable only on the gateway
+		},
+		tree.Overload{
+			Types:      tree.ParamTypes{},
+			ReturnType: tree.FixedReturnType(types.Void),
+			Fn: func(ctx context.Context, evalCtx *eval.Context, args tree.Datums) (tree.Datum, error) {
+				evalCtx.Planner.AwaitStatementHintsCache(ctx, evalCtx.Settings)
+				return tree.DVoidDatum, nil
+			},
+			Info:       `This function is used to await the statement hints cache on the gateway node`,
+			Volatility: volatility.Volatile,
+		},
+	),
 }
 
 var lengthImpls = func(incBitOverload bool) builtinDefinition {

pkg/sql/sem/builtins/fixed_oids.go

@@ -2861,6 +2861,8 @@ var builtinOidsArray = []string{
 	2906: `levenshtein_less_equal(source: string, target: string, ins_cost: int, del_cost: int, sub_cost: int, max_d: int) -> int`,
 	2907: `crdb_internal.request_transaction_bundle(transaction_fingerprint_id: string, sampling_probability: float, min_execution_latency: interval, expires_after: interval, redacted: bool) -> tuple{int AS request_id, bool AS created}`,
 	2908: `crdb_internal.inject_hint(statement_fingerprint: string, donor_sql: string) -> int`,
+	2909: `crdb_internal.clear_statement_hints_cache() -> void`,
+	2910: `crdb_internal.await_statement_hints_cache() -> void`,
 }
 
 var builtinOidsBySignature map[string]oid.Oid

pkg/sql/sem/eval/deps.go

@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/hintpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgnotice"
@@ -457,6 +458,14 @@ type Planner interface {
 	// ClearTableStatsCache removes all entries from the node's table stats cache.
 	ClearTableStatsCache()
 
+	// ClearStatementHintsCache removes all entries from the node's statement
+	// hints cache.
+	ClearStatementHintsCache()
+
+	// AwaitStatementHintsCache waits for the node's statement hints cache to
+	// catch up with recent hint injections.
+	AwaitStatementHintsCache(ctx context.Context, st *cluster.Settings)
+
 	// RetryCounter is the number of times this statement has been retried.
 	RetryCounter() int