Repository: cockroachdb/pebble
Pull request #4210 (2024-12-22 17:11:41)

db: add UserKeyCategories
Add a way to configure Pebble with an informative partitioning of the
key space. Each partition is a named key category; we use these for
pprof labels.

We currently set these pprof labels for compactions. In the future, we
will allow turning on labels in the read/write paths for specific
investigations.

Sample CRDB configuration code:
```
var userKeyCategories = pebble.MakeUserKeyCategories(
  EngineComparer.Compare,
  []pebble.UserKeyCategory{
    {Name: "local-1", UpperBound: toEngineKey(keys.LocalRangeIDPrefix.AsRawKey())},
    {Name: "rangeid", UpperBound: toEngineKey(keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd())},
    {Name: "local-2", UpperBound: toEngineKey(keys.LocalRangePrefix)},
    {Name: "range", UpperBound: toEngineKey(keys.LocalRangePrefix.PrefixEnd())},
    {Name: "local-3", UpperBound: toEngineKey(keys.LocalRangeLockTablePrefix)},
    {Name: "lock", UpperBound: toEngineKey(keys.LocalRangeLockTablePrefix.PrefixEnd())},
    {Name: "local-4", UpperBound: toEngineKey(keys.LocalPrefix.PrefixEnd())},
    {Name: "meta", UpperBound: toEngineKey(keys.MetaMax)},
    {Name: "system", UpperBound: toEngineKey(keys.SystemMax)},
    {Name: "tenant"},
  }...,
)
```

cockroach-teamcity: 

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/cockroachdb/pebble/4210)
<!-- Reviewable:end -->

RaduBerinde: Sample flame graph with `-tagroot key-type`:
<img width="1125" alt="image" src="https://github.com/user-attachments/assets/4606623f-a05b-4ee7-ab8a-72753313c2ed" />


Patch:
diff --git a/compaction.go b/compaction.go
index 8de36afb..f5f66d77 100644
--- a/compaction.go
+++ b/compaction.go
@@ -36,14 +36,6 @@ var errEmptyTable = errors.New("pebble: empty table")
 // concurrent excise or ingest-split operation.
 var ErrCancelledCompaction = errors.New("pebble: compaction cancelled by a concurrent operation, will retry compaction")
 
-var compactLabelsNoLevel = pprof.Labels("pebble", "compact")
-var compactLabelsPerLevel = func() [manifest.NumLevels]pprof.LabelSet {
-	var labels [manifest.NumLevels]pprof.LabelSet
-	for i := 0; i < manifest.NumLevels; i++ {
-		labels[i] = pprof.Labels("pebble", "compact", "output-level", fmt.Sprintf("L%d", i))
-	}
-	return labels
-}()
 var flushLabels = pprof.Labels("pebble", "flush", "output-level", "L0")
 var gcLabels = pprof.Labels("pebble", "gc")
 
@@ -2180,16 +2172,26 @@ func checkDeleteCompactionHints(
 	return compactLevels, resolvedHints, unresolvedHints
 }
 
