perf: seek past the internal-key region instead of stepping through it (#1243)

### Motivation

Item 2.2 from the performance review. Floor/ceiling gets and unbounded
scans that exclude internal keys used pebble's `SkipPoint`. `SkipPoint`
is a **filter, not a seek**: the iterator still steps through — and
loads the blocks of — every internal key it passes. A `CEILING`/`HIGHER`
get past the last user key, or a scan with no upper bound, therefore
walks the **entire live notification backlog** before returning. With
100k notifications retained a single client `Get` took **3 ms**, and the
cost grows without bound with the retention window.

### The subtlety (why the obvious fix is wrong)

The internal keys are contiguous, so the tempting fix is to clamp the
iterator's `UpperBound` to the start of the internal region. That is
**unsafe for natural sorting**: a regular key is stored raw, so one
whose bytes sort after the encoded internal prefix `"\xff\xffoxia/"` —
e.g. `"\xff\xffz"` — lives *after* the internal keys. Bounding the
region away would silently hide those keys. Hierarchical sorting is
different: the marker bit puts internal keys strictly last, so a bound
*is* safe there.

### Changes

`Encoder` gains `InternalKeyRange() (start, end []byte)`, returning the
contiguous internal-key region as a half-open range:

- **hierarchical** → `end == nil` (region runs to the end of the
keyspace): pruned outright with an `UpperBound`, so pebble never visits
it.
- **natural** → real `end`: the region is **seeked over** in a single
jump (it is contiguous, so nothing at/after `end` is internal), in
**both directions** — the reverse paths had the same walk, a floor probe
above the region stepped *backwards* through the whole backlog.

`SkipPoint` is removed on these paths. Besides causing the walk, it
hides the internal keys from the wrapper, so we could not detect the
region in order to seek past it — the two approaches are mutually
exclusive. The skipper is applied at every positioning op (the four
point lookups plus `Next`/`Prev`/`SeekGE`/`SeekLT` in both iterator
wrappers) so no internal key can leak. `InternalKeyRange` is
allocation-free for both encoders (constant, read-only bounds shared as
package vars).

### Benchmark

`BenchmarkPebbleGetCeilingPastUserKeys` — ceiling lookup past the last
user key, N internal keys retained:

| backlog | before (SkipPoint) | after (seek) |
|---|---|---|
| 100 | 4.1 µs | 1.23 µs |
| 10,000 | 295 µs | 1.26 µs |
| 100,000 | **3.05 ms** | **1.27 µs** |

O(N) → O(1).

### Verification

- `TestPebbleScanSkipsInternalRegionButKeepsKeysAfterIt` and
`TestPebbleGetAcrossInternalRegionNatural` plant a regular key **after**
the internal region under natural sorting and assert scans/gets in both
directions still return it. Verified to discriminate: they fail against
an upper-bound prune (`end == nil` for natural) — exactly the mistake
the naive fix would make.
- `TestPebbleGetPastInternalRegionHierarchical` pins that hierarchical
ceilings past the user keys are genuinely `NOT_FOUND` without walking
the region.
- `go test -race` green on `oxiad/dataserver/...` and `common/compare`;
`golangci-lint` clean across all go.work modules on the CI-pinned v2.6.2
(the change spans `common/` and `oxiad/`).

---

Stacked on #1242 (drops a redundant seek in `getHigher`, which this PR
also rewrites); please merge that first. Until it does, its one commit
shows in this PR's diff.

---------

Signed-off-by: Matteo Merli <mmerli@apache.org>

Author: merlimat

common/compare/encode.go

