overlap: rework ingest LSM overlap check code

Instead of using level iterators, we directly check overlap with files
individually. This requires less machinery and there is no chance of
opening files that are outside of the bounds of interest.

We put the new code in a new `internal/overlap` package and
`overlapChecker` is now a very small wrapper which interfaces with
`overlap.Checker`.

Author: RaduBerinde

db.go

@@ -3007,7 +3007,13 @@ func (d *DB) checkVirtualBounds(m *fileMetadata) {
 
 // DebugString returns a debugging string describing the LSM.
 func (d *DB) DebugString() string {
+	return d.DebugCurrentVersion().DebugString()
+}
+
+// DebugCurrentVersion returns the current LSM tree metadata. Should only be
+// used for testing/debugging.
+func (d *DB) DebugCurrentVersion() *manifest.Version {
 	d.mu.Lock()
 	defer d.mu.Unlock()
-	return d.mu.versions.currentVersion().DebugString()
+	return d.mu.versions.currentVersion()
 }

flushable.go

@@ -394,3 +394,33 @@ func determineOverlapAllIters(
 	}
 	return false, nil
 }
+
+func determineOverlapPointIterator(
+	cmp base.Compare, bounds base.UserKeyBounds, iter internalIterator,
+) (bool, error) {
+	kv := iter.SeekGE(bounds.Start, base.SeekGEFlagsNone)
+	if kv == nil {
+		return false, iter.Error()
+	}
+	return bounds.End.IsUpperBoundForInternalKey(cmp, kv.K), nil
+}
+
+func determineOverlapKeyspanIterator(
+	cmp base.Compare, bounds base.UserKeyBounds, iter keyspan.FragmentIterator,
+) (bool, error) {
+	// NB: The spans surfaced by the fragment iterator are non-overlapping.
+	span, err := iter.SeekGE(bounds.Start)
+	if err != nil {
+		return false, err
+	}
+	for ; span != nil; span, err = iter.Next() {
+		if !bounds.End.IsUpperBoundFor(cmp, span.Start) {
+			// The span starts after our bounds.
+			return false, nil
+		}
+		if !span.Empty() {
+			return true, nil
+		}
+	}
+	return false, err
+}

ingest.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/internal/manifest"
+	"github.com/cockroachdb/pebble/internal/overlap"
 	"github.com/cockroachdb/pebble/internal/private"
 	"github.com/cockroachdb/pebble/objstorage"
 	"github.com/cockroachdb/pebble/objstorage/remote"
