Revert "internal/rangekey: remove ForeignSSTTransformer"

This reverts commit 02ec4a8.

internal/keyspan/transformer.go

@@ -48,3 +48,102 @@ func VisibleTransform(snapshot uint64) Transformer {
 		return nil
 	})
 }
+
+// TransformerIter is a FragmentIterator that applies a Transformer on all
+// returned keys. Used for when a caller needs to apply a transformer on an
+// iterator but does not otherwise need the mergingiter's merging ability.
+type TransformerIter struct {
+	FragmentIterator
+
+	// Transformer is applied on every Span returned by this iterator.
+	Transformer Transformer
+	// Comparer in use for this keyspace.
+	Compare base.Compare
+
+	span Span
+	err  error
+}
+
+func (t *TransformerIter) applyTransform(span *Span) *Span {
+	t.span = Span{
+		Start: t.span.Start[:0],
+		End:   t.span.End[:0],
+		Keys:  t.span.Keys[:0],
+	}
+	if err := t.Transformer.Transform(t.Compare, *span, &t.span); err != nil {
+		t.err = err
+		return nil
+	}
+	return &t.span
+}
+
+// SeekGE implements the FragmentIterator interface.
+func (t *TransformerIter) SeekGE(key []byte) *Span {
+	span := t.FragmentIterator.SeekGE(key)
+	if span == nil {
+		return nil
+	}
+	return t.applyTransform(span)
+}
+
+// SeekLT implements the FragmentIterator interface.
+func (t *TransformerIter) SeekLT(key []byte) *Span {
+	span := t.FragmentIterator.SeekLT(key)
+	if span == nil {
+		return nil
+	}
+	return t.applyTransform(span)
+}
+
+// First implements the FragmentIterator interface.
+func (t *TransformerIter) First() *Span {
+	span := t.FragmentIterator.First()
+	if span == nil {
+		return nil
+	}
+	return t.applyTransform(span)
+}
+
+// Last implements the FragmentIterator interface.
+func (t *TransformerIter) Last() *Span {
+	span := t.FragmentIterator.Last()
+	if span == nil {
+		return nil
+	}
+	return t.applyTransform(span)
+}
+
+// Next implements the FragmentIterator interface.
+func (t *TransformerIter) Next() *Span {
+	span := t.FragmentIterator.Next()
+	if span == nil {
+		return nil
+	}
+	return t.applyTransform(span)
+}
+
+// Prev implements the FragmentIterator interface.
+func (t *TransformerIter) Prev() *Span {
+	span := t.FragmentIterator.Prev()
+	if span == nil {
+		return nil
+	}
+	return t.applyTransform(span)
+}
+
+// Error implements the FragmentIterator interface.
+func (t *TransformerIter) Error() error {
+	if t.err != nil {
+		return t.err
+	}
+	return t.FragmentIterator.Error()
+}
+
+// Close implements the FragmentIterator interface.
+func (t *TransformerIter) Close() error {
+	err := t.FragmentIterator.Close()
+	if err != nil {
+		return err
+	}
+	return t.err
+}

internal/rangekey/coalesce.go

@@ -6,6 +6,7 @@ package rangekey
 
 import (
 	"bytes"
+	"fmt"
 	"math"
 	"sort"
 
@@ -378,3 +379,78 @@ func coalesce(
 	}
 	return nil
 }
