Initial commit

vendor/github.com/cockroachdb/pebble/range_keys.go

@@ -0,0 +1,140 @@
+// Copyright 2021 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package pebble
+
+import (
+	"sync"
+
+	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble/internal/arenaskl"
+	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/keyspan"
+	"github.com/cockroachdb/pebble/internal/rangekey"
+)
+
+const rangeKeyArenaSize = 1 << 20
+
+// RangeKeysArena is an in-memory arena in which range keys are stored.
+//
+// This is a temporary type that will eventually be removed.
+type RangeKeysArena struct {
+	once      sync.Once
+	skl       arenaskl.Skiplist
+	arena     *arenaskl.Arena
+	fragCache keySpanCache
+}
+
+// applyFlushedRangeKeys is a temporary hack to support in-memory only range
+// keys.  We use much of the same code that we will use for the memtable, but we
+// use a separate arena that exists beyond the lifetime of any individual
+// memtable.  For as long as we're relying on this hack, a single *pebble.DB may
+// only store as many range keys as fit in this arena.
+//
+// TODO(jackson): Remove applyFlushedRangeKeys when range keys are persisted.
+func (d *DB) applyFlushedRangeKeys(flushable []*flushableEntry) error {
+	var added uint32
+	for i := 0; i < len(flushable); i++ {
+		iter := flushable[i].newRangeKeyIter(nil)
+		if iter == nil {
+			continue
+		}
+		d.maybeInitializeRangeKeys()
+
+		for s := iter.First(); s.Valid(); s = iter.Next() {
+			// flushable.newRangeKeyIter provides a FragmentIterator, which
+			// iterates over parsed keyspan.Spans.
+			//
+			// While we're faking a flush, we just want to write the original
+			// key into the global arena, so we need to re-encode the span's
+			// keys into internal key-value pairs.
+			//
+			// This awkward recombination will be removed when flushes implement
+			// range-key logic, coalescing range keys and constructing internal
+			// keys.
+			err := rangekey.Encode(s, func(k base.InternalKey, v []byte) error {
+				err := d.rangeKeys.skl.Add(k, v)
+				switch {
+				case err == nil:
+					added++
+					return nil
+				case errors.Is(err, arenaskl.ErrRecordExists):
+					// It's possible that we'll try to add a key to the arena twice
+					// during metamorphic tests that reset the synced state. Ignore.
+					// When range keys are actually flushed to stable storage, this
+					// will go away.
+					return nil
+				default:
+					// err != nil
+					return err
+				}
+			})
+			if err != nil {
+				return err
+			}
+		}
+	}
+	if added > 0 {
+		d.rangeKeys.fragCache.invalidate(added)
+	}
+	return nil
+}
+
+func (d *DB) maybeInitializeRangeKeys() {
+	// Lazily construct the global range key arena, so that tests that
+	// don't use range keys don't need to allocate this long-lived
+	// buffer.
+	d.rangeKeys.once.Do(func() {
+		arenaBuf := make([]byte, rangeKeyArenaSize)
+		d.rangeKeys.arena = arenaskl.NewArena(arenaBuf)
+		d.rangeKeys.skl.Reset(d.rangeKeys.arena, d.cmp)
+		d.rangeKeys.fragCache = keySpanCache{
+			cmp:           d.cmp,
+			formatKey:     d.opts.Comparer.FormatKey,
+			skl:           &d.rangeKeys.skl,
+			constructSpan: rangekey.Decode,
+		}
+	})
+}
+
+func (d *DB) newRangeKeyIter(
+	it *iteratorRangeKeyState, seqNum uint64, batch *Batch, readState *readState, opts *IterOptions,
+) keyspan.FragmentIterator {
+	d.maybeInitializeRangeKeys()
+
+	// TODO(jackson): Preallocate iters, mergingIter, rangeKeyIter in a
+	// structure analogous to iterAlloc.
+	var iters []keyspan.FragmentIterator
+
+	// If there's an indexed batch with range keys, include it.
+	if batch != nil {
+		if rki := batch.newRangeKeyIter(opts); rki != nil {
+			iters = append(iters, rki)
+		}
+	}
+
+	// Next are the flushables: memtables and large batches.
+	for i := len(readState.memtables) - 1; i >= 0; i-- {
+		mem := readState.memtables[i]
+		// We only need to read from memtables which contain sequence numbers older
+		// than seqNum.
+		if logSeqNum := mem.logSeqNum; logSeqNum >= seqNum {
+			continue
+		}
+		if rki := mem.newRangeKeyIter(opts); rki != nil {
+			iters = append(iters, rki)
+		}
+	}
+
+	// For now while range keys are not fully integrated into Pebble, all range
+	// keys ever written to the DB are persisted in the d.rangeKeys arena.
+	frags := d.rangeKeys.fragCache.get()
+	if len(frags) > 0 {
+		iters = append(iters, keyspan.NewIter(d.cmp, frags))
+	}
+	it.rangeKeyIter = rangekey.InitUserIteration(
+		d.cmp, seqNum, &it.alloc.merging, &it.alloc.defraging, iters...,
+	)
+	return it.rangeKeyIter
+}