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

Author: DrewKimball

pkg/sql/hints/hint_cache.go

@@ -390,17 +390,18 @@ func (c *StatementHintsCache) GetGeneration() int64 {
 }
 
 // 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()
@@ -410,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)
@@ -452,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
@@ -481,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,
@@ -511,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)
 	}
@@ -528,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
 }
 
 // ============================================================================

pkg/sql/hints/hint_cache_test.go

@@ -59,7 +59,7 @@ func TestHintCacheBasic(t *testing.T) {
 	// Query for hints on a statement that has no hints.
 	nonHintedFingerprint := "SELECT x FROM t WHERE y = $1"
 	require.False(t, hc.TestingHashHasHints(computeHash(t, nonHintedFingerprint)))
-	require.Nil(t, hc.MaybeGetStatementHints(ctx, nonHintedFingerprint))
+	requireHintsCount(t, hc, ctx, nonHintedFingerprint, 0)
 
 	// Add a hint for another statement.
 	fingerprint2 := "SELECT c FROM t WHERE d = $1"
@@ -122,43 +122,43 @@ func TestHintCacheLRU(t *testing.T) {
 
 	// Access the first two fingerprints to populate the cache.
 	// This should result in 2 table reads.
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[0]))
+	requireHintsCount(t, hc, ctx, fingerprints[0], 1)
 	require.Equal(t, 1, hc.TestingNumTableReads())
 
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[1]))
+	requireHintsCount(t, hc, ctx, fingerprints[1], 1)
 	require.Equal(t, 2, hc.TestingNumTableReads())
 
 	// Access the same fingerprints again - should be served from cache with no
 	// additional reads.
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[0]))
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[1]))
+	requireHintsCount(t, hc, ctx, fingerprints[0], 1)
+	requireHintsCount(t, hc, ctx, fingerprints[1], 1)
 	require.Equal(t, 2, hc.TestingNumTableReads())
 
 	// Access the third fingerprint. This should evict the first (LRU) due to
 	// cache size limit of 2, resulting in one more table read.
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[2]))
+	requireHintsCount(t, hc, ctx, fingerprints[2], 1)
 	require.Equal(t, 3, hc.TestingNumTableReads())
 
 	// Access the first fingerprint again. Since it was evicted, this should
 	// result in another table read on the first access.
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[0]))
+	requireHintsCount(t, hc, ctx, fingerprints[0], 1)
 	require.Equal(t, 4, hc.TestingNumTableReads())
 
 	// Access the second fingerprint. It should have been evicted by now, so
 	// another table read on the first access.
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[1]))
+	requireHintsCount(t, hc, ctx, fingerprints[1], 1)
 	require.Equal(t, 5, hc.TestingNumTableReads())
 
 	// The first and second fingerprint should now be cached, so accessing them
 	// again should not increase table reads.
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[0]))
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[1]))
+	requireHintsCount(t, hc, ctx, fingerprints[0], 1)
+	requireHintsCount(t, hc, ctx, fingerprints[1], 1)
 	require.Equal(t, 5, hc.TestingNumTableReads())
 
 	// Access the third fingerprint again - should have been evicted, so the first
 	// access should cause a table read.
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[2]))
-	require.NotNil(t, hc.MaybeGetStatementHints(ctx, fingerprints[2]))
+	requireHintsCount(t, hc, ctx, fingerprints[2], 1)
+	requireHintsCount(t, hc, ctx, fingerprints[2], 1)
 	require.Equal(t, 6, hc.TestingNumTableReads())
 }
 
