Merge #142236

142236: split: add key itself to popular key surface r=angles-n-daemons a=angles-n-daemons

split: add key itself to popular key surface

historically, load based splitters have only ever exposed the frequency with which a popular key is accessed, and not the key itself. this change updates this so that the key is included in the function surface.

this may be used later, as part of the replica's split statistics.

Fixes: #138758
Epic: CRDB-43150

Release note: None

Co-authored-by: Brian Dillmann <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/split/decider.go

@@ -23,6 +23,11 @@ const minSplitSuggestionInterval = time.Minute
 const minNoSplitKeyLoggingMetricsInterval = time.Minute
 const minPerSecondSampleDuration = time.Second
 
+type PopularKey struct {
+	Key       roachpb.Key
+	Frequency float64
+}
+
 type LoadBasedSplitter interface {
 	redact.SafeFormatter
 	// Record informs the LoadBasedSplitter about where the span lies with regard
@@ -43,9 +48,9 @@ type LoadBasedSplitter interface {
 	// empty string.
 	NoSplitKeyCauseLogMsg() redact.RedactableString
 
-	// PopularKeyFrequency returns the percentage that the most popular key
-	// appears in the sampled candidate split keys.
-	PopularKeyFrequency() float64
+	// PopularKey returns the most popular key in the sample dataset in addition
+	// to its frequency..
+	PopularKey() PopularKey
 
 	// String formats the state of the load based splitter.
 	String() string
@@ -262,7 +267,7 @@ func (d *Decider) recordLocked(
 				if now.Sub(d.mu.lastNoSplitKeyLoggingMetrics) > minNoSplitKeyLoggingMetricsInterval {
 					d.mu.lastNoSplitKeyLoggingMetrics = now
 					if causeMsg := d.mu.splitFinder.NoSplitKeyCauseLogMsg(); causeMsg != "" {
-						popularKeyFrequency := d.mu.splitFinder.PopularKeyFrequency()
+						popularKeyFrequency := d.mu.splitFinder.PopularKey().Frequency
 						log.KvDistribution.Infof(ctx, "%s, most popular key occurs in %d%% of samples",
 							causeMsg, int(popularKeyFrequency*100))
 						log.KvDistribution.VInfof(ctx, 3, "splitter_state=%v", (*lockedDecider)(d))

pkg/kv/kvserver/split/unweighted_finder.go

@@ -201,26 +201,31 @@ func (f *UnweightedFinder) NoSplitKeyCauseLogMsg() redact.RedactableString {
 		imbalanceAndTooManyContained)
 }
 
-// PopularKeyFrequency implements the LoadBasedSplitter interface.
-func (f *UnweightedFinder) PopularKeyFrequency() float64 {
+// PopularKey implements the LoadBasedSplitter interface.
+func (f *UnweightedFinder) PopularKey() PopularKey {
 	slices.SortFunc(f.samples[:], func(a, b sample) int {
 		return bytes.Compare(a.key, b.key)
 	})
 
 	currentKeyCount := 1
-	popularKeyCount := 1
+	popularKeyCount := 0
+	var key roachpb.Key
 	for i := 1; i < len(f.samples); i++ {
 		if bytes.Equal(f.samples[i].key, f.samples[i-1].key) {
 			currentKeyCount++
 		} else {
 			currentKeyCount = 1
 		}
 		if popularKeyCount < currentKeyCount {
+			key = f.samples[i].key
 			popularKeyCount = currentKeyCount
 		}
 	}
 
-	return float64(popularKeyCount) / float64(splitKeySampleSize)
+	return PopularKey{
+		Key:       key,
+		Frequency: float64(popularKeyCount) / float64(splitKeySampleSize),
+	}
 }
 
 // SafeFormat implements the redact.SafeFormatter interface.

pkg/kv/kvserver/split/unweighted_finder_test.go

@@ -8,6 +8,7 @@ package split
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"math/rand"
 	"reflect"
 	"testing"
@@ -381,23 +382,25 @@ func TestFinderPopularKeyFrequency(t *testing.T) {
 		}
 	}
 
-	testCases := []struct {
+	randSource := rand.New(rand.NewSource(2022))
+	for i, test := range []struct {
 		samples                     [splitKeySampleSize]sample
+		expectedPopularKey          roachpb.Key
 		expectedPopularKeyFrequency float64
 	}{
-		{uniqueKeySample, 0.05},
-		{twentyPercentPopularKeySample, 0.2},
-		{twentyFivePercentPopularKeySample, 0.25},
-		{fiftyPercentPopularKeySample, 0.5},
-		{fiftyFivePercentPopularKeySample, 0.55},
-		{sameKeySample, 1},
-	}
-
-	randSource := rand.New(rand.NewSource(2022))
-	for i, test := range testCases {
-		finder := NewUnweightedFinder(timeutil.Now(), randSource)
-		finder.samples = test.samples
-		popularKeyFrequency := finder.PopularKeyFrequency()
-		assert.Equal(t, test.expectedPopularKeyFrequency, popularKeyFrequency, "unexpected popular key frequency in test %d", i)
+		{uniqueKeySample, keys.SystemSQLCodec.TablePrefix(1), 0.05},
+		{twentyPercentPopularKeySample, keys.SystemSQLCodec.TablePrefix(6), 0.2},
+		{twentyFivePercentPopularKeySample, keys.SystemSQLCodec.TablePrefix(2), 0.25},
+		{fiftyPercentPopularKeySample, keys.SystemSQLCodec.TablePrefix(0), 0.5},
+		{fiftyFivePercentPopularKeySample, keys.SystemSQLCodec.TablePrefix(0), 0.55},
+		{sameKeySample, keys.SystemSQLCodec.TablePrefix(0), 1},
+	} {
+		t.Run(fmt.Sprintf("popular key test %d", i), func(t *testing.T) {
+			finder := NewUnweightedFinder(timeutil.Now(), randSource)
+			finder.samples = test.samples
+			popularKey := finder.PopularKey()
+			assert.Equal(t, test.expectedPopularKey, popularKey.Key, "unexpected popular key in test %d", i)
+			assert.Equal(t, test.expectedPopularKeyFrequency, popularKey.Frequency, "unexpected popular key frequency in test %d", i)
+		})
 	}
 }

pkg/kv/kvserver/split/weighted_finder.go

@@ -252,8 +252,8 @@ func (f *WeightedFinder) NoSplitKeyCauseLogMsg() redact.RedactableString {
 		insufficientCounters, imbalance)
 }
 
-// PopularKeyFrequency implements the LoadBasedSplitter interface.
-func (f *WeightedFinder) PopularKeyFrequency() float64 {
+// PopularKey implements the LoadBasedSplitter interface.
+func (f *WeightedFinder) PopularKey() PopularKey {
 	// Sort the sample slice to determine the frequency that a popular key
 	// appears. We could copy the slice, however it would require an allocation.
 	// The probability a sample is replaced doesn't change as it is independent
@@ -263,6 +263,7 @@ func (f *WeightedFinder) PopularKeyFrequency() float64 {
 	})
 
 	weight := f.samples[0].weight
+	key := f.samples[0].key
 	currentKeyWeight := weight
 	popularKeyWeight := weight
 	totalWeight := weight
@@ -274,12 +275,16 @@ func (f *WeightedFinder) PopularKeyFrequency() float64 {
 			currentKeyWeight = weight
 		}
 		if popularKeyWeight < currentKeyWeight {
+			key = f.samples[i].key
 			popularKeyWeight = currentKeyWeight
 		}
 		totalWeight += weight
 	}
 
-	return popularKeyWeight / totalWeight
+	return PopularKey{
+		Key:       key,
+		Frequency: popularKeyWeight / totalWeight,
+	}
 }
 
 // SafeFormat implements the redact.SafeFormatter interface.

pkg/kv/kvserver/split/weighted_finder_test.go

@@ -335,7 +335,7 @@ func TestWeightedFinderNoSplitKeyCause(t *testing.T) {
 	assert.Equal(t, 13, imbalance, "unexpected imbalance counters")
 }
 
-func TestWeightedFinderPopularKeyFrequency(t *testing.T) {
+func TestWeightedFinderPopularKey(t *testing.T) {
 	uniqueKeyUnweightedSample := [splitKeySampleSize]weightedSample{}
 	for i, idx := range rand.Perm(splitKeySampleSize) {
 		uniqueKeyUnweightedSample[idx] = weightedSample{
@@ -387,22 +387,24 @@ func TestWeightedFinderPopularKeyFrequency(t *testing.T) {
 	const eps = 1e-3
 	testCases := []struct {
 		samples                     [splitKeySampleSize]weightedSample
+		expectedPopularKey          roachpb.Key
 		expectedPopularKeyFrequency float64
 	}{
-		{uniqueKeyUnweightedSample, 1.0 / 20.0},
-		{uniqueKeyWeightedSample, 20.0 / 210.0}, // 20/(1+2+...+20)
-		{duplicateKeyUnweightedSample, 5.0 / 20.0},
-		{duplicateKeyWeightedSample, 84.0 / 210.0}, // (9+10+...+15)/(1+2+...+20)
-		{sameKeySample, 1},
+		{uniqueKeyUnweightedSample, keys.SystemSQLCodec.TablePrefix(uint32(0)), 1.0 / 20.0},
+		{uniqueKeyWeightedSample, keys.SystemSQLCodec.TablePrefix(uint32(19)), 20.0 / 210.0}, // 20/(1+2+...+20)
+		{duplicateKeyUnweightedSample, keys.SystemSQLCodec.TablePrefix(uint32(2)), 5.0 / 20.0},
+		{duplicateKeyWeightedSample, keys.SystemSQLCodec.TablePrefix(uint32(2)), 84.0 / 210.0}, // (9+10+...+15)/(1+2+...+20)
+		{sameKeySample, keys.SystemSQLCodec.TablePrefix(uint32(0)), 1},
 	}
 
 	randSource := rand.New(rand.NewSource(2022))
 	for i, test := range testCases {
 		weightedFinder := NewWeightedFinder(timeutil.Now(), randSource)
 		weightedFinder.samples = test.samples
-		popularKeyFrequency := weightedFinder.PopularKeyFrequency()
-		assert.True(t, math.Abs(test.expectedPopularKeyFrequency-popularKeyFrequency) < eps,
+		popularKey := weightedFinder.PopularKey()
+		assert.Equal(t, test.expectedPopularKey, popularKey.Key)
+		assert.True(t, math.Abs(test.expectedPopularKeyFrequency-popularKey.Frequency) < eps,
 			"%d: expected popular key frequency %f, got %f",
-			i, test.expectedPopularKeyFrequency, popularKeyFrequency)
+			i, test.expectedPopularKeyFrequency, popularKey.Frequency)
 	}
 }