db: add probe tests to TestIterHistories

This commit adds probe tests to iterator_histories, also adds unit tests
for the new optimization.

Author: Sachuman

file_cache.go

@@ -732,7 +732,7 @@ func (h *fileCacheHandle) newPointIter(
 				ValueFetcher: internalOpts.blobValueFetcher,
 				References:   blobReferences,
 			},
-			MaximumSuffixProperty: sstable.MaxTestKeysSuffixProperty{},
+			MaximumSuffixProperty: opts.GetMaximumSuffixProperty(),
 		})
 	}
 	if err != nil {

internal/base/internal.go

@@ -96,8 +96,7 @@ const (
 	//InternalKeyKindRollbackXID              InternalKeyKind = 12
 	//InternalKeyKindNoop                     InternalKeyKind = 13
 	//InternalKeyKindColumnFamilyRangeDelete  InternalKeyKind = 14
-	InternalKeyKindSyntheticKey InternalKeyKind = 14
-	InternalKeyKindRangeDelete  InternalKeyKind = 15
+	InternalKeyKindRangeDelete InternalKeyKind = 15
 	//InternalKeyKindColumnFamilyBlobIndex    InternalKeyKind = 16
 	//InternalKeyKindBlobIndex                InternalKeyKind = 17
 
@@ -148,6 +147,9 @@ const (
 	// appear amongst other key kinds in a batch (with the exception of alongside
 	// InternalKeyKindIngestSST), or in an sstable.
 	InternalKeyKindExcise InternalKeyKind = 24
+	// InternalKeyKindSyntheticKey is a key used to mark synthetic keys in the
+	// sstable. This is used to perform optimization during SeekPrefixGE.
+	InternalKeyKindSyntheticKey InternalKeyKind = 25
 
 	// This maximum value isn't part of the file format. Future extensions may
 	// increase this value.
@@ -158,7 +160,7 @@ const (
 	// which sorts 'less than or equal to' any other valid internalKeyKind, when
 	// searching for any kind of internal key formed by a certain user key and
 	// seqNum.
-	InternalKeyKindMax InternalKeyKind = 24
+	InternalKeyKindMax InternalKeyKind = 25
 
 	// InternalKeyKindMaxForSSTable is the largest valid key kind that can exist
 	// in an SSTable. This should usually equal InternalKeyKindMax, except
@@ -201,7 +203,6 @@ var internalKeyKindNames = []string{
 	InternalKeyKindMerge:          "MERGE",
 	InternalKeyKindLogData:        "LOGDATA",
 	InternalKeyKindSingleDelete:   "SINGLEDEL",
-	InternalKeyKindSyntheticKey:   "SYNTHETIC",
 	InternalKeyKindRangeDelete:    "RANGEDEL",
 	InternalKeyKindSeparator:      "SEPARATOR",
 	InternalKeyKindSetWithDelete:  "SETWITHDEL",
@@ -211,6 +212,7 @@ var internalKeyKindNames = []string{
 	InternalKeyKindIngestSST:      "INGESTSST",
 	InternalKeyKindDeleteSized:    "DELSIZED",
 	InternalKeyKindExcise:         "EXCISE",
+	InternalKeyKindSyntheticKey:   "SYNTHETIC",
 	InternalKeyKindInvalid:        "INVALID",
 }
 
@@ -309,7 +311,6 @@ func MakeExclusiveSentinelKey(kind InternalKeyKind, userKey []byte) InternalKey
 var kindsMap = map[string]InternalKeyKind{
 	"DEL":           InternalKeyKindDelete,
 	"SINGLEDEL":     InternalKeyKindSingleDelete,
-	"SYNTHETIC":     InternalKeyKindSyntheticKey,
 	"RANGEDEL":      InternalKeyKindRangeDelete,
 	"LOGDATA":       InternalKeyKindLogData,
 	"SET":           InternalKeyKindSet,
@@ -323,6 +324,7 @@ var kindsMap = map[string]InternalKeyKind{
 	"INGESTSST":     InternalKeyKindIngestSST,
 	"DELSIZED":      InternalKeyKindDeleteSized,
 	"EXCISE":        InternalKeyKindExcise,
+	"SYNTHETIC":     InternalKeyKindSyntheticKey,
 }
 
 // ParseSeqNum parses the string representation of a sequence number.

internal/base/internal_test.go

@@ -41,7 +41,7 @@ func TestInvalidInternalKey(t *testing.T) {
 		"\x01\x02\x03\x04\x05\x06\x07",
 		"foo",
 		"foo\x08\x07\x06\x05\x04\x03\x02",
-		"foo\x19\x07\x06\x05\x04\x03\x02\x01",
+		"foo\x20\x07\x06\x05\x04\x03\x02\x01",
 	}
 	for _, tc := range testCases {
 		k := DecodeInternalKey([]byte(tc))

internal/itertest/testdata/probes

@@ -21,6 +21,33 @@ bar#2,SET:value
 bar#2,SET:value
 bar#2,SET:value
 
+# Create probes to exercise the UserKey predicate. On 2 SeekGE Index return a KV
+# with user key "bar"; on other ops return a KV with user key "foo". Then if
+# the returned KV's user key is "bar", override the result with a specific KV.
+# In this case "ok"
+
+new
+(If (And OpSeekGE (OnIndex 2)) (ReturnKV "bar.SET.2" "value") (ReturnKV "foo.SET.1" "value"))
+(If (UserKey "bar") (ReturnKV "ok.SET.1" "ok") noop)
+----
+
+iter
+first
+next
+seek-ge foo
+prev
+seek-ge foo
+next
+seek-ge foo
+----
+foo#1,SET:value
+foo#1,SET:value
+foo#1,SET:value
+foo#1,SET:value
+foo#1,SET:value
+foo#1,SET:value
+ok#1,SET:ok
+
 # Create an identical probe as above, and then wrap it in a "noop" probe.
 # A noop probe does nothing; it passes through the child iterator's results
 # verbatim, so the iteration results should be the same.

iterator_histories_test.go

@@ -5,6 +5,7 @@ package pebble
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"path/filepath"
 	"strconv"
@@ -15,6 +16,8 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/invariants"
+	"github.com/cockroachdb/pebble/internal/itertest"
+	"github.com/cockroachdb/pebble/internal/manifest"
 	"github.com/cockroachdb/pebble/internal/testkeys"
 	"github.com/cockroachdb/pebble/internal/testutils"
 	"github.com/cockroachdb/pebble/sstable"
@@ -38,6 +41,7 @@ func TestIterHistories(t *testing.T) {
 		var buf bytes.Buffer
 		iters := map[string]*Iterator{}
 		batches := map[string]*Batch{}
+		parser := itertest.NewParser()
 		newIter := func(name string, reader Reader, o *IterOptions) *Iterator {
 			it, _ := reader.NewIter(o)
 			iters[name] = it
@@ -285,9 +289,11 @@ func TestIterHistories(t *testing.T) {
 				delete(batches, name)
 				return fmt.Sprintf("committed %d keys\n", count)
 			case "combined-iter":
-				o := &IterOptions{KeyTypes: IterKeyTypePointsAndRanges}
+				o := &IterOptions{KeyTypes: IterKeyTypePointsAndRanges, MaximumSuffixProperty: sstable.MaxTestKeysSuffixProperty{}}
 				var reader Reader = d
 				var name string
+				pointProbes := make(map[base.TableNum][]itertest.Probe)
+				rangeDelProbes := make(map[base.TableNum][]keyspanProbe)
 				for _, arg := range td.CmdArgs {
 					switch arg.Key {
 					case "mask-suffix":
@@ -348,6 +354,40 @@ func TestIterHistories(t *testing.T) {
 						}()
 					case "use-l6-filter":
 						o.UseL6Filters = true
+					case "probe-points":
+						fileNumStr := arg.Vals[0]
+						i, err := strconv.Atoi(fileNumStr)
+						if err != nil {
+							require.NoError(t, err)
+						}
+						pointProbes[base.TableNum(i)] = itertest.MustParseProbes(parser, arg.Vals[1:]...)
+					case "probe-rangedels":
+						fileNumStr := arg.Vals[0]
+						i, err := strconv.Atoi(fileNumStr)
+						if err != nil {
+							require.NoError(t, err)
+						}
+						rangeDelProbes[base.TableNum(i)] = parseKeyspanProbes(arg.Vals[1:]...)
+					}
+				}
+				// If we have probes, override d.newIters to attach them
+				if (len(pointProbes) > 0 || len(rangeDelProbes) > 0) && reader == d {
+					oldNewIters := d.newIters
+					d.newIters = func(
+						ctx context.Context, file *manifest.TableMetadata, iopts *IterOptions,
+						iio internalIterOpts, kinds iterKinds,
+					) (iterSet, error) {
+						set, err := oldNewIters(ctx, file, iopts, iio, kinds)
+						if err != nil {
+							return set, err
+						}
+						if probes := pointProbes[file.TableNum]; len(probes) > 0 {
+							set.point = itertest.Attach(set.point, itertest.ProbeState{Comparer: testkeys.Comparer, Log: &buf}, probes...)
+						}
+						if rangedelProbes := rangeDelProbes[file.TableNum]; len(rangedelProbes) > 0 {
+							set.rangeDeletion = attachKeyspanProbes(set.rangeDeletion, keyspanProbeContext{log: &buf}, rangedelProbes...)
+						}
+						return set, nil
 					}
 				}
 				var iter *Iterator
@@ -370,7 +410,12 @@ func TestIterHistories(t *testing.T) {
 				if err != nil {
 					return err.Error()
 				}
-				return runIterCmd(td, iter, name == "" /* close iter */)
+				out := runIterCmd(td, iter, name == "" /* close iter */)
+				// Append probe logs if any
+				if len(pointProbes) > 0 && buf.Len() > 0 {
+					out += buf.String()
+				}
+				return out
 			case "rangekey-iter":
 				name := pluckStringCmdArg(td, "name")
 				iter := newIter(name, d, &IterOptions{KeyTypes: IterKeyTypeRangesOnly})

level_iter.go

@@ -154,6 +154,7 @@ func (l *levelIter) init(
 	if len(opts.PointKeyFilters) == 0 {
 		l.tableOpts.PointKeyFilters = l.filtersBuf[:0:1]
 	}
+	l.tableOpts.MaximumSuffixProperty = opts.MaximumSuffixProperty
 	l.tableOpts.UseL6Filters = opts.UseL6Filters
 	l.tableOpts.Category = opts.Category
 	l.tableOpts.layer = l.layer

options.go

@@ -67,6 +67,9 @@ type BlockPropertyCollector = sstable.BlockPropertyCollector
 // BlockPropertyFilter exports the sstable.BlockPropertyFilter type.
 type BlockPropertyFilter = base.BlockPropertyFilter
 
+// MaximumSuffixProperty exports the sstable.MaximumSuffixProperty type.
+type MaximumSuffixProperty = sstable.MaximumSuffixProperty
+
 // ShortAttributeExtractor exports the base.ShortAttributeExtractor type.
 type ShortAttributeExtractor = base.ShortAttributeExtractor
 
@@ -218,6 +221,9 @@ type IterOptions struct {
 	// implies obsolete points should not be hidden.
 	snapshotForHideObsoletePoints base.SeqNum
 
+	// MaximumSuffixProperty is the maximum suffix property for the iterator.
+	// This is used to perform the synthetic key optimization.
+	MaximumSuffixProperty MaximumSuffixProperty
 	// NB: If adding new Options, you must account for them in iterator
 	// construction and Iterator.SetOptions.
 }
@@ -238,6 +244,14 @@ func (o *IterOptions) GetUpperBound() []byte {
 	return o.UpperBound
 }
 
+// GetMaximumSuffixProperty returns the MaximumSuffixProperty.
+func (o *IterOptions) GetMaximumSuffixProperty() MaximumSuffixProperty {
+	if o == nil {
+		return nil
+	}
+	return o.MaximumSuffixProperty
+}
+
 func (o *IterOptions) pointKeys() bool {
 	if o == nil {
 		return true

sstable/block_property_test_utils.go

@@ -141,12 +141,15 @@ func testKeysSuffixToInterval(suffix []byte) BlockInterval {
 
 type MaxTestKeysSuffixProperty struct{}
 
+// Name is part of the MaxTestKeysSuffixProperty interface.
 func (testprop MaxTestKeysSuffixProperty) Name() string {
 	return `pebble.internal.testkeys.suffixes`
 }
 
+// Extract is part of the MaxTestKeysSuffixProperty interface used to extract the
+// latest suffix from the block property.
 func (testprop MaxTestKeysSuffixProperty) Extract(
-	encodedProperty []byte,
+	dst []byte, encodedProperty []byte,
 ) (suffix []byte, ok bool, err error) {
 	if len(encodedProperty) <= 1 {
 		return nil, false, nil
@@ -158,6 +161,7 @@ func (testprop MaxTestKeysSuffixProperty) Extract(
 	} else if interval.IsEmpty() {
 		return nil, false, nil
 	}
-	ret := strconv.AppendUint([]byte("@"), (interval.Upper - 1), 10)
-	return ret, true, nil
+	dst = append(dst, '@')
+	dst = strconv.AppendUint(dst, (interval.Upper - 1), 10)
+	return dst, true, nil
 }

sstable/reader.go

@@ -78,7 +78,8 @@ type ReadEnv struct {
 	// only be set when Virtual is non-nil.
 	IsSharedIngested bool
 	Block            block.ReadEnv
-	InternalBounds   *base.InternalKeyBounds
+	// InternalBounds is the bounds of the sstable. This is currently used during synthetic key optimization.
+	InternalBounds *base.InternalKeyBounds
 }
 
 var NoReadEnv = ReadEnv{}
@@ -106,15 +107,23 @@ type IterOptions struct {
 	MaximumSuffixProperty MaximumSuffixProperty
 }
 
+// MaximumSuffixProperty is an interface for the maximum suffix property.
+// This is used to perform the synthetic key optimization.
 type MaximumSuffixProperty interface {
+	// Name returns the name of the maximum suffix property.
 	Name() string
-	Extract(encodedProperty []byte) (suffix []byte, ok bool, err error)
+	// Extract extracts the maximum suffix from the encoded property and
+	// appends it to the provided dst slice, if it exists. If no such property is found,
+	// Extract returns ok=false. The Extract method must NOT mutate the encoded property.
+	Extract(dst []byte, encodedProperty []byte) (suffix []byte, ok bool, err error)
 }
 
-type SyntheticKey struct {
+// syntheticKey is a struct that contains the synthetic key and the seek key.
+type syntheticKey struct {
 	kv             base.InternalKV
 	seekKey        []byte
 	atSyntheticKey bool
+	maxSuffixBuf   []byte
 }
 
 // NewPointIter returns an iterator for the point keys in the table.