db: add ability to collect a profile of value retrievals

With the introduction of value separation, retrievals of separated values are
more expensive than retrievals of inlined values. Pebble exports an aggregate
metric recording the total count of value retrievals that must load a value
from an external blob file. This can be used to help attribute performance
issues to excessive separated value retrievals, but it doesn't help with
understanding what codepaths are performing the retrievals.

This commit adds support for collection of a 'separated value retrieval'
profile, aggregating the count and byte-total of value retrievals per unique
stack trace while profile collection is active. This additional instrumentation
adds some overhead and is too expensive to perform by default.

Users may call DB.RecordSeparatedValueRetrievals to initiate collection of a
profile. All subsequent opened iterators and compactions will record retrievals
of separated values into the profile until the caller stops the collection with
the invocation of a callback.

For now this profile can only be dumped to a text format through a String
method. We could export it to a pprof file but it would require a fair bit of
work.

Close #5535.
Informs cockroachdb/cockroach#144715.

Author: jbowens

compaction.go

@@ -3390,6 +3390,7 @@ func (d *DB) compactAndWrite(
 			uint64(uintptr(unsafe.Pointer(c))),
 			categoryCompaction,
 		),
+		ValueRetrievalProfile: d.valueRetrievalProfile.Load(),
 	}
 	if c.version != nil {
 		c.iterationState.valueFetcher.Init(&c.version.BlobFiles, d.fileCache, blockReadEnv, suggestedCacheReaders)

db.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/arenaskl"
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/bytesprofile"
 	"github.com/cockroachdb/pebble/internal/cache"
 	"github.com/cockroachdb/pebble/internal/deletepacer"
 	"github.com/cockroachdb/pebble/internal/inflight"
@@ -554,7 +555,8 @@ type DB struct {
 
 	compressionCounters block.CompressionCounters
 
-	iterTracker *inflight.Tracker
+	iterTracker           *inflight.Tracker
+	valueRetrievalProfile atomic.Pointer[bytesprofile.Profile]
 }
 
 var _ Reader = (*DB)(nil)
@@ -565,6 +567,27 @@ func (d *DB) TestOnlyWaitForCleaning() {
 	d.deletePacer.WaitForTesting()
 }
 
+// RecordSeparatedValueRetrievals begins collection of stack traces that are
+// retrieving separated values. If a profile has already begun, an error is
+// returned. If successful, the stop function must be called to stop the
+// collection and return the resulting profile.
+func (d *DB) RecordSeparatedValueRetrievals() (
+	stop func() *metrics.ValueRetrievalProfile,
+	err error,
+) {
+	p := bytesprofile.NewProfile()
+	swapped := d.valueRetrievalProfile.CompareAndSwap(nil, p)
+	if !swapped {
+		return nil, errors.New("separated value retrieval profile is already in progress")
+	}
+	return func() *metrics.ValueRetrievalProfile {
+		if !d.valueRetrievalProfile.CompareAndSwap(p, nil) {
+			panic(errors.AssertionFailedf("profile already stopped"))
+		}
+		return p
+	}, nil
+}
+
 // Set sets the value for the given key. It overwrites any previous value
 // for that key; a DB is not a multi-map.
 //
@@ -1106,6 +1129,7 @@ func (d *DB) newIter(
 	dbi.fc = d.fileCache
 	dbi.newIters = newIters
 	dbi.newIterRangeKey = newIterRangeKey
+	dbi.valueRetrievalProfile = d.valueRetrievalProfile.Load()
 	dbi.seqNum = seqNum
 	dbi.batchOnlyIter = newIterOpts.batch.batchOnly
 	if o != nil {
@@ -1255,6 +1279,7 @@ func (i *Iterator) constructPointIter(
 			uint64(uintptr(unsafe.Pointer(i))),
 			i.opts.Category,
 		),
+		ValueRetrievalProfile: i.valueRetrievalProfile,
 	}
 	if i.readState != nil {
 		i.blobValueFetcher.Init(&i.readState.current.BlobFiles, i.fc, readEnv,

internal/bytesprofile/bytesprofile.go

@@ -0,0 +1,91 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package bytesprofile
+
+import (
+	"cmp"
+	"fmt"
+	"maps"
+	"runtime"
+	"slices"
+	"strings"
+	"sync"
+
+	"github.com/cockroachdb/pebble/internal/humanize"
+)
+
+// Profile is a profile mapping stack traces to a cumulative count and byte sum.
+type Profile struct {
+	mu      sync.Mutex
+	samples map[stack]aggSamples
+}
+
+// NewProfile creates a new profile.
+func NewProfile() *Profile {
+	return &Profile{samples: make(map[stack]aggSamples)}
+}
+
+type stack [20]uintptr
+
+// trimmed returns the non-zero stack frames of stack.
+func (s stack) trimmed() []uintptr {
+	for i := range s {
+		if s[i] == 0 {
+			return s[:i]
+		}
+	}
+	return s[:]
+}
+
+type aggSamples struct {
+	bytes int64
+	count int64
+}
+
+// Record records a sample of the given number of bytes with the calling stack trace.
+func (p *Profile) Record(bytes int64) {
+	var stack stack
+	runtime.Callers(2, stack[:])
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	curr := p.samples[stack]
+	curr.bytes += bytes
+	curr.count++
+	p.samples[stack] = curr
+}
+
+// TODO(jackson): We could add the ability to export the profile to a pprof file
+// (which internally is just a protocol buffer). Ideally the Go standard library
+// would provide facilities for this (e.g., golang/go#18454). The runtime/pprof
+// library comes close with its definition of custom profiles, but they only
+// support profiles tracking in-use resources.
+
+// String returns a string representation of the stacks captured by the profile.
+func (p *Profile) String() string {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	// Sort the stacks by bytes in descending order.
+	uniqueStacks := slices.SortedFunc(maps.Keys(p.samples), func(a, b stack) int {
+		return -cmp.Compare(p.samples[a].bytes, p.samples[b].bytes)
+	})
+	var sb strings.Builder
+	for i, stack := range uniqueStacks {
+		if i > 0 {
+			sb.WriteString("\n")
+		}
+		fmt.Fprintf(&sb, "%d: Count: %d (%s), Bytes: %d (%s)\n", i,
+			p.samples[stack].count, humanize.Count.Int64(p.samples[stack].count),
+			p.samples[stack].bytes, humanize.Bytes.Int64(p.samples[stack].bytes))
+		frames := runtime.CallersFrames(stack.trimmed())
+		for {
+			frame, more := frames.Next()
+			fmt.Fprintf(&sb, "  %s\n   %s:%d\n", frame.Function, frame.File, frame.Line)
+			if !more {
+				break
+			}
+		}
+	}
+	return sb.String()
+}

internal/bytesprofile/bytesprofile_test.go

@@ -0,0 +1,25 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package bytesprofile
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBytesProfile(t *testing.T) {
+	p := NewProfile()
+	for i := 0; i < 100; i++ {
+		p.Record(int64(i))
+	}
+	for i := 0; i < 10; i++ {
+		p.Record(11 * int64(i))
+	}
+	s := p.String()
+	require.Contains(t, s, "0: Count: 100 (100), Bytes: 4950 (4.8KB)")
+	require.Contains(t, s, "1: Count: 10 (10), Bytes: 495 (495B)")
+	t.Log(s)
+}

iterator.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/bytealloc"
+	"github.com/cockroachdb/pebble/internal/bytesprofile"
 	"github.com/cockroachdb/pebble/internal/humanize"
 	"github.com/cockroachdb/pebble/internal/inflight"
 	"github.com/cockroachdb/pebble/internal/invariants"
@@ -260,11 +261,12 @@ type Iterator struct {
 	externalIter        *externalIterState
 	// Following fields used when constructing an iterator stack, eg, in Clone
 	// and SetOptions or when re-fragmenting a batch's range keys/range dels.
-	fc               *fileCacheHandle
-	newIters         tableNewIters
-	newIterRangeKey  keyspanimpl.TableNewSpanIter
-	lazyCombinedIter lazyCombinedIter
-	seqNum           base.SeqNum
+	fc                    *fileCacheHandle
+	newIters              tableNewIters
+	newIterRangeKey       keyspanimpl.TableNewSpanIter
+	valueRetrievalProfile *bytesprofile.Profile
+	lazyCombinedIter      lazyCombinedIter
+	seqNum                base.SeqNum
 	// batch is non-nil if this Iterator includes an indexed batch. Batch
 	// contains all the state pertaining to iterating over the indexed batch.
 	// The iteratorBatchState struct is bundled within the iterAlloc struct to
@@ -2889,21 +2891,22 @@ func (i *Iterator) CloneWithContext(ctx context.Context, opts CloneOptions) (*It
 	buf := newIterAlloc()
 	dbi := &buf.dbi
 	*dbi = Iterator{
-		ctx:                 ctx,
-		opts:                *opts.IterOptions,
-		alloc:               buf,
-		merge:               i.merge,
-		comparer:            i.comparer,
-		readState:           readState,
-		version:             vers,
-		keyBuf:              buf.keyBuf,
-		prefixOrFullSeekKey: buf.prefixOrFullSeekKey,
-		boundsBuf:           buf.boundsBuf,
-		fc:                  i.fc,
-		newIters:            i.newIters,
-		newIterRangeKey:     i.newIterRangeKey,
-		seqNum:              i.seqNum,
-		tracker:             i.tracker,
+		ctx:                   ctx,
+		opts:                  *opts.IterOptions,
+		alloc:                 buf,
+		merge:                 i.merge,
+		comparer:              i.comparer,
+		readState:             readState,
+		version:               vers,
+		keyBuf:                buf.keyBuf,
+		prefixOrFullSeekKey:   buf.prefixOrFullSeekKey,
+		boundsBuf:             buf.boundsBuf,
+		fc:                    i.fc,
+		newIters:              i.newIters,
+		newIterRangeKey:       i.newIterRangeKey,
+		valueRetrievalProfile: i.valueRetrievalProfile,
+		seqNum:                i.seqNum,
+		tracker:               i.tracker,
 	}
 	if i.tracker != nil && !dbi.opts.ExemptFromTracking {
 		dbi.trackerHandle = i.tracker.Start()

iterator_test.go

@@ -1415,6 +1415,39 @@ func TestIteratorStatsMerge(t *testing.T) {
 	require.Equal(t, expected, s)
 }
 
+func TestIteratorValueRetrievalProfile(t *testing.T) {
+	opts := &Options{}
+	opts.FormatMajorVersion = internalFormatNewest
+	opts.Experimental.ValueSeparationPolicy = func() ValueSeparationPolicy {
+		return ValueSeparationPolicy{
+			Enabled:               true,
+			MinimumSize:           1,
+			MaxBlobReferenceDepth: 5,
+		}
+	}
+	d := newTestkeysDatabase(t, opts, testkeys.Alpha(2), rand.New(rand.NewPCG(1, 1)))
+	defer func() { require.NoError(t, d.Close()) }()
+	require.NoError(t, d.Flush())
+	stop, err := d.RecordSeparatedValueRetrievals()
+	require.NoError(t, err)
+	{
+		_, err := d.RecordSeparatedValueRetrievals()
+		require.Error(t, err, "should not be able to start a second profile")
+	}
+	it, err := d.NewIter(nil)
+	require.NoError(t, err)
+	defer it.Close()
+	require.True(t, it.SeekGE([]byte("a")))
+	val, err := it.ValueAndErr()
+	require.NoError(t, err)
+	require.NotNil(t, val)
+	p := stop()
+	require.NotNil(t, p)
+	s := p.String()
+	require.NotEmpty(t, s)
+	t.Log(s)
+}
+
 // TestSetOptionsEquivalence tests equivalence between SetOptions to mutate an
 // iterator and constructing a new iterator with NewIter. The long-lived
 // iterator and the new iterator should surface identical iterator states.
@@ -1428,7 +1461,7 @@ func TestSetOptionsEquivalence(t *testing.T) {
 func testSetOptionsEquivalence(t *testing.T, seed uint64) {
 	rng := rand.New(rand.NewPCG(seed, seed))
 	ks := testkeys.Alpha(2)
-	d := newTestkeysDatabase(t, ks, rng)
+	d := newTestkeysDatabase(t, nil /* dbOpts */, ks, rng)
 	defer func() { require.NoError(t, d.Close()) }()
 
 	var o IterOptions
@@ -1571,12 +1604,13 @@ func iterOptionsString(o *IterOptions) string {
 	return buf.String()
 }
 
-func newTestkeysDatabase(t *testing.T, ks testkeys.Keyspace, rng *rand.Rand) *DB {
-	dbOpts := &Options{
-		Comparer: testkeys.Comparer,
-		FS:       vfs.NewMem(),
-		Logger:   panicLogger{},
+func newTestkeysDatabase(t *testing.T, dbOpts *Options, ks testkeys.Keyspace, rng *rand.Rand) *DB {
+	if dbOpts == nil {
+		dbOpts = &Options{}
 	}
+	dbOpts.Comparer = testkeys.Comparer
+	dbOpts.FS = vfs.NewMem()
+	dbOpts.Logger = panicLogger{}
 	dbOpts.randomizeForTesting(t)
 	d, err := Open("", dbOpts)
 	require.NoError(t, err)
@@ -2272,7 +2306,7 @@ func TestIteratorSeekPrefixGERandomized(t *testing.T) {
 	t.Logf("seed: %d", seed)
 	rng := rand.New(rand.NewPCG(seed, seed))
 	ks := testkeys.Alpha(2)
-	d := newTestkeysDatabase(t, ks, rng)
+	d := newTestkeysDatabase(t, nil /* dbOpts */, ks, rng)
 	defer func() { require.NoError(t, d.Close()) }()
 
 	// Scan through the keys and construct a map from unique prefix to the

metrics/value_retrieval_profile.go

@@ -0,0 +1,10 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package metrics
+
+import "github.com/cockroachdb/pebble/internal/bytesprofile"
+
+// ValueRetrievalProfile is a profile of separated value retrievals.
+type ValueRetrievalProfile = bytesprofile.Profile

scan_internal.go

@@ -201,7 +201,8 @@ func (d *DB) newInternalIter(
 	if d.iterTracker != nil {
 		dbi.trackerHandle = d.iterTracker.Start()
 	}
-	dbi.blobValueFetcher.Init(&vers.BlobFiles, d.fileCache, block.ReadEnv{},
+	dbi.blobValueFetcher.Init(&vers.BlobFiles, d.fileCache,
+		block.ReadEnv{ValueRetrievalProfile: d.valueRetrievalProfile.Load()},
 		blob.SuggestedCachedReaders(vers.MaxReadAmp()))
 
 	dbi.opts = *o

sstable/blob/fetcher.go

@@ -340,6 +340,14 @@ func (cr *cachedReader) GetUnsafeValue(
 		cr.currentValueBlock.loaded = true
 	}
 
+	// If the ReadEnv is configured with a value-retrieval profile, record the
+	// value retrieval to it. This is used to allow runtime profiling of value
+	// retrievals, providing observability into which codepaths are responsible
+	// for the comparatively expensive value retrievals.
+	if env.ValueRetrievalProfile != nil {
+		env.ValueRetrievalProfile.Record(int64(vh.ValueLen))
+	}
+
 	// Convert the ValueID to an index into the block's values. When a blob file
 	// is first constructed, the ValueID == the index. However when a blob file
 	// is rewritten, multiple blocks from the original blob file may be combined

sstable/block/block.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/bitflip"
+	"github.com/cockroachdb/pebble/internal/bytesprofile"
 	"github.com/cockroachdb/pebble/internal/cache"
 	"github.com/cockroachdb/pebble/internal/crc"
 	"github.com/cockroachdb/pebble/internal/invariants"
@@ -284,6 +285,10 @@ type ReadEnv struct {
 	// more details.
 	ReportCorruptionFn  func(base.ObjectInfo, error) error
 	ReportCorruptionArg base.ObjectInfo
+
+	// ValueRetrievalProfile, if non-nil, is used to record stack traces of
+	// retrievals of separated values. It's only used by the blob.ValueFetcher.
+	ValueRetrievalProfile *bytesprofile.Profile
 }
 
 // BlockServedFromCache updates the stats when a block was found in the cache.