@@ -204,15 +204,15 @@ func TestHintCacheInitialScan(t *testing.T) {
 	require.True(t, hc.TestingHashHasHints(hashes[0]))
 	require.True(t, hc.TestingHashHasHints(hashes[1]))
 	require.True(t, hc.TestingHashHasHints(hashes[2]))
-	require.Len(t, hc.MaybeGetStatementHints(ctx, fingerprints[0]), 2)
-	require.Len(t, hc.MaybeGetStatementHints(ctx, fingerprints[1]), 1)
-	require.Len(t, hc.MaybeGetStatementHints(ctx, fingerprints[2]), 1)
+	requireHintsCount(t, hc, ctx, fingerprints[0], 2)
+	requireHintsCount(t, hc, ctx, fingerprints[1], 1)
+	requireHintsCount(t, hc, ctx, fingerprints[2], 1)
 
 	// Query for a fingerprint with no hints.
 	nonHintedFingerprint := "SELECT x FROM t WHERE y = $1"
 	nonHintedHash := computeHash(t, nonHintedFingerprint)
 	require.False(t, hc.TestingHashHasHints(nonHintedHash))
-	require.Nil(t, hc.MaybeGetStatementHints(ctx, nonHintedFingerprint))
+	requireHintsCount(t, hc, ctx, nonHintedFingerprint, 0)
 
 	// After the initial scan, new hints should still be detected via rangefeed.
 	fingerprint4 := "SELECT g FROM t WHERE h = $1"
@@ -274,8 +274,8 @@ func TestHintCacheMultiNode(t *testing.T) {
 	require.Equal(t, 1, hc.TestingHashCount())
 
 	// Verify that fingerprint1 still has hints but fingerprint2 doesn't.
-	require.Len(t, hc.MaybeGetStatementHints(ctx, fingerprint1), 1)
-	require.Nil(t, hc.MaybeGetStatementHints(ctx, fingerprint2))
+	requireHintsCount(t, hc, ctx, fingerprint1, 1)
+	requireHintsCount(t, hc, ctx, fingerprint2, 0)
 }
 
 // TestHintCacheMultiTenant tests that hint caches for different tenants are
@@ -317,24 +317,24 @@ func TestHintCacheMultiTenant(t *testing.T) {
 	insertStatementHint(t, r1, fingerprint1)
 	waitForUpdateOnFingerprintHash(t, ctx, hc1, fingerprint1, 1)
 	require.Equal(t, 1, hc1.TestingHashCount())
-	require.Len(t, hc1.MaybeGetStatementHints(ctx, fingerprint1), 1)
+	requireHintsCount(t, hc1, ctx, fingerprint1, 1)
 
 	// Tenant2's cache should remain empty.
 	require.Equal(t, 0, hc2.TestingHashCount())
 	require.False(t, hc2.TestingHashHasHints(computeHash(t, fingerprint1)))
-	require.Nil(t, hc2.MaybeGetStatementHints(ctx, fingerprint1))
+	requireHintsCount(t, hc2, ctx, fingerprint1, 0)
 
 	// Insert a different hint for tenant2.
 	fingerprint2 := "SELECT c FROM t WHERE d = $1"
 	insertStatementHint(t, r2, fingerprint2)
 	waitForUpdateOnFingerprintHash(t, ctx, hc2, fingerprint2, 1)
 	require.Equal(t, 1, hc2.TestingHashCount())
-	require.Len(t, hc2.MaybeGetStatementHints(ctx, fingerprint2), 1)
+	requireHintsCount(t, hc2, ctx, fingerprint2, 1)
 
 	// Tenant1's cache should still only have its original hint.
 	require.Equal(t, 1, hc1.TestingHashCount())
 	require.False(t, hc1.TestingHashHasHints(computeHash(t, fingerprint2)))
-	require.Nil(t, hc1.MaybeGetStatementHints(ctx, fingerprint2))
+	requireHintsCount(t, hc1, ctx, fingerprint2, 0)
 
 	// Insert the same fingerprint in both tenants - they should be independent.
 	fingerprintShared := "SELECT x FROM t WHERE y = $1"
@@ -352,8 +352,8 @@ func TestHintCacheMultiTenant(t *testing.T) {
 	// Verify that the shared fingerprint has different counts in each tenant.
 	require.True(t, hc1.TestingHashHasHints(computeHash(t, fingerprintShared)))
 	require.True(t, hc2.TestingHashHasHints(computeHash(t, fingerprintShared)))
-	require.Len(t, hc1.MaybeGetStatementHints(ctx, fingerprintShared), 2)
-	require.Len(t, hc2.MaybeGetStatementHints(ctx, fingerprintShared), 1)
+	requireHintsCount(t, hc1, ctx, fingerprintShared, 2)
+	requireHintsCount(t, hc2, ctx, fingerprintShared, 1)
 
 	// Delete one hint from tenant1's shared fingerprint.
 	deleteStatementHints(t, r1, fingerprintShared, 1)
@@ -370,8 +370,8 @@ func TestHintCacheMultiTenant(t *testing.T) {
 	// Tenant1 should still have the hint, but tenant2 should not.
 	require.True(t, hc1.TestingHashHasHints(computeHash(t, fingerprintShared)))
 	require.False(t, hc2.TestingHashHasHints(computeHash(t, fingerprintShared)))
-	require.Len(t, hc1.MaybeGetStatementHints(ctx, fingerprintShared), 1)
-	require.Nil(t, hc2.MaybeGetStatementHints(ctx, fingerprintShared))
+	requireHintsCount(t, hc1, ctx, fingerprintShared, 1)
+	requireHintsCount(t, hc2, ctx, fingerprintShared, 0)
 }
 
 // TestHintCacheGeneration tests that the cache generation is incremented
@@ -505,14 +505,39 @@ func waitForUpdateOnFingerprintHash(
 		if hasHints := hc.TestingHashHasHints(hash); (expected > 0) != hasHints {
 			return errors.Errorf("expected hash %d with hasHints=%t, got hasHints=%t", hash, expected > 0, hasHints)
 		}
-		hints := hc.MaybeGetStatementHints(ctx, fingerprint)
+		hints, ids := hc.MaybeGetStatementHints(ctx, fingerprint)
 		if len(hints) != expected {
 			return errors.Errorf("expected %d hints for fingerprint %q, got %d", expected, fingerprint, len(hints))
 		}
+		checkIDOrder(t, ids)
 		return nil
 	})
 }
 
+// requireHintsCount verifies that MaybeGetStatementHints returns the expected number of hints and IDs.
+func requireHintsCount(
+	t *testing.T,
+	hc *hints.StatementHintsCache,
+	ctx context.Context,
+	fingerprint string,
+	expectedCount int,
+) {
+	hints, ids := hc.MaybeGetStatementHints(ctx, fingerprint)
+	require.Len(t, hints, expectedCount)
+	require.Len(t, ids, expectedCount)
+	checkIDOrder(t, ids)
+}
+
+// checkIDOrder verifies that the row IDs are in ascending order.
+func checkIDOrder(t *testing.T, ids []int64) {
+	t.Helper()
+	for i := 1; i < len(ids); i++ {
+		if ids[i] <= ids[i-1] {
+			t.Fatalf("expected IDs to be in ascending order, got %v", ids)
+		}
+	}
+}
+
 // insertStatementHint inserts an empty statement hint into the
 // system.statement_hints table.
 func insertStatementHint(t *testing.T, r *sqlutils.SQLRunner, fingerprint string) {