[WIP] Initial commit

2025-05-08 18:12:52 +03:00
commit 5379ae6eb3
1516 changed files with 648142 additions and 0 deletions
--- a/vendor/github.com/cockroachdb/pebble/sstable/suffix_rewriter.go
+++ b/vendor/github.com/cockroachdb/pebble/sstable/suffix_rewriter.go
@@ -0,0 +1,589 @@
+package sstable
+
+import (
+	"bytes"
+	"context"
+	"math"
+	"sync"
+
+	"github.com/cespare/xxhash/v2"
+	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/bytealloc"
+	"github.com/cockroachdb/pebble/internal/invariants"
+	"github.com/cockroachdb/pebble/internal/rangekey"
+	"github.com/cockroachdb/pebble/objstorage"
+)
+
+// RewriteKeySuffixes is deprecated.
+//
+// TODO(sumeer): remove after switching CockroachDB to RewriteKeySuffixesAndReturnFormat.
+func RewriteKeySuffixes(
+	sst []byte,
+	rOpts ReaderOptions,
+	out objstorage.Writable,
+	o WriterOptions,
+	from, to []byte,
+	concurrency int,
+) (*WriterMetadata, error) {
+	meta, _, err := RewriteKeySuffixesAndReturnFormat(sst, rOpts, out, o, from, to, concurrency)
+	return meta, err
+}
+
+// RewriteKeySuffixesAndReturnFormat copies the content of the passed SSTable
+// bytes to a new sstable, written to `out`, in which the suffix `from` has is
+// replaced with `to` in every key. The input sstable must consist of only
+// Sets or RangeKeySets and every key must have `from` as its suffix as
+// determined by the Split function of the Comparer in the passed
+// WriterOptions. Range deletes must not exist in this sstable, as they will
+// be ignored.
+//
+// Data blocks are rewritten in parallel by `concurrency` workers and then
+// assembled into a final SST. Filters are copied from the original SST without
+// modification as they are not affected by the suffix, while block and table
+// properties are only minimally recomputed.
+//
+// TODO(sumeer): document limitations, if any, due to this limited
+// re-computation of properties (is there any loss of fidelity?).
+//
+// Any block and table property collectors configured in the WriterOptions must
+// implement SuffixReplaceableTableCollector/SuffixReplaceableBlockCollector.
+//
+// The WriterOptions.TableFormat is ignored, and the output sstable has the
+// same TableFormat as the input, which is returned in case the caller wants
+// to do some error checking. Suffix rewriting is meant to be efficient, and
+// allowing changes in the TableFormat detracts from that efficiency.
+//
+// Any obsolete bits that key-value pairs may be annotated with are ignored
+// and lost during the rewrite. Additionally, the output sstable has the
+// pebble.obsolete.is_strict property set to false. These limitations could be
+// removed if needed. The current use case for
+// RewriteKeySuffixesAndReturnFormat in CockroachDB is for MVCC-compliant file
+// ingestion, where these files do not contain RANGEDELs and have one
+// key-value pair per userkey -- so they trivially satisfy the strict
+// criteria, and we don't need the obsolete bit as a performance optimization.
+// For disaggregated storage, strict obsolete sstables are needed for L5 and
+// L6, but at the time of writing, we expect such MVCC-compliant file
+// ingestion to only ingest into levels L4 and higher. If this changes, we can
+// do one of two things to get rid of this limitation:
+//   - Validate that there are no duplicate userkeys and no RANGEDELs/MERGEs
+//     in the sstable to be rewritten. Validating no duplicate userkeys is
+//     non-trivial when rewriting blocks in parallel, so we could encode the
+//     pre-existing condition in the (existing) SnapshotPinnedKeys property --
+//     we need to update the external sst writer to calculate and encode this
+//     property.
+//   - Preserve the obsolete bit (with changes to the blockIter).
+func RewriteKeySuffixesAndReturnFormat(
+	sst []byte,
+	rOpts ReaderOptions,
+	out objstorage.Writable,
+	o WriterOptions,
+	from, to []byte,
+	concurrency int,
+) (*WriterMetadata, TableFormat, error) {
+	r, err := NewMemReader(sst, rOpts)
+	if err != nil {
+		return nil, TableFormatUnspecified, err
+	}
+	defer r.Close()
+	return rewriteKeySuffixesInBlocks(r, out, o, from, to, concurrency)
+}
+
+func rewriteKeySuffixesInBlocks(
+	r *Reader, out objstorage.Writable, o WriterOptions, from, to []byte, concurrency int,
+) (*WriterMetadata, TableFormat, error) {
+	if o.Comparer == nil || o.Comparer.Split == nil {
+		return nil, TableFormatUnspecified,
+			errors.New("a valid splitter is required to rewrite suffixes")
+	}
+	if concurrency < 1 {
+		return nil, TableFormatUnspecified, errors.New("concurrency must be >= 1")
+	}
+	// Even though NumValueBlocks = 0 => NumValuesInValueBlocks = 0, check both
+	// as a defensive measure.
+	if r.Properties.NumValueBlocks > 0 || r.Properties.NumValuesInValueBlocks > 0 {
+		return nil, TableFormatUnspecified,
+			errors.New("sstable with a single suffix should not have value blocks")
+	}
+
+	tableFormat := r.tableFormat
+	o.TableFormat = tableFormat
+	w := NewWriter(out, o)
+	defer func() {
+		if w != nil {
+			w.Close()
+		}
+	}()
+
+	for _, c := range w.propCollectors {
+		if _, ok := c.(SuffixReplaceableTableCollector); !ok {
+			return nil, TableFormatUnspecified,
+				errors.Errorf("property collector %s does not support suffix replacement", c.Name())
+		}
+	}
+	for _, c := range w.blockPropCollectors {
+		if _, ok := c.(SuffixReplaceableBlockCollector); !ok {
+			return nil, TableFormatUnspecified,
+				errors.Errorf("block property collector %s does not support suffix replacement", c.Name())
+		}
+	}
+
+	l, err := r.Layout()
+	if err != nil {
+		return nil, TableFormatUnspecified, errors.Wrap(err, "reading layout")
+	}
+
+	if err := rewriteDataBlocksToWriter(r, w, l.Data, from, to, w.split, concurrency); err != nil {
+		return nil, TableFormatUnspecified, errors.Wrap(err, "rewriting data blocks")
+	}
+
+	// Copy over the range key block and replace suffixes in it if it exists.
+	if err := rewriteRangeKeyBlockToWriter(r, w, from, to); err != nil {
+		return nil, TableFormatUnspecified, errors.Wrap(err, "rewriting range key blocks")
+	}
+
+	// Copy over the filter block if it exists (rewriteDataBlocksToWriter will
+	// already have ensured this is valid if it exists).
+	if w.filter != nil && l.Filter.Length > 0 {
+		filterBlock, _, err := readBlockBuf(r, l.Filter, nil)
+		if err != nil {
+			return nil, TableFormatUnspecified, errors.Wrap(err, "reading filter")
+		}
+		w.filter = copyFilterWriter{
+			origPolicyName: w.filter.policyName(), origMetaName: w.filter.metaName(), data: filterBlock,
+		}
+	}
+
+	if err := w.Close(); err != nil {
+		w = nil
+		return nil, TableFormatUnspecified, err
+	}
+	writerMeta, err := w.Metadata()
+	w = nil
+	return writerMeta, tableFormat, err
+}
+
+var errBadKind = errors.New("key does not have expected kind (set)")
+
+type blockWithSpan struct {
+	start, end InternalKey
+	data       []byte
+}
+
+func rewriteBlocks(
+	r *Reader,
+	restartInterval int,
+	checksumType ChecksumType,
+	compression Compression,
+	input []BlockHandleWithProperties,
+	output []blockWithSpan,
+	totalWorkers, worker int,
+	from, to []byte,
+	split Split,
+) error {
+	bw := blockWriter{
+		restartInterval: restartInterval,
+	}
+	buf := blockBuf{checksummer: checksummer{checksumType: checksumType}}
+	if checksumType == ChecksumTypeXXHash {
+		buf.checksummer.xxHasher = xxhash.New()
+	}
+
+	var blockAlloc bytealloc.A
+	var keyAlloc bytealloc.A
+	var scratch InternalKey
+
+	var inputBlock, inputBlockBuf []byte
+
+	iter := &blockIter{}
+
+	// We'll assume all blocks are _roughly_ equal so round-robin static partition
+	// of each worker doing every ith block is probably enough.
+	for i := worker; i < len(input); i += totalWorkers {
+		bh := input[i]
+
+		var err error
+		inputBlock, inputBlockBuf, err = readBlockBuf(r, bh.BlockHandle, inputBlockBuf)
+		if err != nil {
+			return err
+		}
+		if err := iter.init(r.Compare, inputBlock, r.Properties.GlobalSeqNum, false); err != nil {
+			return err
+		}
+
+		if cap(bw.restarts) < int(iter.restarts) {
+			bw.restarts = make([]uint32, 0, iter.restarts)
+		}
+		if cap(bw.buf) == 0 {
+			bw.buf = make([]byte, 0, len(inputBlock))
+		}
+		if cap(bw.restarts) < int(iter.numRestarts) {
+			bw.restarts = make([]uint32, 0, iter.numRestarts)
+		}
+
+		for key, val := iter.First(); key != nil; key, val = iter.Next() {
+			if key.Kind() != InternalKeyKindSet {
+				return errBadKind
+			}
+			si := split(key.UserKey)
+			oldSuffix := key.UserKey[si:]
+			if !bytes.Equal(oldSuffix, from) {
+				err := errors.Errorf("key has suffix %q, expected %q", oldSuffix, from)
+				return err
+			}
+			newLen := si + len(to)
+			if cap(scratch.UserKey) < newLen {
+				scratch.UserKey = make([]byte, 0, len(key.UserKey)*2+len(to)-len(from))
+			}
+
+			scratch.Trailer = key.Trailer
+			scratch.UserKey = scratch.UserKey[:newLen]
+			copy(scratch.UserKey, key.UserKey[:si])
+			copy(scratch.UserKey[si:], to)
+
+			// NB: for TableFormatPebblev3 and higher, since
+			// !iter.lazyValueHandling.hasValuePrefix, it will return the raw value
+			// in the block, which includes the 1-byte prefix. This is fine since bw
+			// also does not know about the prefix and will preserve it in bw.add.
+			v := val.InPlaceValue()
+			if invariants.Enabled && r.tableFormat >= TableFormatPebblev3 &&
+				key.Kind() == InternalKeyKindSet {
+				if len(v) < 1 {
+					return errors.Errorf("value has no prefix")
+				}
+				prefix := valuePrefix(v[0])
+				if isValueHandle(prefix) {
+					return errors.Errorf("value prefix is incorrect")
+				}
+				if setHasSamePrefix(prefix) {
+					return errors.Errorf("multiple keys with same key prefix")
+				}
+			}
+			bw.add(scratch, v)
+			if output[i].start.UserKey == nil {
+				keyAlloc, output[i].start = cloneKeyWithBuf(scratch, keyAlloc)
+			}
+		}
+		*iter = iter.resetForReuse()
+
+		keyAlloc, output[i].end = cloneKeyWithBuf(scratch, keyAlloc)
+
+		finished := compressAndChecksum(bw.finish(), compression, &buf)
+
+		// copy our finished block into the output buffer.
+		blockAlloc, output[i].data = blockAlloc.Alloc(len(finished) + blockTrailerLen)
+		copy(output[i].data, finished)
+		copy(output[i].data[len(finished):], buf.tmp[:blockTrailerLen])
+	}
+	return nil
+}
+
+func rewriteDataBlocksToWriter(
+	r *Reader,
+	w *Writer,
+	data []BlockHandleWithProperties,
+	from, to []byte,
+	split Split,
+	concurrency int,
+) error {
+	if r.Properties.NumEntries == 0 {
+		// No point keys.
+		return nil
+	}
+	blocks := make([]blockWithSpan, len(data))
+
+	if w.filter != nil {
+		if r.Properties.FilterPolicyName != w.filter.policyName() {
+			return errors.New("mismatched filters")
+		}
+		if was, is := r.Properties.ComparerName, w.props.ComparerName; was != is {
+			return errors.Errorf("mismatched Comparer %s vs %s, replacement requires same splitter to copy filters", was, is)
+		}
+	}
+
+	g := &sync.WaitGroup{}
+	g.Add(concurrency)
+	errCh := make(chan error, concurrency)
+	for i := 0; i < concurrency; i++ {
+		worker := i
+		go func() {
+			defer g.Done()
+			err := rewriteBlocks(
+				r,
+				w.dataBlockBuf.dataBlock.restartInterval,
+				w.blockBuf.checksummer.checksumType,
+				w.compression,
+				data,
+				blocks,
+				concurrency,
+				worker,
+				from, to,
+				split,
+			)
+			if err != nil {
+				errCh <- err
+			}
+		}()
+	}
+	g.Wait()
+	close(errCh)
+	if err, ok := <-errCh; ok {
+		return err
+	}
+
+	for _, p := range w.propCollectors {
+		if err := p.(SuffixReplaceableTableCollector).UpdateKeySuffixes(r.Properties.UserProperties, from, to); err != nil {
+			return err
+		}
+	}
+
+	var decoder blockPropertiesDecoder
+	var oldShortIDs []shortID
+	var oldProps [][]byte
+	if len(w.blockPropCollectors) > 0 {
+		oldProps = make([][]byte, len(w.blockPropCollectors))
+		oldShortIDs = make([]shortID, math.MaxUint8)
+		for i, p := range w.blockPropCollectors {
+			if prop, ok := r.Properties.UserProperties[p.Name()]; ok {
+				was, is := shortID(byte(prop[0])), shortID(i)
+				oldShortIDs[was] = is
+			}
+		}
+	}
+
+	for i := range blocks {
+		// Write the rewritten block to the file.
+		if err := w.writable.Write(blocks[i].data); err != nil {
+			return err
+		}
+
+		n := len(blocks[i].data)
+		bh := BlockHandle{Offset: w.meta.Size, Length: uint64(n) - blockTrailerLen}
+		// Update the overall size.
+		w.meta.Size += uint64(n)
+
+		// Load any previous values for our prop collectors into oldProps.
+		for i := range oldProps {
+			oldProps[i] = nil
+		}
+		decoder.props = data[i].Props
+		for !decoder.done() {
+			id, val, err := decoder.next()
+			if err != nil {
+				return err
+			}
+			oldProps[oldShortIDs[id]] = val
+		}
+
+		for i, p := range w.blockPropCollectors {
+			if err := p.(SuffixReplaceableBlockCollector).UpdateKeySuffixes(oldProps[i], from, to); err != nil {
+				return err
+			}
+		}
+
+		bhp, err := w.maybeAddBlockPropertiesToBlockHandle(bh)
+		if err != nil {
+			return err
+		}
+		var nextKey InternalKey
+		if i+1 < len(blocks) {
+			nextKey = blocks[i+1].start
+		}
+		if err = w.addIndexEntrySync(blocks[i].end, nextKey, bhp, w.dataBlockBuf.tmp[:]); err != nil {
+			return err
+		}
+	}
+
+	w.meta.updateSeqNum(blocks[0].start.SeqNum())
+	w.props.NumEntries = r.Properties.NumEntries
+	w.props.RawKeySize = r.Properties.RawKeySize
+	w.props.RawValueSize = r.Properties.RawValueSize
+	w.meta.SetSmallestPointKey(blocks[0].start)
+	w.meta.SetLargestPointKey(blocks[len(blocks)-1].end)
+	return nil
+}
+
+func rewriteRangeKeyBlockToWriter(r *Reader, w *Writer, from, to []byte) error {
+	iter, err := r.NewRawRangeKeyIter()
+	if err != nil {
+		return err
+	}
+	if iter == nil {
+		// No range keys.
+		return nil
+	}
+	defer iter.Close()
+
+	for s := iter.First(); s != nil; s = iter.Next() {
+		if !s.Valid() {
+			break
+		}
+		for i := range s.Keys {
+			if s.Keys[i].Kind() != base.InternalKeyKindRangeKeySet {
+				return errBadKind
+			}
+			if !bytes.Equal(s.Keys[i].Suffix, from) {
+				return errors.Errorf("key has suffix %q, expected %q", s.Keys[i].Suffix, from)
+			}
+			s.Keys[i].Suffix = to
+		}
+
+		err := rangekey.Encode(s, func(k base.InternalKey, v []byte) error {
+			// Calling AddRangeKey instead of addRangeKeySpan bypasses the fragmenter.
+			// This is okay because the raw fragments off of `iter` are already
+			// fragmented, and suffix replacement should not affect fragmentation.
+			return w.AddRangeKey(k, v)
+		})
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+type copyFilterWriter struct {
+	origMetaName   string
+	origPolicyName string
+	data           []byte
+}
+
+func (copyFilterWriter) addKey(key []byte)         { panic("unimplemented") }
+func (c copyFilterWriter) finish() ([]byte, error) { return c.data, nil }
+func (c copyFilterWriter) metaName() string        { return c.origMetaName }
+func (c copyFilterWriter) policyName() string      { return c.origPolicyName }
+
+// RewriteKeySuffixesViaWriter is similar to RewriteKeySuffixes but uses just a
+// single loop over the Reader that writes each key to the Writer with the new
+// suffix. The is significantly slower than the parallelized rewriter, and does
+// more work to rederive filters, props, etc.
+//
+// Any obsolete bits that key-value pairs may be annotated with are ignored
+// and lost during the rewrite. Some of the obsolete bits may be recreated --
+// specifically when there are multiple keys with the same user key.
+// Additionally, the output sstable has the pebble.obsolete.is_strict property
+// set to false. See the longer comment at RewriteKeySuffixesAndReturnFormat.
+func RewriteKeySuffixesViaWriter(
+	r *Reader, out objstorage.Writable, o WriterOptions, from, to []byte,
+) (*WriterMetadata, error) {
+	if o.Comparer == nil || o.Comparer.Split == nil {
+		return nil, errors.New("a valid splitter is required to rewrite suffixes")
+	}
+
+	o.IsStrictObsolete = false
+	w := NewWriter(out, o)
+	defer func() {
+		if w != nil {
+			w.Close()
+		}
+	}()
+	i, err := r.NewIter(nil, nil)
+	if err != nil {
+		return nil, err
+	}
+	defer i.Close()
+
+	k, v := i.First()
+	var scratch InternalKey
+	for k != nil {
+		if k.Kind() != InternalKeyKindSet {
+			return nil, errors.New("invalid key type")
+		}
+		oldSuffix := k.UserKey[r.Split(k.UserKey):]
+		if !bytes.Equal(oldSuffix, from) {
+			return nil, errors.Errorf("key has suffix %q, expected %q", oldSuffix, from)
+		}
+		scratch.UserKey = append(scratch.UserKey[:0], k.UserKey[:len(k.UserKey)-len(from)]...)
+		scratch.UserKey = append(scratch.UserKey, to...)
+		scratch.Trailer = k.Trailer
+
+		val, _, err := v.Value(nil)
+		if err != nil {
+			return nil, err
+		}
+		if w.addPoint(scratch, val, false); err != nil {
+			return nil, err
+		}
+		k, v = i.Next()
+	}
+	if err := rewriteRangeKeyBlockToWriter(r, w, from, to); err != nil {
+		return nil, err
+	}
+	if err := w.Close(); err != nil {
+		w = nil
+		return nil, err
+	}
+	writerMeta, err := w.Metadata()
+	w = nil
+	return writerMeta, err
+}
+
+// NewMemReader opens a reader over the SST stored in the passed []byte.
+func NewMemReader(sst []byte, o ReaderOptions) (*Reader, error) {
+	return NewReader(newMemReader(sst), o)
+}
+
+func readBlockBuf(r *Reader, bh BlockHandle, buf []byte) ([]byte, []byte, error) {
+	raw := r.readable.(*memReader).b[bh.Offset : bh.Offset+bh.Length+blockTrailerLen]
+	if err := checkChecksum(r.checksumType, raw, bh, 0); err != nil {
+		return nil, buf, err
+	}
+	typ := blockType(raw[bh.Length])
+	raw = raw[:bh.Length]
+	if typ == noCompressionBlockType {
+		return raw, buf, nil
+	}
+	decompressedLen, prefix, err := decompressedLen(typ, raw)
+	if err != nil {
+		return nil, buf, err
+	}
+	if cap(buf) < decompressedLen {
+		buf = make([]byte, decompressedLen)
+	}
+	res, err := decompressInto(typ, raw[prefix:], buf[:decompressedLen])
+	return res, buf, err
+}
+
+// memReader is a thin wrapper around a []byte such that it can be passed to
+// sstable.Reader. It supports concurrent use, and does so without locking in
+// contrast to the heavier read/write vfs.MemFile.
+type memReader struct {
+	b  []byte
+	r  *bytes.Reader
+	rh objstorage.NoopReadHandle
+}
+
+var _ objstorage.Readable = (*memReader)(nil)
+
+func newMemReader(b []byte) *memReader {
+	r := &memReader{
+		b: b,
+		r: bytes.NewReader(b),
+	}
+	r.rh = objstorage.MakeNoopReadHandle(r)
+	return r
+}
+
+// ReadAt is part of objstorage.Readable.
+func (m *memReader) ReadAt(_ context.Context, p []byte, off int64) error {
+	n, err := m.r.ReadAt(p, off)
+	if invariants.Enabled && err == nil && n != len(p) {
+		panic("short read")
+	}
+	return err
+}
+
+// Close is part of objstorage.Readable.
+func (*memReader) Close() error {
+	return nil
+}
+
+// Stat is part of objstorage.Readable.
+func (m *memReader) Size() int64 {
+	return int64(len(m.b))
+}
+
+// NewReadHandle is part of objstorage.Readable.
+func (m *memReader) NewReadHandle(_ context.Context) objstorage.ReadHandle {
+	return &m.rh
+}