Merge #155251

155251: sql/hints: add facilities to invalidate query plans when their hints change r=DrewKimball a=DrewKimball

#### sql/hints: add statement hint cache generation

This commit adds a generation to the statement hints cache, which is
incremented every time the cache processes a rangefeed event. The generation
is used as part of the `cat.DependencyDigest`, used to decide whether an
expensive staleness check is necessary for a cached plan.

Informs #148160

Release note: None

#### sql/hints: include row ID when fetching hints

The row ID of each hint is now included with the hints fetched from the
hint cache. In a future PR, cached query plans will store the row IDs of the
hints they retrieved from the cache, and will compare against the re-fetched
IDs to determine if the query plan is stale (e.g. hints were added or removed).

Informs #148160

Release note: None

Co-authored-by: Drew Kimball <[email protected]>

Author: craig[bot]

pkg/sql/hints/hint_cache.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"hash/fnv"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -87,6 +88,10 @@ type StatementHintsCache struct {
 
 	// Used to read hints from the database when they are not in the cache.
 	db descs.DB
+
+	// generation is incremented any time the hint cache is updated by the
+	// rangefeed.
+	generation atomic.Int64
 }
 
 // cacheSize is the size of the entries to store in the cache.
@@ -214,6 +219,7 @@ func (c *StatementHintsCache) onUpdate(
 // handleInitialScan builds the hintedHashes map and adds it to
 // StatementHintsCache after the initial scan completes.
 func (c *StatementHintsCache) handleInitialScan(update rangefeedcache.Update[*bufferEvent]) {
+	defer c.generation.Add(1)
 	hintedHashes := make(map[int64]hlc.Timestamp)
 	for _, ev := range update.Events {
 		if ev.del {
@@ -246,6 +252,7 @@ func (c *StatementHintsCache) handleIncrementalUpdate(
 		// Avoid synchronization when we're just bumping the resolved timestamp.
 		return
 	}
+	defer c.generation.Add(1)
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	for _, ev := range update.Events {
@@ -309,8 +316,6 @@ func (c *StatementHintsCache) checkHashHasHintsAsync(
 					// Unset the timestamp to indicate that the refresh is done.
 					c.mu.hintedHashes[hash] = hlc.Timestamp{}
 				} else {
-					// There may be a hintCache entry with no hints; drop it.
-					c.mu.hintCache.Del(hash)
 					delete(c.mu.hintedHashes, hash)
 				}
 				return true
@@ -378,18 +383,25 @@ var _ rangefeedbuffer.Event = &bufferEvent{}
 // Reading from the cache.
 // ============================================================================
 
+// GetGeneration returns the current generation, which will change if any
+// modifications happen to the cache.
+func (c *StatementHintsCache) GetGeneration() int64 {
+	return c.generation.Load()
+}
+
 // MaybeGetStatementHints attempts to retrieve the hints for the given statement
-// fingerprint. It returns nil if the statement has no hints, or there was an
-// error retrieving them.
+// fingerprint, along with the unique ID of each hint (for invalidating cached
+// plans). It returns nil if the statement has no hints, or there was an error
+// retrieving them.
 func (c *StatementHintsCache) MaybeGetStatementHints(
 	ctx context.Context, statementFingerprint string,
-) (hints []StatementHint) {
+) (hints []StatementHint, ids []int64) {
 	hash := fnv.New64()
 	_, err := hash.Write([]byte(statementFingerprint))
 	if err != nil {
 		// This should never happen for 64-bit FNV-1.
 		log.Dev.Errorf(ctx, "failed to compute hash for statement fingerprint: %v", err)
-		return nil
+		return nil, nil
 	}
 	statementHash := int64(hash.Sum64())
 	c.mu.Lock()
@@ -399,7 +411,7 @@ func (c *StatementHintsCache) MaybeGetStatementHints(
 		// unconditionally check hintCache.
 		if _, ok := c.mu.hintedHashes[statementHash]; !ok {
 			// There are no plan hints for this query.
-			return nil
+			return nil, nil
 		}
 	}
 	e, ok := c.mu.hintCache.Get(statementHash)
@@ -441,7 +453,7 @@ func (c *StatementHintsCache) maybeWaitForRefreshLocked(
 // released while reading from the db, and then reacquired.
 func (c *StatementHintsCache) addCacheEntryLocked(
 	ctx context.Context, statementHash int64, statementFingerprint string,
-) []StatementHint {
+) (hints []StatementHint, ids []int64) {
 	c.mu.AssertHeld()
 
 	// Add a cache entry that other queries can find and wait on until we have the
@@ -470,20 +482,24 @@ func (c *StatementHintsCache) addCacheEntryLocked(
 		// Remove the cache entry so that the next caller retries the query.
 		c.mu.hintCache.Del(statementHash)
 		log.VEventf(ctx, 1, "encountered error while reading hints for query %s", statementFingerprint)
-		return nil
+		return nil, nil
 	}
 	return entry.getMatchingHints(statementFingerprint)
 }
 
 // getStatementHintsFromDB queries the system.statement_hints table for hints
 // matching the given fingerprint hash. It is able to handle the case when
 // multiple fingerprints match the hash, as well as the case when there are no
-// hints for the fingerprint.
+// hints for the fingerprint. Results are ordered by row ID.
 func (c *StatementHintsCache) getStatementHintsFromDB(
 	ctx context.Context, statementHash int64, entry *cacheEntry,
 ) (retErr error) {
 	const opName = "get-plan-hints"
-	const getHintsStmt = `SELECT "fingerprint", "hint" FROM system.statement_hints WHERE "hash" = $1`
+	const getHintsStmt = `
+    SELECT "row_id", "fingerprint", "hint"
+    FROM system.statement_hints
+    WHERE "hash" = $1
+    ORDER BY "row_id" ASC`
 	it, err := c.db.Executor().QueryIteratorEx(
 		ctx, opName, nil /* txn */, sessiondata.NodeUserSessionDataOverride,
 		getHintsStmt, statementHash,
@@ -500,11 +516,13 @@ func (c *StatementHintsCache) getStatementHintsFromDB(
 			return err
 		}
 		datums := it.Cur()
-		fingerprint := string(tree.MustBeDString(datums[0]))
-		hint, err := NewStatementHint([]byte(tree.MustBeDBytes(datums[1])))
+		rowID := int64(tree.MustBeDInt(datums[0]))
+		fingerprint := string(tree.MustBeDString(datums[1]))
+		hint, err := NewStatementHint([]byte(tree.MustBeDBytes(datums[2])))
 		if err != nil {
 			return err
 		}
+		entry.ids = append(entry.ids, rowID)
 		entry.fingerprints = append(entry.fingerprints, fingerprint)
 		entry.hints = append(entry.hints, hint)
 	}
@@ -517,30 +535,33 @@ type cacheEntry struct {
 	mustWait bool
 	waitCond sync.Cond
 
-	// fingerprints and hints have the same length. fingerprints[i] is the
-	// statement fingerprint to which hints[i] applies. We have to keep track of
-	// the fingerprints to resolve hash collisions. Note that a single fingerprint
-	// can have multiple hints, in which case there will be duplicate entries in
-	// the fingerprints slice.
+	// hints, fingerprints, and ids have the same length. fingerprints[i] is the
+	// statement fingerprint to which hints[i] applies, while ids[i] uniquely
+	// identifies a hint in the system table. They are kept in order of id.
+	//
+	// We track the hint ID for invalidating cached query plans after a hint is
+	// added or removed. We track the fingerprint to resolve hash collisions. Note
+	// that a single fingerprint can have multiple hints, in which case there will
+	// be duplicate entries in the fingerprints slice.
 	// TODO(drewk): consider de-duplicating the fingerprint strings to reduce
 	// memory usage.
-	fingerprints []string
 	hints        []StatementHint
+	fingerprints []string
+	ids          []int64
 }
 
-// getMatchingHints returns the plan hints for the given fingerprint, or nil if
-// they don't exist.
-func (entry *cacheEntry) getMatchingHints(statementFingerprint string) []StatementHint {
-	var res []StatementHint
+// getMatchingHints returns the plan hints and row IDs for the given
+// fingerprint, or nil if they don't exist. The results are in order of row ID.
+func (entry *cacheEntry) getMatchingHints(
+	statementFingerprint string,
+) (hints []StatementHint, ids []int64) {
 	for i := range entry.hints {
 		if entry.fingerprints[i] == statementFingerprint {
-			// TODO(drewk,michae2): we could do something smarter here to avoid
-			// allocations. For example, order by fingerprint in the original query
-			// so that we can return a slice here.
-			res = append(res, entry.hints[i])
+			hints = append(hints, entry.hints[i])
+			ids = append(ids, entry.ids[i])
 		}
 	}
-	return res
+	return hints, ids
 }
 
 // ============================================================================