ceremonyclient/pebble/sstable/suffix_rewriter.go

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
}