+func (d *DB) compactionPprofLabels(c *compaction) pprof.LabelSet {
+	activity := "compact"
+	if len(c.flushing) != 0 {
+		activity = "flush"
+	}
+	level := "L?"
+	// Delete-only compactions don't have an output level.
+	if c.outputLevel != nil {
+		level = fmt.Sprintf("L%d", c.outputLevel.level)
+	}
+	if kc := d.opts.Experimental.UserKeyCategories; kc.Len() > 0 {
+		cat := kc.CategorizeKeyRange(c.smallest.UserKey, c.largest.UserKey)
+		return pprof.Labels("pebble", activity, "output-level", level, "key-type", cat)
+	}
+	return pprof.Labels("pebble", activity, "output-level", level)
+}
+
 // compact runs one compaction and maybe schedules another call to compact.
 func (d *DB) compact(c *compaction, errChannel chan error) {
-	var labels pprof.LabelSet
-	if c.outputLevel == nil {
-		// Delete-only compactions don't have an output level.
-		labels = compactLabelsNoLevel
-	} else {
-		labels = compactLabelsPerLevel[c.outputLevel.level]
-	}
-	pprof.Do(context.Background(), labels, func(context.Context) {
+	pprof.Do(context.Background(), d.compactionPprofLabels(c), func(context.Context) {
 		d.mu.Lock()
 		defer d.mu.Unlock()
 		if err := d.compact1(c, errChannel); err != nil {
diff --git a/options.go b/options.go
index 3726ba21..e8a8455e 100644
--- a/options.go
+++ b/options.go
@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"io"
 	"runtime"
+	"sort"
 	"strconv"
 	"strings"
 	"time"
@@ -859,6 +860,8 @@ type Options struct {
 		// no limiting or pacing happens other than that controlled by other options
 		// like L0CompactionConcurrency and CompactionDebtConcurrency.
 		CompactionLimiter CompactionLimiter
+
+		UserKeyCategories UserKeyCategories
 	}
 
 	// Filters is a map from filter policy name to filter policy. It is used for
@@ -2200,3 +2203,99 @@ func resolveDefaultCompression(c Compression) Compression {
 	}
 	return c
 }
+
+// UserKeyCategories describes a partitioning of the user key space. Each
+// partition is a category with a name. The categories are used for informative
+// purposes only (like pprof labels). Pebble does not treat keys differently
+// based on the UserKeyCategories.
+//
+// The partitions are defined by their upper bounds. The last partition is
+// assumed to go until the end of keyspace; its UpperBound is ignored. The rest
+// of the partitions are ordered by their UpperBound.
+type UserKeyCategories struct {
+	categories []UserKeyCategory
+	cmp        base.Compare
+	// rangeNames[i][j] contains the string referring to the categories in the
+	// range [i, j], with j > i.
+	rangeNames [][]string
+}
+
+// UserKeyCategory describes a partition of the user key space.
+//
+// User keys >= the previous category's UpperBound and < this category's
+// UpperBound are part of this category.
+type UserKeyCategory struct {
+	Name string
+	// UpperBound is the exclusive upper bound of the category. All user keys >= the
+	// previous category's UpperBound and < this UpperBound are part of this
+	// category.
+	UpperBound []byte
+}
+
+// MakeUserKeyCategories creates a UserKeyCategories object with the given
+// categories. The object is immutable and can be reused across different
+// stores.
+func MakeUserKeyCategories(cmp base.Compare, categories ...UserKeyCategory) UserKeyCategories {
+	n := len(categories)
+	if n == 0 {
+		return UserKeyCategories{}
+	}
+	if categories[n-1].UpperBound != nil {
+		panic("last category UpperBound must be nil")
+	}
+	// Verify that the partitions are ordered as expected.
+	for i := 1; i < n-1; i++ {
+		if cmp(categories[i-1].UpperBound, categories[i].UpperBound) >= 0 {
+			panic("invalid UserKeyCategories: key prefixes must be sorted")
+		}
+	}
+
+	// Precalculate a table of range names to avoid allocations in the
+	// categorization path.
+	rangeNamesBuf := make([]string, n*n)
+	rangeNames := make([][]string, n)
+	for i := range rangeNames {
+		rangeNames[i] = rangeNamesBuf[:n]
+		rangeNamesBuf = rangeNamesBuf[n:]
+		for j := i + 1; j < n; j++ {
+			rangeNames[i][j] = categories[i].Name + "-" + categories[j].Name
+		}
+	}
+	return UserKeyCategories{
+		categories: categories,
+		cmp:        cmp,
+		rangeNames: rangeNames,
+	}
+}
+
+// Len returns the number of categories defined.
+func (kc *UserKeyCategories) Len() int {
+	return len(kc.categories)
+}
+
+// CategorizeKey returns the name of the category containing the key.
+func (kc *UserKeyCategories) CategorizeKey(userKey []byte) string {
+	idx := sort.Search(len(kc.categories)-1, func(i int) bool {
+		return kc.cmp(userKey, kc.categories[i].UpperBound) < 0
+	})
+	return kc.categories[idx].Name
+}
+
+// CategorizeKeyRange returns the name of the category containing the key range.
+// If the key range spans multiple categories, the result shows the first and
+// last category separated by a dash, e.g. `cat1-cat5`.
+func (kc *UserKeyCategories) CategorizeKeyRange(startUserKey, endUserKey []byte) string {
+	n := len(kc.categories)
+	p := sort.Search(n-1, func(i int) bool {
+		return kc.cmp(startUserKey, kc.categories[i].UpperBound) < 0
+	})
+	if p == n-1 || kc.cmp(endUserKey, kc.categories[p].UpperBound) < 0 {
+		// Fast path for a single category.
+		return kc.categories[p].Name
+	}
+	// Binary search among the remaining categories.
+	q := p + 1 + sort.Search(n-2-p, func(i int) bool {
+		return kc.cmp(endUserKey, kc.categories[p+1+i].UpperBound) < 0
+	})
+	return kc.rangeNames[p][q]
+}


Test patch:
diff --git a/options_test.go b/options_test.go
index d840366c..3eaf4380 100644
--- a/options_test.go
+++ b/options_test.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"math/rand/v2"
 	"runtime"
+	"strings"
 	"testing"
 	"time"
 
@@ -394,3 +395,65 @@ func TestOptionsValidateCache(t *testing.T) {
 		t.Errorf("Unexpected error message")
 	}
 }
+
+func TestKeyCategories(t *testing.T) {
+	kc := MakeUserKeyCategories(base.DefaultComparer.Compare, []UserKeyCategory{
+		{Name: "b", UpperBound: []byte("b")},
+		{Name: "dd", UpperBound: []byte("dd")},
+		{Name: "e", UpperBound: []byte("e")},
+		{Name: "h", UpperBound: []byte("h")},
+		{Name: "o", UpperBound: nil},
+	}...)
+
+	for _, tc := range []struct {
+		key      string
+		expected string
+	}{
+		{key: "a", expected: "b"},
+		{key: "a123", expected: "b"},
+		{key: "az", expected: "b"},
+		{key: "b", expected: "dd"},
+		{key: "b123", expected: "dd"},
+		{key: "c", expected: "dd"},
+		{key: "d", expected: "dd"},
+		{key: "dd", expected: "e"},
+		{key: "dd0", expected: "e"},
+		{key: "de", expected: "e"},
+		{key: "e", expected: "h"},
+		{key: "e1", expected: "h"},
+		{key: "f", expected: "h"},
+		{key: "h", expected: "o"},
+		{key: "h1", expected: "o"},
+		{key: "z", expected: "o"},
+	} {
+		t.Run(tc.key, func(t *testing.T) {
+			res := kc.CategorizeKey([]byte(tc.key))
+			require.Equal(t, tc.expected, res)
+		})
+	}
+
+	for _, tc := range []struct {
+		rng      string
+		expected string
+	}{
+		{rng: "a-a", expected: "b"},
+		{rng: "a-a1", expected: "b"},
+		{rng: "a-b", expected: "b-dd"},
+		{rng: "b1-b5", expected: "dd"},
+		{rng: "b1-dd0", expected: "dd-e"},
+		{rng: "b1-g", expected: "dd-h"},
+		{rng: "b1-j", expected: "dd-o"},
+		{rng: "b1-z", expected: "dd-o"},
+		{rng: "dd-e", expected: "e-h"},
+		{rng: "h-i", expected: "o"},
+		{rng: "i-i", expected: "o"},
+		{rng: "i-j", expected: "o"},
+		{rng: "a-z", expected: "b-o"},
+	} {
+		t.Run(tc.rng, func(t *testing.T) {
+			keys := strings.SplitN(tc.rng, "-", 2)
+			res := kc.CategorizeKeyRange([]byte(keys[0]), []byte(keys[1]))
+			require.Equal(t, tc.expected, res)
+		})
+	}
+}