@@ -832,7 +833,7 @@ func ingestUpdateSeqNum(
 func ingestTargetLevel(
 	ctx context.Context,
 	cmp base.Compare,
-	lsmOverlap lsmOverlap,
+	lsmOverlap overlap.WithLSM,
 	baseLevel int,
 	compactions map[*compaction]struct{},
 	meta *fileMetadata,
@@ -904,32 +905,32 @@ func ingestTargetLevel(
 	// existing point that falls within the ingested table bounds as being "data
 	// overlap".
 
-	if lsmOverlap[0].result == dataOverlap {
+	if lsmOverlap[0].Result == overlap.Data {
 		return 0, nil, nil
 	}
 	targetLevel = 0
 	splitFile = nil
 	for level := baseLevel; level < numLevels; level++ {
 		var candidateSplitFile *fileMetadata
-		switch lsmOverlap[level].result {
-		case dataOverlap:
+		switch lsmOverlap[level].Result {
+		case overlap.Data:
 			// We cannot ingest into or under this level; return the best target level
 			// so far.
 			return targetLevel, splitFile, nil
 
-		case noDataOverlap:
-			if !suggestSplit || lsmOverlap[level].splitFile == nil {
+		case overlap.OnlyBoundary:
+			if !suggestSplit || lsmOverlap[level].SplitFile == nil {
 				// We can ingest under this level, but not into this level.
 				continue
 			}
 			// We can ingest into this level if we split this file.
-			candidateSplitFile = lsmOverlap[level].splitFile
+			candidateSplitFile = lsmOverlap[level].SplitFile
 
-		case noBoundaryOverlap:
+		case overlap.None:
 		// We can ingest into this level.
 
 		default:
-			return 0, nil, base.AssertionFailedf("unexpected lsmOverlap result: %v", lsmOverlap[level].result)
+			return 0, nil, base.AssertionFailedf("unexpected WithLevel.Result: %v", lsmOverlap[level].Result)
 		}
 
 		// Check boundary overlap with any ongoing compactions. We consider an
@@ -2192,15 +2193,15 @@ func (d *DB) ingestApply(
 				// We check overlap against the LSM without holding DB.mu. Note that we
 				// are still holding the log lock, so the version cannot change.
 				// TODO(radu): perform this check optimistically outside of the log lock.
-				var overlap lsmOverlap
-				overlap, err = func() (lsmOverlap, error) {
+				var lsmOverlap overlap.WithLSM
+				lsmOverlap, err = func() (overlap.WithLSM, error) {
 					d.mu.Unlock()
 					defer d.mu.Lock()
 					return overlapChecker.DetermineLSMOverlap(ctx, m.UserKeyBounds())
 				}()
 				if err == nil {
 					f.Level, splitFile, err = ingestTargetLevel(
-						ctx, d.cmp, overlap, baseLevel, d.mu.compact.inProgress, m, shouldIngestSplit,
+						ctx, d.cmp, lsmOverlap, baseLevel, d.mu.compact.inProgress, m, shouldIngestSplit,
 					)
 				}
 			}

internal/base/key_bounds.go

@@ -73,6 +73,23 @@ func (eb UserKeyBoundary) IsUpperBoundForInternalKey(cmp Compare, key InternalKe
 	return c < 0 || (c == 0 && (eb.Kind == Inclusive || key.IsExclusiveSentinel()))
 }
 
+// CompareUpperBounds compares two UserKeyBoundaries as upper bounds (e.g. when
+// they are used for UserKeyBounds.End).
+func (eb UserKeyBoundary) CompareUpperBounds(cmp Compare, other UserKeyBoundary) int {
+	switch c := cmp(eb.Key, other.Key); {
+	case c != 0:
+		return c
+	case eb.Kind == other.Kind:
+		return 0
+	case eb.Kind == Inclusive:
+		// eb is inclusive, other is exclusive.
+		return 1
+	default:
+		// eb is exclusive, other is inclusive.
+		return -1
+	}
+}
+
 // UserKeyBounds is a user key interval with an inclusive start boundary and
 // with an end boundary that can be either inclusive or exclusive.
 type UserKeyBounds struct {
@@ -132,8 +149,7 @@ func (b *UserKeyBounds) ContainsBounds(cmp Compare, other *UserKeyBounds) bool {
 	if cmp(b.Start, other.Start) > 0 {
 		return false
 	}
-	c := cmp(other.End.Key, b.End.Key)
-	return c < 0 || (c == 0 && (b.End.Kind == Inclusive || other.End.Kind == Exclusive))
+	return other.End.CompareUpperBounds(cmp, b.End) <= 0
 }
 
 // ContainsUserKey returns true if the user key is within the bounds.

internal/base/key_bounds_test.go

@@ -10,6 +10,42 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestUserKeyBoundary(t *testing.T) {
+	cmp := DefaultComparer.Compare
+	a := []byte("a")
+	b := []byte("b")
+
+	aExclusive := UserKeyExclusive(a)
+	require.False(t, aExclusive.IsUpperBoundFor(cmp, a))
+	require.False(t, aExclusive.IsUpperBoundFor(cmp, b))
+
+	aInclusive := UserKeyInclusive(a)
+	require.True(t, aInclusive.IsUpperBoundFor(cmp, a))
+	require.False(t, aInclusive.IsUpperBoundFor(cmp, b))
+
+	bExclusive := UserKeyExclusive(b)
+	require.True(t, bExclusive.IsUpperBoundFor(cmp, a))
+	require.False(t, bExclusive.IsUpperBoundFor(cmp, b))
+
+	bInclusive := UserKeyInclusive(b)
+	require.True(t, bInclusive.IsUpperBoundFor(cmp, a))
+	require.True(t, bInclusive.IsUpperBoundFor(cmp, b))
+
+	ordered := []UserKeyBoundary{aExclusive, aInclusive, bExclusive, bInclusive}
+	for i := range ordered {
+		for j := range ordered {
+			expected := 0
+			if i < j {
+				expected = -1
+			} else if i > j {
+				expected = 1
+			}
+			require.Equalf(t, expected, ordered[i].CompareUpperBounds(cmp, ordered[j]),
+				"%v, %v", ordered[i], ordered[j])
+		}
+	}
+}
+
 func TestUserKeyBounds(t *testing.T) {
 	var ukb UserKeyBounds
 	cmp := DefaultComparer.Compare

internal/base/test_utils.go

@@ -6,6 +6,7 @@ package base
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"strconv"
 	"strings"
@@ -249,3 +250,14 @@ func (f *FakeIter) SetBounds(lower, upper []byte) {
 
 // SetContext is part of the InternalIterator interface.
 func (f *FakeIter) SetContext(_ context.Context) {}
+
+// ParseUserKeyBounds parses UserKeyBounds from a string representation of the
+// form "[foo, bar]" or "[foo, bar)".
+func ParseUserKeyBounds(s string) UserKeyBounds {
+	first, last, s := s[0], s[len(s)-1], s[1:len(s)-1]
+	start, end, ok := strings.Cut(s, ", ")
+	if !ok || first != '[' || (last != ']' && last != ')') {
+		panic(fmt.Sprintf("invalid bounds %q", s))
+	}
+	return UserKeyBoundsEndExclusiveIf([]byte(start), []byte(end), last == ')')
+}

internal/manifest/l0_sublevels_test.go

@@ -289,7 +289,7 @@ func TestL0Sublevels(t *testing.T) {
 				SortBySmallest(fileMetas[i], base.DefaultComparer.Compare)
 			}
 
-			levelMetadata := makeLevelMetadata(base.DefaultComparer.Compare, 0, fileMetas[0])
+			levelMetadata := MakeLevelMetadata(base.DefaultComparer.Compare, 0, fileMetas[0])
 			if initialize {
 				if addL0FilesOpt {
 					SortBySeqNum(addedL0Files)
@@ -550,7 +550,7 @@ func TestAddL0FilesEquivalence(t *testing.T) {
 			continue
 		}
 
-		levelMetadata := makeLevelMetadata(testkeys.Comparer.Compare, 0, fileMetas)
+		levelMetadata := MakeLevelMetadata(testkeys.Comparer.Compare, 0, fileMetas)
 		var err error
 
 		if s2 == nil {

internal/manifest/level_metadata.go

@@ -41,7 +41,8 @@ func (lm *LevelMetadata) release() (obsolete []*FileBacking) {
 	return lm.tree.Release()
 }
 
-func makeLevelMetadata(cmp Compare, level int, files []*FileMetadata) LevelMetadata {
+// MakeLevelMetadata creates a LevelMetadata with the given files.
+func MakeLevelMetadata(cmp Compare, level int, files []*FileMetadata) LevelMetadata {
 	bcmp := btreeCmpSeqNum
 	if level > 0 {
 		bcmp = btreeCmpSmallestKey(cmp)

internal/manifest/version.go

@@ -301,6 +301,21 @@ func (m *FileMetadata) UserKeyBounds() base.UserKeyBounds {
 	return base.UserKeyBoundsFromInternal(m.Smallest, m.Largest)
 }
 
+// UserKeyBoundsByType returns the user key bounds for the given key types.
+// Note that the returned bounds are invalid when requesting KeyTypePoint but
+// HasPointKeys is false, or when requesting KeyTypeRange and HasRangeKeys is
+// false.
+func (m *FileMetadata) UserKeyBoundsByType(keyType KeyType) base.UserKeyBounds {
+	switch keyType {
+	case KeyTypePoint:
+		return base.UserKeyBoundsFromInternal(m.SmallestPointKey, m.LargestPointKey)
+	case KeyTypeRange:
+		return base.UserKeyBoundsFromInternal(m.SmallestRangeKey, m.LargestRangeKey)
+	default:
+		return base.UserKeyBoundsFromInternal(m.Smallest, m.Largest)
+	}
+}
+
 // SyntheticSeqNum returns a SyntheticSeqNum which is set when SmallestSeqNum
 // equals LargestSeqNum.
 func (m *FileMetadata) SyntheticSeqNum() sstable.SyntheticSeqNum {

internal/manifest/version_edit.go

@@ -963,12 +963,12 @@ func (b *BulkVersionEdit) Apply(
 
 	for level := range v.Levels {
 		if curr == nil || curr.Levels[level].tree.root == nil {
-			v.Levels[level] = makeLevelMetadata(comparer.Compare, level, nil /* files */)
+			v.Levels[level] = MakeLevelMetadata(comparer.Compare, level, nil /* files */)
 		} else {
 			v.Levels[level] = curr.Levels[level].clone()
 		}
 		if curr == nil || curr.RangeKeyLevels[level].tree.root == nil {
-			v.RangeKeyLevels[level] = makeLevelMetadata(comparer.Compare, level, nil /* files */)
+			v.RangeKeyLevels[level] = MakeLevelMetadata(comparer.Compare, level, nil /* files */)
 		} else {
 			v.RangeKeyLevels[level] = curr.RangeKeyLevels[level].clone()
 		}