+
+// ForeignSSTTransformer implements a keyspan.Transformer for range keys in
+// foreign sstables (i.e. shared sstables not created by us). It is largely
+// similar to the Transform function implemented in UserIteratorConfig in that
+// it calls coalesce to remove range keys shadowed by other range keys, but also
+// retains the range key that does the shadowing. In addition, it elides
+// RangeKey unsets/dels in L6 as they are inapplicable when reading from a
+// different Pebble instance.
+type ForeignSSTTransformer struct {
+	Comparer *base.Comparer
+	Level    int
+	sortBuf  keyspan.KeysBySuffix
+}
+
+// Transform implements the Transformer interface.
+func (f *ForeignSSTTransformer) Transform(
+	cmp base.Compare, s keyspan.Span, dst *keyspan.Span,
+) error {
+	// Apply shadowing of keys.
+	dst.Start = s.Start
+	dst.End = s.End
+	f.sortBuf = keyspan.KeysBySuffix{
+		Cmp:  cmp,
+		Keys: f.sortBuf.Keys[:0],
+	}
+	if err := coalesce(f.Comparer.Equal, &f.sortBuf, math.MaxUint64, s.Keys); err != nil {
+		return err
+	}
+	keys := f.sortBuf.Keys
+	dst.Keys = dst.Keys[:0]
+	for i := range keys {
+		seqNum := keys[i].SeqNum()
+		switch keys[i].Kind() {
+		case base.InternalKeyKindRangeKeySet:
+			if invariants.Enabled && len(dst.Keys) > 0 && cmp(dst.Keys[len(dst.Keys)-1].Suffix, keys[i].Suffix) > 0 {
+				panic("pebble: keys unexpectedly not in ascending suffix order")
+			}
+			switch f.Level {
+			case 5:
+				fallthrough
+			case 6:
+				if seqNum != base.SeqNumForLevel(f.Level) {
+					panic(fmt.Sprintf("pebble: expected range key iter to return seqnum %d, got %d", base.SeqNumForLevel(f.Level), seqNum))
+				}
+			}
+		case base.InternalKeyKindRangeKeyUnset:
+			if invariants.Enabled && len(dst.Keys) > 0 && cmp(dst.Keys[len(dst.Keys)-1].Suffix, keys[i].Suffix) > 0 {
+				panic("pebble: keys unexpectedly not in ascending suffix order")
+			}
+			fallthrough
+		case base.InternalKeyKindRangeKeyDelete:
+			switch f.Level {
+			case 5:
+				// Emit this key.
+				if seqNum != base.SeqNumForLevel(f.Level) {
+					panic(fmt.Sprintf("pebble: expected range key iter to return seqnum %d, got %d", base.SeqNumForLevel(f.Level), seqNum))
+				}
+			case 6:
+				// Skip this key, as foreign sstable in L6 do not need to emit range key
+				// unsets/dels as they do not apply to any other sstables.
+				continue
+			}
+		default:
+			return base.CorruptionErrorf("pebble: unrecognized range key kind %s", keys[i].Kind())
+		}
+		dst.Keys = append(dst.Keys, keyspan.Key{
+			Trailer: base.MakeTrailer(seqNum, keys[i].Kind()),
+			Suffix:  keys[i].Suffix,
+			Value:   keys[i].Value,
+		})
+	}
+	// coalesce results in dst.Keys being sorted by Suffix.
+	dst.KeysOrder = keyspan.BySuffixAsc
+	return nil
+}

table_cache.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/pebble/internal/keyspan"
 	"github.com/cockroachdb/pebble/internal/manifest"
 	"github.com/cockroachdb/pebble/internal/private"
+	"github.com/cockroachdb/pebble/internal/rangekey"
 	"github.com/cockroachdb/pebble/objstorage"
 	"github.com/cockroachdb/pebble/objstorage/objstorageprovider/objiotracing"
 	"github.com/cockroachdb/pebble/sstable"
@@ -649,6 +650,29 @@ func (c *tableCacheShard) newRangeKeyIter(
 		return emptyKeyspanIter, nil
 	}
 
+	objMeta, err := dbOpts.objProvider.Lookup(fileTypeTable, file.FileBacking.DiskFileNum)
+	if err != nil {
+		return nil, err
+	}
+	if dbOpts.objProvider.IsForeign(objMeta) {
+		if opts.Level == 0 {
+			panic("unexpected zero level when reading foreign file")
+		}
+		transform := &rangekey.ForeignSSTTransformer{
+			Comparer: dbOpts.opts.Comparer,
+			Level:    manifest.LevelToInt(opts.Level),
+		}
+		if iter == nil {
+			iter = emptyKeyspanIter
+		}
+		transformIter := &keyspan.TransformerIter{
+			FragmentIterator: iter,
+			Transformer:      transform,
+			Compare:          dbOpts.opts.Comparer.Compare,
+		}
+		return transformIter, nil
+	}
+
 	return iter, nil
 }