@@ -30,11 +30,39 @@ type Encoder interface {
 	Decode(encodedKey []byte) string
 
 	IsInternalKey(encodedKey []byte) bool
+
+	// InternalKeyRange returns the encoded half-open range [start, end) that
+	// contains every internal key. The region is always contiguous, so an
+	// iterator can seek past it in one jump instead of stepping through it.
+	//
+	// end is nil when the region runs to the end of the keyspace — no user key
+	// can sort after it — which lets the region be pruned with an upper bound
+	// instead. That is the case for the hierarchical encoder, but NOT for the
+	// natural one, where a user key is stored raw and can therefore sort after
+	// the internal keys.
+	InternalKeyRange() (start []byte, end []byte)
+}
+
+// keyRegionSuccessor returns the smallest key that sorts after every key with
+// the given prefix, or nil when no such key exists (an all-0xff prefix).
+func keyRegionSuccessor(prefix []byte) []byte {
+	end := bytes.Clone(prefix)
+	for i := len(end) - 1; i >= 0; i-- {
+		if end[i] != maxByteValue {
+			end[i]++
+			return end[:i+1]
+		}
+	}
+	return nil
 }
 
 const (
 	encodedSeparator      = 0xff
 	internalKeysBitMarker = 1 << 15
+
+	// maxByteValue is the largest byte value; a prefix made only of these has
+	// no successor.
+	maxByteValue = 0xff
 )
 
 var (
@@ -105,6 +133,19 @@ func (encoderHierarchical) IsInternalKey(encodedKey []byte) bool {
 	return encodedKey[0]&(1<<7) != 0
 }
 
+// InternalKeyRange: the internal-key marker is the top bit of the 2-byte
+// prefix, so every internal key sorts after every regular one (a regular key
+// would need 32768 separators to reach that region) — the region runs to the
+// end of the keyspace.
+func (encoderHierarchical) InternalKeyRange() (start []byte, end []byte) {
+	return hierarchicalInternalKeyStart, nil
+}
+
+// The bounds returned by InternalKeyRange are constant and read-only, so build
+// them once. Callers use them as pebble iterator bounds and seek keys, never
+// mutating them (same contract as encodedInternalKeyPrefixBytes).
+var hierarchicalInternalKeyStart = []byte{internalKeysBitMarker >> 8}
+
 // EncoderHierarchical ensure that we can sort keys from same level together
 // and thus we can easily return the children of a given path
 // The encoding is done by prepending 2 bytes with the count of
@@ -163,6 +204,16 @@ func (encoderNatural) IsInternalKey(encodedKey []byte) bool {
 	return bytes.HasPrefix(encodedKey, encodedInternalKeyPrefixBytes)
 }
 
+// InternalKeyRange: internal keys are the "\xff\xffoxia/" prefix region. A
+// regular key is stored raw, so one whose bytes sort after that prefix (e.g.
+// "\xff\xffz") lives *after* the internal keys — the region has a real upper
+// end and cannot be pruned away with an upper bound.
+func (encoderNatural) InternalKeyRange() (start []byte, end []byte) {
+	return encodedInternalKeyPrefixBytes, encodedInternalKeyPrefixSuccessor
+}
+
+var encodedInternalKeyPrefixSuccessor = keyRegionSuccessor(encodedInternalKeyPrefixBytes)
+
 var EncoderNatural Encoder = &encoderNatural{}
 
 // /////////////////////////////////////////////////////////////////////////////////////////////////////////////////

oxiad/dataserver/database/kvstore/kv_pebble.go

@@ -430,14 +430,13 @@ func (p *Pebble) getFloor(key []byte, itOpts IteratorOpts) (returnedKey string,
 }
 
 func (p *Pebble) getCeiling(key []byte, itOpts IteratorOpts) (returnedKey string, value []byte, closer io.Closer, err error) {
-	opts := newIterOptions(p.keyEncoder, itOpts)
-	opts.LowerBound = key
-	it, err := p.db.NewIter(opts)
+	it, err := p.db.NewIter(newIterOptions(p.keyEncoder, itOpts, key, nil))
 	if err != nil {
 		return "", nil, nil, err
 	}
+	skipper := newInternalRegionSkipper(p.keyEncoder, itOpts)
 
-	if !it.First() {
+	if !it.First() || !skipper.forward(it) {
 		return "", nil, nil, multierr.Combine(it.Close(), pebble.ErrNotFound)
 	}
 
@@ -447,14 +446,15 @@ func (p *Pebble) getCeiling(key []byte, itOpts IteratorOpts) (returnedKey string
 }
 
 func (p *Pebble) getLower(key []byte, itOpts IteratorOpts) (returnedKey string, value []byte, closer io.Closer, err error) {
-	opts := newIterOptions(p.keyEncoder, itOpts)
-	opts.UpperBound = key
-	it, err := p.db.NewIter(opts)
+	it, err := p.db.NewIter(newIterOptions(p.keyEncoder, itOpts, nil, key))
 	if err != nil {
 		return "", nil, nil, err
 	}
+	skipper := newInternalRegionSkipper(p.keyEncoder, itOpts)
 
-	if !it.Last() {
+	// Backwards, the internal region is just as costly to step through: a floor
+	// probe above it would walk the whole backlog in reverse.
+	if !it.Last() || !skipper.backward(it) {
 		return "", nil, nil, multierr.Combine(it.Close(), pebble.ErrNotFound)
 	}
 
@@ -464,21 +464,20 @@ func (p *Pebble) getLower(key []byte, itOpts IteratorOpts) (returnedKey string,
 }
 
 func (p *Pebble) getHigher(key []byte, itOpts IteratorOpts) (returnedKey string, value []byte, closer io.Closer, err error) {
-	opts := newIterOptions(p.keyEncoder, itOpts)
-	opts.LowerBound = key
-	it, err := p.db.NewIter(opts)
+	it, err := p.db.NewIter(newIterOptions(p.keyEncoder, itOpts, key, nil))
 	if err != nil {
 		return "", nil, nil, err
 	}
+	skipper := newInternalRegionSkipper(p.keyEncoder, itOpts)
 
-	if !it.First() {
+	if !it.First() || !skipper.forward(it) {
 		return "", nil, nil, multierr.Combine(it.Close(), pebble.ErrNotFound)
 	}
 
 	// The lower bound is inclusive, so the iterator may be positioned exactly on
 	// the key. We are looking for a strict `x > y`, so step over it.
 	if bytes.Equal(it.Key(), key) {
-		if !it.Next() {
+		if !it.Next() || !skipper.forward(it) {
 			return "", nil, nil, multierr.Combine(it.Close(), pebble.ErrNotFound)
 		}
 	}
@@ -521,62 +520,52 @@ func (p *Pebble) KeyRangeScan(lowerBound, upperBound string, opts IteratorOpts)
 }
 
 func (p *Pebble) KeyIterator(itOpts IteratorOpts) (KeyIterator, error) {
-	opts := &pebble.IterOptions{}
-	if !itOpts.IncludeInternalKeys {
-		opts.SkipPoint = func(encodedKey []byte) bool {
-			return p.keyEncoder.IsInternalKey(encodedKey)
-		}
-	}
-	pbit, err := p.db.NewIter(opts)
+	pbit, err := p.db.NewIter(newIterOptions(p.keyEncoder, itOpts, nil, nil))
 	if err != nil {
 		return nil, err
 	}
 
-	return &PebbleIterator{p, pbit}, nil
+	return &PebbleIterator{p, pbit, newInternalRegionSkipper(p.keyEncoder, itOpts)}, nil
 }
 
 func (p *Pebble) KeyRangeScanReverse(lowerBound, upperBound string, itOpts IteratorOpts) (ReverseKeyIterator, error) {
-	opts := &pebble.IterOptions{}
-	if !itOpts.IncludeInternalKeys {
-		opts.SkipPoint = func(encodedKey []byte) bool {
-			return p.keyEncoder.IsInternalKey(encodedKey)
-		}
-	}
+	var lb, ub []byte
 	if lowerBound != "" {
-		opts.LowerBound = p.keyEncoder.Encode(lowerBound)
+		lb = p.keyEncoder.Encode(lowerBound)
 	}
 	if upperBound != "" {
-		opts.UpperBound = p.keyEncoder.Encode(upperBound)
+		ub = p.keyEncoder.Encode(upperBound)
 	}
-	pbit, err := p.db.NewIter(opts)
+	pbit, err := p.db.NewIter(newIterOptions(p.keyEncoder, itOpts, lb, ub))
 	if err != nil {
 		return nil, err
 	}
-	pbit.Last()
-	return &PebbleReverseIterator{p, pbit}, nil
+	skipper := newInternalRegionSkipper(p.keyEncoder, itOpts)
+	if pbit.Last() {
+		skipper.backward(pbit)
+	}
+	return &PebbleReverseIterator{p, pbit, skipper}, nil
 }
 
 func (p *Pebble) RangeScan(lowerBound, upperBound string, itOpts IteratorOpts) (KeyValueIterator, error) {
-	opts := &pebble.IterOptions{}
-	if !itOpts.IncludeInternalKeys {
-		opts.SkipPoint = func(encodedKey []byte) bool {
-			return p.keyEncoder.IsInternalKey(encodedKey)
-		}
-	}
+	var lb, ub []byte
 	if lowerBound != "" {
-		opts.LowerBound = p.keyEncoder.Encode(lowerBound)
+		lb = p.keyEncoder.Encode(lowerBound)
 	}
 	if upperBound != "" {
-		opts.UpperBound = p.keyEncoder.Encode(upperBound)
+		ub = p.keyEncoder.Encode(upperBound)
 	}
 
-	pbit, err := p.db.NewIter(opts)
+	pbit, err := p.db.NewIter(newIterOptions(p.keyEncoder, itOpts, lb, ub))
 	if err != nil {
 		return nil, err
 	}
 
-	pbit.First()
-	return &PebbleIterator{p, pbit}, nil
+	skipper := newInternalRegionSkipper(p.keyEncoder, itOpts)
+	if pbit.First() {
+		skipper.forward(pbit)
+	}
+	return &PebbleIterator{p, pbit, skipper}, nil
 }
 
 func (p *Pebble) Snapshot() (Snapshot, error) {
@@ -620,7 +609,7 @@ func (b *PebbleBatch) RangeScan(lowerBound, upperBound string) (KeyValueIterator
 		return nil, err
 	}
 	pbit.SeekGE(lb)
-	return &PebbleIterator{b.p, pbit}, nil
+	return &PebbleIterator{b.p, pbit, internalRegionSkipper{}}, nil
 }
 
 func (b *PebbleBatch) Close() error {
@@ -714,8 +703,9 @@ func (b *PebbleBatch) Commit() error {
 // Iterator wrapper methods
 
 type PebbleIterator struct {
-	p  *Pebble
-	pi *pebble.Iterator
+	p       *Pebble
+	pi      *pebble.Iterator
+	skipper internalRegionSkipper
 }
 
 func (p *PebbleIterator) Close() error {
@@ -731,19 +721,19 @@ func (p *PebbleIterator) Key() string {
 }
 
 func (p *PebbleIterator) Next() bool {
-	return p.pi.Next()
+	return p.pi.Next() && p.skipper.forward(p.pi)
 }
 
 func (p *PebbleIterator) Prev() bool {
-	return p.pi.Prev()
+	return p.pi.Prev() && p.skipper.backward(p.pi)
 }
 
 func (p *PebbleIterator) SeekGE(key string) bool {
-	return p.pi.SeekGE(p.p.keyEncoder.Encode(key))
+	return p.pi.SeekGE(p.p.keyEncoder.Encode(key)) && p.skipper.forward(p.pi)
 }
 
 func (p *PebbleIterator) SeekLT(key string) bool {
-	return p.pi.SeekLT(p.p.keyEncoder.Encode(key))
+	return p.pi.SeekLT(p.p.keyEncoder.Encode(key)) && p.skipper.backward(p.pi)
 }
 
 func (p *PebbleIterator) Value() ([]byte, error) {
@@ -757,8 +747,9 @@ func (p *PebbleIterator) Value() ([]byte, error) {
 // Iterator wrapper methods
 
 type PebbleReverseIterator struct {
-	p  *Pebble
-	pi *pebble.Iterator
+	p       *Pebble
+	pi      *pebble.Iterator
+	skipper internalRegionSkipper
 }
 
 func (p *PebbleReverseIterator) Close() error {
@@ -774,7 +765,7 @@ func (p *PebbleReverseIterator) Key() string {
 }
 
 func (p *PebbleReverseIterator) Prev() bool {
-	return p.pi.Prev()
+	return p.pi.Prev() && p.skipper.backward(p.pi)
 }
 
 func (p *PebbleReverseIterator) Value() ([]byte, error) {
@@ -855,12 +846,70 @@ func (sl *pebbleSnapshotLoader) Complete() {
 	sl.complete = true
 }
 
-func newIterOptions(keyEncoder compare.Encoder, itOpts IteratorOpts) *pebble.IterOptions {
-	opts := &pebble.IterOptions{}
-	if !itOpts.IncludeInternalKeys {
-		opts.SkipPoint = func(encodedKey []byte) bool {
-			return keyEncoder.IsInternalKey(encodedKey)
+// newIterOptions builds the Pebble iterator options for a scan that may have to
+// exclude the internal-key region.
+//
+// It deliberately does not use pebble's SkipPoint. SkipPoint is a *filter*: the
+// iterator still steps through — and loads the blocks of — every internal key it
+// passes, so a ceiling/higher lookup past the last user key, or a scan with no
+// upper bound, walks the entire live notification backlog before returning. It
+// also hides those keys from us, so we could not detect the region to seek past
+// it. Instead: when the internal region runs to the end of the keyspace it is
+// pruned outright with an upper bound, and otherwise internalRegionSkipper
+// jumps over it with a single seek.
+func newIterOptions(enc compare.Encoder, itOpts IteratorOpts, lowerBound, upperBound []byte) *pebble.IterOptions {
+	opts := &pebble.IterOptions{LowerBound: lowerBound, UpperBound: upperBound}
+	if itOpts.IncludeInternalKeys {
+		return opts
+	}
+	if start, end := enc.InternalKeyRange(); end == nil {
+		// No user key can sort after the internal keys: bound them away.
+		if opts.UpperBound == nil || bytes.Compare(start, opts.UpperBound) < 0 {
+			opts.UpperBound = start
 		}
 	}
 	return opts
 }
+
+// internalRegionSkipper keeps an iterator out of the contiguous internal-key
+// region, jumping over it in a single seek rather than stepping through it.
+//
+// It is a no-op when internal keys are wanted, and when newIterOptions already
+// pruned the region with an upper bound (hierarchical). It carries its weight
+// for the natural encoder, where a user key is stored raw and can sort *after*
+// the internal keys, so the region cannot be bounded away.
+type internalRegionSkipper struct {
+	enc        compare.Encoder
+	start, end []byte
+}
+
+func newInternalRegionSkipper(enc compare.Encoder, itOpts IteratorOpts) internalRegionSkipper {
+	if itOpts.IncludeInternalKeys {
+		return internalRegionSkipper{}
+	}
+	start, end := enc.InternalKeyRange()
+	if end == nil {
+		// Already pruned by the upper bound in newIterOptions.
+		return internalRegionSkipper{}
+	}
+	return internalRegionSkipper{enc: enc, start: start, end: end}
+}
+
+// forward moves the iterator past the internal region when it is positioned
+// inside it, and reports whether it is still valid. A single seek suffices: the
+// region is contiguous, so nothing at or after end is an internal key.
+func (s internalRegionSkipper) forward(it *pebble.Iterator) bool {
+	if s.end == nil || !it.Valid() || !s.enc.IsInternalKey(it.Key()) {
+		return it.Valid()
+	}
+	return it.SeekGE(s.end)
+}
+
+// backward moves the iterator before the internal region when it is positioned
+// inside it, and reports whether it is still valid.
+func (s internalRegionSkipper) backward(it *pebble.Iterator) bool {
+	if s.end == nil || !it.Valid() || !s.enc.IsInternalKey(it.Key()) {
+		return it.Valid()
+	}
+	return it.SeekLT(s.start)
+}