db: add iterator stats for iterators below the top-level

The current IteratorStats only include information from the top
of the tree (pebble.Iterator). The new InternalIteratorStats
includes block bytes loaded (and cached), and stats for points
iterated over by the mergingIter. The latter includes points
that were ignored because of range tombstones.

These stats can be used to make inferences about root causes
of slow scans.

The stats are introduced by adding an InternalIteratorWithStats
interface that extends InternalIterator. Not all iterators need
to implement this interface (block iter, memtable iter etc.).
Flexibility of plugging in InternalIterators in various places
(especially for testing) is preserved by adding a trivial wrapper,
internalIteratorWithEmptyStats. This wrapping introduces a bug
risk that an InternalIteratorWithStats could be wrapped with an
InternalIterator and then wrapped using this trivial wrapper,
thereby losing the real stats. To reduce this risk, mergingIter
and Iterator expect to be provided an InternalIteratorWithStats
for their typical initialization paths.

Informs cockroachdb#1342
Informs cockroachdb/cockroach#59069

data_test.go

@@ -312,6 +312,18 @@ func runInternalIterCmd(d *datadriven.TestData, iter internalIterator, opts ...i
 			}
 			iter.SetBounds(lower, upper)
 			continue
+		case "stats":
+			ii, ok := iter.(internalIteratorWithStats)
+			if ok {
+				fmt.Fprintf(&b, "%+v\n", ii.Stats())
+			}
+			continue
+		case "reset-stats":
+			ii, ok := iter.(internalIteratorWithStats)
+			if ok {
+				ii.ResetStats()
+			}
+			continue
 		default:
 			return fmt.Sprintf("unknown op: %s", parts[0])
 		}

db.go

@@ -532,7 +532,7 @@ func (d *DB) getInternal(key []byte, b *Batch, s *Snapshot) ([]byte, io.Closer,
 		getIterAlloc: buf,
 		cmp:          d.cmp,
 		equal:        d.equal,
-		iter:         get,
+		iter:         base.WrapIterWithStats(get),
 		merge:        d.merge,
 		split:        d.split,
 		readState:    readState,
@@ -986,7 +986,7 @@ func finishInitializingIter(buf *iterAlloc) *Iterator {
 		// Top-level is the batch, if any.
 		if batch != nil {
 			mlevels = append(mlevels, mergingIterLevel{
-				iter:         batch.newInternalIter(&dbi.opts),
+				iter:         base.WrapIterWithStats(batch.newInternalIter(&dbi.opts)),
 				rangeDelIter: batch.newRangeDelIter(&dbi.opts),
 			})
 		}
@@ -1000,7 +1000,7 @@ func finishInitializingIter(buf *iterAlloc) *Iterator {
 				continue
 			}
 			mlevels = append(mlevels, mergingIterLevel{
-				iter:         mem.newIter(&dbi.opts),
+				iter:         base.WrapIterWithStats(mem.newIter(&dbi.opts)),
 				rangeDelIter: mem.newRangeDelIter(&dbi.opts),
 			})
 		}

external_iterator.go

@@ -112,7 +112,7 @@ func NewExternalIter(
 				return nil, err
 			}
 			mlevels = append(mlevels, mergingIterLevel{
-				iter:         pointIter,
+				iter:         base.WrapIterWithStats(pointIter),
 				rangeDelIter: rangeDelIter,
 			})
 		}

get_iter.go

@@ -37,6 +37,9 @@ type getIter struct {
 	err          error
 }
 
+// TODO(sumeer): CockroachDB code doesn't use getIter, but, for completeness,
+// make this implement InternalIteratorWithStats.
+
 // getIter implements the base.InternalIterator interface.
 var _ base.InternalIterator = (*getIter)(nil)
 

get_iter_test.go

@@ -541,7 +541,7 @@ func TestGetIter(t *testing.T) {
 			i.cmp = cmp
 			i.equal = equal
 			i.merge = DefaultMerger.Merge
-			i.iter = get
+			i.iter = base.WrapIterWithStats(get)
 
 			defer i.Close()
 			if !i.First() {

internal.go

@@ -32,3 +32,5 @@ const (
 type InternalKey = base.InternalKey
 
 type internalIterator = base.InternalIterator
+
+type internalIteratorWithStats = base.InternalIteratorWithStats

internal/base/iterator.go

@@ -213,3 +213,74 @@ type InternalIterator interface {
 
 	fmt.Stringer
 }
+
+// InternalIteratorWithStats extends InternalIterator to expose stats.
+type InternalIteratorWithStats interface {
+	InternalIterator
+	Stats() InternalIteratorStats
+	ResetStats()
+}
+
+// InternalIteratorStats contains miscellaneous stats produced by
+// InternalIterators that are part of the InternalIterator tree. Not every
+// field is relevant for an InternalIterator implementation. The field values
+// are aggregated as one goes up the InternalIterator tree.
+type InternalIteratorStats struct {
+	// Bytes in the loaded blocks. If the block was compressed, this is the
+	// compressed bytes. Currently, only the second-level index and data blocks
+	// containing points are included.
+	BlockBytes uint64
+	// Subset of BlockBytes that were in the block cache.
+	BlockBytesInCache uint64
+
+	// The following can repeatedly count the same points if they are iterated
+	// over multiple times. Additionally, they may count a point twice when
+	// switching directions. The latter could be improved if needed.
+
+	// Bytes in keys that were iterated over. Currently, only point keys are
+	// included.
+	KeyBytes uint64
+	// Bytes in values that were iterated over. Currently, only point values are
+	// included.
+	ValueBytes uint64
+	// The count of points iterated over.
+	PointCount uint64
+	// Points that were iterated over that were covered by range tombstones. It
+	// can be useful for discovering instances of
+	// https://github.com/cockroachdb/pebble/issues/1070.
+	PointsCoveredByRangeTombstones uint64
+}
+
+// Merge merges the stats in from into the given stats.
+func (s *InternalIteratorStats) Merge(from InternalIteratorStats) {
+	s.BlockBytes += from.BlockBytes
+	s.BlockBytesInCache += from.BlockBytesInCache
+	s.PointsCoveredByRangeTombstones += from.PointsCoveredByRangeTombstones
+}
+
+type internalIteratorWithEmptyStats struct {
+	InternalIterator
+}
+
+var _ InternalIteratorWithStats = internalIteratorWithEmptyStats{}
+
+// Stats implements InternalIteratorWithStats.
+func (i internalIteratorWithEmptyStats) Stats() InternalIteratorStats {
+	return InternalIteratorStats{}
+}
+
+// ResetStats implements InternalIteratorWithStats.
+func (i internalIteratorWithEmptyStats) ResetStats() {}
+
+// WrapIterWithStats ensures that either iter implements the stats methods or
+// wraps it, such that the return value implements InternalIteratorWithStats.
+func WrapIterWithStats(iter InternalIterator) InternalIteratorWithStats {
+	if iter == nil {
+		return nil
+	}
+	i, ok := iter.(InternalIteratorWithStats)
+	if ok {
+		return i
+	}
+	return internalIteratorWithEmptyStats{InternalIterator: iter}
+}

internal/rangekey/interleaving_iter.go

@@ -94,7 +94,7 @@ import (
 type InterleavingIter struct {
 	cmp                    base.Compare
 	split                  base.Split
-	pointIter              base.InternalIterator
+	pointIter              base.InternalIteratorWithStats
 	rangeKeyIter           Iterator
 	maskingThresholdSuffix []byte
 	maskSuffix             []byte
@@ -164,7 +164,7 @@ var _ base.InternalIterator = &InterleavingIter{}
 func (i *InterleavingIter) Init(
 	cmp base.Compare,
 	split base.Split,
-	pointIter base.InternalIterator,
+	pointIter base.InternalIteratorWithStats,
 	rangeKeyIter Iterator,
 	maskingThresholdSuffix []byte,
 ) {
@@ -744,6 +744,18 @@ func (i *InterleavingIter) String() string {
 	return fmt.Sprintf("range-key-interleaving(%q)", i.pointIter.String())
 }
 
+var _ base.InternalIteratorWithStats = &InterleavingIter{}
+
+// Stats implements InternalIteratorWithStats.
+func (i *InterleavingIter) Stats() base.InternalIteratorStats {
+	return i.pointIter.Stats()
+}
+
+// ResetStats implements InternalIteratorWithStats.
+func (i *InterleavingIter) ResetStats() {
+	i.pointIter.ResetStats()
+}
+
 func firstError(err0, err1 error) error {
 	if err0 != nil {
 		return err0

internal/rangekey/interleaving_iter_test.go

@@ -55,7 +55,8 @@ func runInterleavingIterTest(t *testing.T, filename string) {
 		switch td.Cmd {
 		case "set-masking-threshold":
 			maskingThreshold = []byte(strings.TrimSpace(td.Input))
-			iter.Init(cmp, testkeys.Comparer.Split, &pointIter, &rangeKeyIter, maskingThreshold)
+			iter.Init(cmp, testkeys.Comparer.Split, base.WrapIterWithStats(&pointIter), &rangeKeyIter,
+				maskingThreshold)
 			return "OK"
 		case "define-rangekeys":
 			var spans []keyspan.Span
@@ -71,7 +72,8 @@ func runInterleavingIterTest(t *testing.T, filename string) {
 				})
 			}
 			rangeKeyIter.Init(cmp, testkeys.Comparer.FormatKey, base.InternalKeySeqNumMax, keyspan.NewIter(cmp, spans))
-			iter.Init(cmp, testkeys.Comparer.Split, &pointIter, &rangeKeyIter, maskingThreshold)
+			iter.Init(cmp, testkeys.Comparer.Split, base.WrapIterWithStats(&pointIter), &rangeKeyIter,
+				maskingThreshold)
 			return "OK"
 		case "define-pointkeys":
 			var points []base.InternalKey
@@ -80,7 +82,8 @@ func runInterleavingIterTest(t *testing.T, filename string) {
 				points = append(points, base.ParseInternalKey(line))
 			}
 			pointIter = pointIterator{cmp: cmp, keys: points}
-			iter.Init(cmp, testkeys.Comparer.Split, &pointIter, &rangeKeyIter, maskingThreshold)
+			iter.Init(cmp, testkeys.Comparer.Split, base.WrapIterWithStats(&pointIter), &rangeKeyIter,
+				maskingThreshold)
 			return "OK"
 		case "iter":
 			buf.Reset()

iterator.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/fastrand"
+	"github.com/cockroachdb/pebble/internal/humanize"
 	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/internal/manifest"
 	"github.com/cockroachdb/pebble/internal/rangekey"
@@ -114,10 +115,15 @@ type IteratorStats struct {
 	ForwardStepCount [NumStatsKind]int
 	// ReverseStepCount includes Prev.
 	ReverseStepCount [NumStatsKind]int
+	InternalStats    InternalIteratorStats
 }
 
 var _ redact.SafeFormatter = &IteratorStats{}
 
+// InternalIteratorStats contains miscellaneous stats produced by internal
+// iterators.
+type InternalIteratorStats = base.InternalIteratorStats
+
 // Iterator iterates over a DB's key/value pairs in key order.
 //
 // An iterator must be closed after use, but it is not necessary to read an
@@ -143,7 +149,7 @@ type Iterator struct {
 	equal     Equal
 	merge     Merge
 	split     Split
-	iter      internalIterator
+	iter      internalIteratorWithStats
 	readState *readState
 	rangeKey  *iteratorRangeKeyState
 	err       error
@@ -1604,11 +1610,14 @@ func (i *Iterator) Metrics() IteratorMetrics {
 // ResetStats resets the stats to 0.
 func (i *Iterator) ResetStats() {
 	i.stats = IteratorStats{}
+	i.iter.ResetStats()
 }
 
 // Stats returns the current stats.
 func (i *Iterator) Stats() IteratorStats {
-	return i.stats
+	stats := i.stats
+	stats.InternalStats = i.iter.Stats()
+	return stats
 }
 
 // Clone creates a new Iterator over the same underlying data, i.e., over the
@@ -1681,4 +1690,16 @@ func (stats *IteratorStats) SafeFormat(s redact.SafePrinter, verb rune) {
 			redact.Safe(stats.ForwardSeekCount[i]), redact.Safe(stats.ForwardStepCount[i]),
 			redact.Safe(stats.ReverseSeekCount[i]), redact.Safe(stats.ReverseStepCount[i]))
 	}
+	if stats.InternalStats != (InternalIteratorStats{}) {
+		s.SafeString(",\n(internal-stats: ")
+		s.Printf("(block-bytes: (total %s, cached %s)), "+
+			"(points: (count %s, key-bytes %s, value-bytes %s, tombstoned: %s))",
+			humanize.IEC.Uint64(stats.InternalStats.BlockBytes),
+			humanize.IEC.Uint64(stats.InternalStats.BlockBytesInCache),
+			humanize.SI.Uint64(stats.InternalStats.PointCount),
+			humanize.SI.Uint64(stats.InternalStats.KeyBytes),
+			humanize.SI.Uint64(stats.InternalStats.ValueBytes),
+			humanize.SI.Uint64(stats.InternalStats.PointsCoveredByRangeTombstones),
+		)
+	}
 }

iterator_test.go

@@ -219,14 +219,14 @@ func (f *fakeIter) SetBounds(lower, upper []byte) {
 // invalidatingIter tests unsafe key/value slice reuse by modifying the last
 // returned key/value to all 1s.
 type invalidatingIter struct {
-	iter        internalIterator
+	iter        internalIteratorWithStats
 	lastKey     *InternalKey
 	lastValue   []byte
 	ignoreKinds [base.InternalKeyKindMax + 1]bool
 }
 
 func newInvalidatingIter(iter internalIterator) *invalidatingIter {
-	return &invalidatingIter{iter: iter}
+	return &invalidatingIter{iter: base.WrapIterWithStats(iter)}
 }
 
 func (i *invalidatingIter) ignoreKind(kind base.InternalKeyKind) {
@@ -314,6 +314,14 @@ func (i *invalidatingIter) String() string {
 	return i.iter.String()
 }
 
+func (i *invalidatingIter) Stats() base.InternalIteratorStats {
+	return i.iter.Stats()
+}
+
+func (i *invalidatingIter) ResetStats() {
+	i.iter.ResetStats()
+}
+
 // testIterator tests creating a combined iterator from a number of sub-
 // iterators. newFunc is a constructor function. splitFunc returns a random
 // split of the testKeyValuePairs slice such that walking a combined iterator
@@ -992,6 +1000,9 @@ func TestIteratorSeekOpt(t *testing.T) {
 			}
 			iterOutput := runIterCmd(td, iter, false)
 			stats := iter.Stats()
+			// InternalStats are non-deterministic since they depend on how data is
+			// distributed across memtables and sstables in the DB.
+			stats.InternalStats = InternalIteratorStats{}
 			var builder strings.Builder
 			fmt.Fprintf(&builder, "%sstats: %s\n", iterOutput, stats.String())
 			fmt.Fprintf(&builder, "SeekGEs with trySeekUsingNext: %d\n", seekGEUsingNext)
@@ -1106,7 +1117,7 @@ func TestIteratorSeekOptErrors(t *testing.T) {
 			equal: equal,
 			split: split,
 			merge: DefaultMerger.Merge,
-			iter:  &errorIter,
+			iter:  base.WrapIterWithStats(&errorIter),
 		}
 	}
 
@@ -1459,7 +1470,7 @@ func BenchmarkIteratorSeekGE(b *testing.B) {
 	iter := &Iterator{
 		cmp:   DefaultComparer.Compare,
 		equal: DefaultComparer.Equal,
-		iter:  m.newIter(nil),
+		iter:  base.WrapIterWithStats(m.newIter(nil)),
 	}
 	rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano())))
 
@@ -1475,7 +1486,7 @@ func BenchmarkIteratorNext(b *testing.B) {
 	iter := &Iterator{
 		cmp:   DefaultComparer.Compare,
 		equal: DefaultComparer.Equal,
-		iter:  m.newIter(nil),
+		iter:  base.WrapIterWithStats(m.newIter(nil)),
 	}
 
 	b.ResetTimer()
@@ -1492,7 +1503,7 @@ func BenchmarkIteratorPrev(b *testing.B) {
 	iter := &Iterator{
 		cmp:   DefaultComparer.Compare,
 		equal: DefaultComparer.Equal,
-		iter:  m.newIter(nil),
+		iter:  base.WrapIterWithStats(m.newIter(nil)),
 	}
 
 	b.ResetTimer()