zapdb/table/builder.go

/*
 * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc.
 * SPDX-License-Identifier: Apache-2.0
 */

package table

import (
	"crypto/aes"
	"errors"
	"math"
	"runtime"
	"sync"
	"sync/atomic"
	"unsafe"

	fbs "github.com/google/flatbuffers/go"
	"github.com/klauspost/compress/s2"
	"google.golang.org/protobuf/proto"

	"github.com/dgraph-io/badger/v4/fb"
	"github.com/dgraph-io/badger/v4/options"
	"github.com/dgraph-io/badger/v4/pb"
	"github.com/dgraph-io/badger/v4/y"
	"github.com/dgraph-io/ristretto/v2/z"
)

const (
	KB = 1024
	MB = KB * 1024

	// When a block is encrypted, it's length increases. We add 256 bytes of padding to
	// handle cases when block size increases. This is an approximate number.
	padding = 256
)

type header struct {
	overlap uint16 // Overlap with base key.
	diff    uint16 // Length of the diff.
}

const headerSize = uint16(unsafe.Sizeof(header{}))

// Encode encodes the header.
func (h header) Encode() []byte {
	var b [4]byte
	*(*header)(unsafe.Pointer(&b[0])) = h
	return b[:]
}

// Decode decodes the header.
func (h *header) Decode(buf []byte) {
	// Copy over data from buf into h. Using *h=unsafe.pointer(...) leads to
	// pointer alignment issues. See https://github.com/dgraph-io/badger/issues/1096
	// and comment https://github.com/dgraph-io/badger/pull/1097#pullrequestreview-307361714
	copy(((*[headerSize]byte)(unsafe.Pointer(h))[:]), buf[:headerSize])
}

// bblock represents a block that is being compressed/encrypted in the background.
type bblock struct {
	data         []byte
	baseKey      []byte   // Base key for the current block.
	entryOffsets []uint32 // Offsets of entries present in current block.
	end          int      // Points to the end offset of the block.
}

// Builder is used in building a table.
type Builder struct {
	// Typically tens or hundreds of meg. This is for one single file.
	alloc            *z.Allocator
	curBlock         *bblock
	compressedSize   atomic.Uint32
	uncompressedSize atomic.Uint32

	lenOffsets    uint32
	keyHashes     []uint32 // Used for building the bloomfilter.
	opts          *Options
	maxVersion    uint64
	onDiskSize    uint32
	staleDataSize int

	// Used to concurrently compress/encrypt blocks.
	wg        sync.WaitGroup
	blockChan chan *bblock
	blockList []*bblock
}

func (b *Builder) allocate(need int) []byte {
	bb := b.curBlock
	if len(bb.data[bb.end:]) < need {
		// We need to reallocate. 1GB is the max size that the allocator can allocate.
		// While reallocating, if doubling exceeds that limit, then put the upper bound on it.
		sz := 2 * len(bb.data)
		if sz > (1 << 30) {
			sz = 1 << 30
		}
		if bb.end+need > sz {
			sz = bb.end + need
		}
		tmp := b.alloc.Allocate(sz)
		copy(tmp, bb.data)
		bb.data = tmp
	}
	bb.end += need
	return bb.data[bb.end-need : bb.end]
}

// append appends to curBlock.data
func (b *Builder) append(data []byte) {
	dst := b.allocate(len(data))
	y.AssertTrue(len(data) == copy(dst, data))
}

const maxAllocatorInitialSz = 256 << 20

// NewTableBuilder makes a new TableBuilder.
func NewTableBuilder(opts Options) *Builder {
	sz := 2 * int(opts.TableSize)
	if sz > maxAllocatorInitialSz {
		sz = maxAllocatorInitialSz
	}
	b := &Builder{
		alloc: opts.AllocPool.Get(sz, "TableBuilder"),
		opts:  &opts,
	}
	b.alloc.Tag = "Builder"
	b.curBlock = &bblock{
		data: b.alloc.Allocate(opts.BlockSize + padding),
	}
	b.opts.tableCapacity = uint64(float64(b.opts.TableSize) * 0.95)

	// If encryption or compression is not enabled, do not start compression/encryption goroutines
	// and write directly to the buffer.
	if b.opts.Compression == options.None && b.opts.DataKey == nil {
		return b
	}

	count := 2 * runtime.NumCPU()
	b.blockChan = make(chan *bblock, count*2)

	b.wg.Add(count)
	for i := 0; i < count; i++ {
		go b.handleBlock()
	}
	return b
}

func maxEncodedLen(ctype options.CompressionType, sz int) int {
	switch ctype {
	case options.Snappy:
		return s2.MaxEncodedLen(sz)
	case options.ZSTD:
		return y.ZSTDCompressBound(sz)
	}
	return sz
}

func (b *Builder) handleBlock() {
	defer b.wg.Done()

	doCompress := b.opts.Compression != options.None
	for item := range b.blockChan {
		// Extract the block.
		blockBuf := item.data[:item.end]
		// Compress the block.
		if doCompress {
			out, err := b.compressData(blockBuf)
			y.Check(err)
			blockBuf = out
		}
		if b.shouldEncrypt() {
			out, err := b.encrypt(blockBuf)
			y.Check(y.Wrapf(err, "Error while encrypting block in table builder."))
			blockBuf = out
		}

		// BlockBuf should always less than or equal to allocated space. If the blockBuf is greater
		// than allocated space that means the data from this block cannot be stored in its
		// existing location.
		allocatedSpace := maxEncodedLen(b.opts.Compression, (item.end)) + padding + 1
		y.AssertTrue(len(blockBuf) <= allocatedSpace)

		// blockBuf was allocated on allocator. So, we don't need to copy it over.
		item.data = blockBuf
		item.end = len(blockBuf)
		b.compressedSize.Add(uint32(len(blockBuf)))
	}
}

// Close closes the TableBuilder.
func (b *Builder) Close() {
	b.opts.AllocPool.Return(b.alloc)
}

// Empty returns whether it's empty.
func (b *Builder) Empty() bool { return len(b.keyHashes) == 0 }

// keyDiff returns a suffix of newKey that is different from b.baseKey.
func (b *Builder) keyDiff(newKey []byte) []byte {
	var i int
	for i = 0; i < len(newKey) && i < len(b.curBlock.baseKey); i++ {
		if newKey[i] != b.curBlock.baseKey[i] {
			break
		}
	}
	return newKey[i:]
}

func (b *Builder) addHelper(key []byte, v y.ValueStruct, vpLen uint32) {
	b.keyHashes = append(b.keyHashes, y.Hash(y.ParseKey(key)))

	if version := y.ParseTs(key); version > b.maxVersion {
		b.maxVersion = version
	}

	// diffKey stores the difference of key with baseKey.
	var diffKey []byte
	if len(b.curBlock.baseKey) == 0 {
		// Make a copy. Builder should not keep references. Otherwise, caller has to be very careful
		// and will have to make copies of keys every time they add to builder, which is even worse.
		b.curBlock.baseKey = append(b.curBlock.baseKey[:0], key...)
		diffKey = key
	} else {
		diffKey = b.keyDiff(key)
	}

	y.AssertTrue(len(key)-len(diffKey) <= math.MaxUint16)
	y.AssertTrue(len(diffKey) <= math.MaxUint16)

	h := header{
		overlap: uint16(len(key) - len(diffKey)),
		diff:    uint16(len(diffKey)),
	}

	// store current entry's offset
	b.curBlock.entryOffsets = append(b.curBlock.entryOffsets, uint32(b.curBlock.end))

	// Layout: header, diffKey, value.
	b.append(h.Encode())
	b.append(diffKey)

	dst := b.allocate(int(v.EncodedSize()))
	v.Encode(dst)

	// Add the vpLen to the onDisk size. We'll add the size of the block to
	// onDisk size in Finish() function.
	b.onDiskSize += vpLen
}

/*
Structure of Block.
+-------------------+---------------------+--------------------+--------------+------------------+
| Entry1            | Entry2              | Entry3             | Entry4       | Entry5           |
+-------------------+---------------------+--------------------+--------------+------------------+
| Entry6            | ...                 | ...                | ...          | EntryN           |
+-------------------+---------------------+--------------------+--------------+------------------+
| Block Meta(contains list of offsets used| Block Meta Size    | Block        | Checksum Size    |
| to perform binary search in the block)  | (4 Bytes)          | Checksum     | (4 Bytes)        |
+-----------------------------------------+--------------------+--------------+------------------+
*/
// In case the data is encrypted, the "IV" is added to the end of the block.
func (b *Builder) finishBlock() {
	if len(b.curBlock.entryOffsets) == 0 {
		return
	}
	// Append the entryOffsets and its length.
	b.append(y.U32SliceToBytes(b.curBlock.entryOffsets))
	b.append(y.U32ToBytes(uint32(len(b.curBlock.entryOffsets))))

	checksum := b.calculateChecksum(b.curBlock.data[:b.curBlock.end])

	// Append the block checksum and its length.
	b.append(checksum)
	b.append(y.U32ToBytes(uint32(len(checksum))))

	b.blockList = append(b.blockList, b.curBlock)
	b.uncompressedSize.Add(uint32(b.curBlock.end))

	// Add length of baseKey (rounded to next multiple of 4 because of alignment).
	// Add another 40 Bytes, these additional 40 bytes consists of
	// 12 bytes of metadata of flatbuffer
	// 8 bytes for Key in flat buffer
	// 8 bytes for offset
	// 8 bytes for the len
	// 4 bytes for the size of slice while SliceAllocate
	b.lenOffsets += uint32(int(math.Ceil(float64(len(b.curBlock.baseKey))/4))*4) + 40

	// If compression/encryption is enabled, we need to send the block to the blockChan.
	if b.blockChan != nil {
		b.blockChan <- b.curBlock
	}
}

func (b *Builder) shouldFinishBlock(key []byte, value y.ValueStruct) bool {
	// If there is no entry till now, we will return false.
	if len(b.curBlock.entryOffsets) <= 0 {
		return false
	}

	// Integer overflow check for statements below.
	y.AssertTrue((uint32(len(b.curBlock.entryOffsets))+1)*4+4+8+4 < math.MaxUint32)
	// We should include current entry also in size, that's why +1 to len(b.entryOffsets).
	entriesOffsetsSize := uint32((len(b.curBlock.entryOffsets)+1)*4 +
		4 + // size of list
		8 + // Sum64 in checksum proto
		4) // checksum length
	estimatedSize := uint32(b.curBlock.end) + uint32(6 /*header size for entry*/) +
		uint32(len(key)) + value.EncodedSize() + entriesOffsetsSize

	if b.shouldEncrypt() {
		// IV is added at the end of the block, while encrypting.
		// So, size of IV is added to estimatedSize.
		estimatedSize += aes.BlockSize
	}

	// Integer overflow check for table size.
	y.AssertTrue(uint64(b.curBlock.end)+uint64(estimatedSize) < math.MaxUint32)

	return estimatedSize > uint32(b.opts.BlockSize)
}

// AddStaleKey is same is Add function but it also increments the internal
// staleDataSize counter. This value will be used to prioritize this table for
// compaction.
func (b *Builder) AddStaleKey(key []byte, v y.ValueStruct, valueLen uint32) {
	// Rough estimate based on how much space it will occupy in the SST.
	b.staleDataSize += len(key) + len(v.Value) + 4 /* entry offset */ + 4 /* header size */
	b.addInternal(key, v, valueLen, true)
}

// Add adds a key-value pair to the block.
func (b *Builder) Add(key []byte, value y.ValueStruct, valueLen uint32) {
	b.addInternal(key, value, valueLen, false)
}

func (b *Builder) addInternal(key []byte, value y.ValueStruct, valueLen uint32, isStale bool) {
	if b.shouldFinishBlock(key, value) {
		if isStale {
			// This key will be added to tableIndex and it is stale.
			b.staleDataSize += len(key) + 4 /* len */ + 4 /* offset */
		}
		b.finishBlock()
		// Create a new block and start writing.
		b.curBlock = &bblock{
			data: b.alloc.Allocate(b.opts.BlockSize + padding),
		}
	}
	b.addHelper(key, value, valueLen)
}

// TODO: vvv this was the comment on ReachedCapacity.
// FinalSize returns the *rough* final size of the array, counting the header which is
// not yet written.
// TODO: Look into why there is a discrepancy. I suspect it is because of Write(empty, empty)
// at the end. The diff can vary.

// ReachedCapacity returns true if we... roughly (?) reached capacity?
func (b *Builder) ReachedCapacity() bool {
	// If encryption/compression is enabled then use the compressed size.
	sumBlockSizes := b.compressedSize.Load()
	if b.opts.Compression == options.None && b.opts.DataKey == nil {
		sumBlockSizes = b.uncompressedSize.Load()
	}
	blocksSize := sumBlockSizes + // actual length of current buffer
		uint32(len(b.curBlock.entryOffsets)*4) + // all entry offsets size
		4 + // count of all entry offsets
		8 + // checksum bytes
		4 // checksum length

	estimateSz := blocksSize +
		4 + // Index length
		b.lenOffsets

	return uint64(estimateSz) > b.opts.tableCapacity
}

// Finish finishes the table by appending the index.
/*
The table structure looks like
+---------+------------+-----------+---------------+
| Block 1 | Block 2    | Block 3   | Block 4       |
+---------+------------+-----------+---------------+
| Block 5 | Block 6    | Block ... | Block N       |
+---------+------------+-----------+---------------+
| Index   | Index Size | Checksum  | Checksum Size |
+---------+------------+-----------+---------------+
*/
// In case the data is encrypted, the "IV" is added to the end of the index.
func (b *Builder) Finish() []byte {
	bd := b.Done()
	buf := make([]byte, bd.Size)
	written := bd.Copy(buf)
	y.AssertTrue(written == len(buf))
	return buf
}

type buildData struct {
	blockList []*bblock
	index     []byte
	checksum  []byte
	Size      int
	alloc     *z.Allocator
}

func (bd *buildData) Copy(dst []byte) int {
	var written int
	for _, bl := range bd.blockList {
		written += copy(dst[written:], bl.data[:bl.end])
	}
	written += copy(dst[written:], bd.index)
	written += copy(dst[written:], y.U32ToBytes(uint32(len(bd.index))))

	written += copy(dst[written:], bd.checksum)
	written += copy(dst[written:], y.U32ToBytes(uint32(len(bd.checksum))))
	return written
}

func (b *Builder) Done() buildData {
	b.finishBlock() // This will never start a new block.
	if b.blockChan != nil {
		close(b.blockChan)
	}
	// Wait for block handler to finish.
	b.wg.Wait()

	if len(b.blockList) == 0 {
		return buildData{}
	}
	bd := buildData{
		blockList: b.blockList,
		alloc:     b.alloc,
	}

	var f y.Filter
	if b.opts.BloomFalsePositive > 0 {
		bits := y.BloomBitsPerKey(len(b.keyHashes), b.opts.BloomFalsePositive)
		f = y.NewFilter(b.keyHashes, bits)
	}
	index, dataSize := b.buildIndex(f)

	var err error
	if b.shouldEncrypt() {
		index, err = b.encrypt(index)
		y.Check(err)
	}
	checksum := b.calculateChecksum(index)

	bd.index = index
	bd.checksum = checksum
	bd.Size = int(dataSize) + len(index) + len(checksum) + 4 + 4
	return bd
}

func (b *Builder) calculateChecksum(data []byte) []byte {
	// Build checksum for the index.
	checksum := pb.Checksum{
		// TODO: The checksum type should be configurable from the
		// options.
		// We chose to use CRC32 as the default option because
		// it performed better compared to xxHash64.
		// See the BenchmarkChecksum in table_test.go file
		// Size     =>   1024 B        2048 B
		// CRC32    => 63.7 ns/op     112 ns/op
		// xxHash64 => 87.5 ns/op     158 ns/op
		Sum:  y.CalculateChecksum(data, pb.Checksum_CRC32C),
		Algo: pb.Checksum_CRC32C,
	}

	// Write checksum to the file.
	chksum, err := proto.Marshal(&checksum)
	y.Check(err)
	// Write checksum size.
	return chksum
}

// DataKey returns datakey of the builder.
func (b *Builder) DataKey() *pb.DataKey {
	return b.opts.DataKey
}

func (b *Builder) Opts() *Options {
	return b.opts
}

// encrypt will encrypt the given data and appends IV to the end of the encrypted data.
// This should be only called only after checking shouldEncrypt method.
func (b *Builder) encrypt(data []byte) ([]byte, error) {
	iv, err := y.GenerateIV()
	if err != nil {
		return data, y.Wrapf(err, "Error while generating IV in Builder.encrypt")
	}
	needSz := len(data) + len(iv)
	dst := b.alloc.Allocate(needSz)

	if err = y.XORBlock(dst[:len(data)], data, b.DataKey().Data, iv); err != nil {
		return data, y.Wrapf(err, "Error while encrypting in Builder.encrypt")
	}

	y.AssertTrue(len(iv) == copy(dst[len(data):], iv))
	return dst, nil
}

// shouldEncrypt tells us whether to encrypt the data or not.
// We encrypt only if the data key exist. Otherwise, not.
func (b *Builder) shouldEncrypt() bool {
	return b.opts.DataKey != nil
}

// compressData compresses the given data.
func (b *Builder) compressData(data []byte) ([]byte, error) {
	switch b.opts.Compression {
	case options.None:
		return data, nil
	case options.Snappy:
		sz := s2.MaxEncodedLen(len(data))
		dst := b.alloc.Allocate(sz)
		return s2.EncodeSnappy(dst, data), nil
	case options.ZSTD:
		sz := y.ZSTDCompressBound(len(data))
		dst := b.alloc.Allocate(sz)
		return y.ZSTDCompress(dst, data, b.opts.ZSTDCompressionLevel)
	}
	return nil, errors.New("Unsupported compression type")
}

func (b *Builder) buildIndex(bloom []byte) ([]byte, uint32) {
	builder := fbs.NewBuilder(3 << 20)

	boList, dataSize := b.writeBlockOffsets(builder)
	// Write block offset vector the the idxBuilder.
	fb.TableIndexStartOffsetsVector(builder, len(boList))

	// Write individual block offsets in reverse order to work around how Flatbuffers expects it.
	for i := len(boList) - 1; i >= 0; i-- {
		builder.PrependUOffsetT(boList[i])
	}
	boEnd := builder.EndVector(len(boList))

	var bfoff fbs.UOffsetT
	// Write the bloom filter.
	if len(bloom) > 0 {
		bfoff = builder.CreateByteVector(bloom)
	}
	b.onDiskSize += dataSize
	fb.TableIndexStart(builder)
	fb.TableIndexAddOffsets(builder, boEnd)
	fb.TableIndexAddBloomFilter(builder, bfoff)
	fb.TableIndexAddMaxVersion(builder, b.maxVersion)
	fb.TableIndexAddUncompressedSize(builder, b.uncompressedSize.Load())
	fb.TableIndexAddKeyCount(builder, uint32(len(b.keyHashes)))
	fb.TableIndexAddOnDiskSize(builder, b.onDiskSize)
	fb.TableIndexAddStaleDataSize(builder, uint32(b.staleDataSize))
	builder.Finish(fb.TableIndexEnd(builder))

	buf := builder.FinishedBytes()
	index := fb.GetRootAsTableIndex(buf, 0)
	// Mutate the ondisk size to include the size of the index as well.
	y.AssertTrue(index.MutateOnDiskSize(index.OnDiskSize() + uint32(len(buf))))
	return buf, dataSize
}

// writeBlockOffsets writes all the blockOffets in b.offsets and returns the
// offsets for the newly written items.
func (b *Builder) writeBlockOffsets(builder *fbs.Builder) ([]fbs.UOffsetT, uint32) {
	var startOffset uint32
	var uoffs []fbs.UOffsetT
	for _, bl := range b.blockList {
		uoff := b.writeBlockOffset(builder, bl, startOffset)
		uoffs = append(uoffs, uoff)
		startOffset += uint32(bl.end)
	}
	return uoffs, startOffset
}

// writeBlockOffset writes the given key,offset,len triple to the indexBuilder.
// It returns the offset of the newly written blockoffset.
func (b *Builder) writeBlockOffset(
	builder *fbs.Builder, bl *bblock, startOffset uint32) fbs.UOffsetT {
	// Write the key to the buffer.
	k := builder.CreateByteVector(bl.baseKey)

	// Build the blockOffset.
	fb.BlockOffsetStart(builder)
	fb.BlockOffsetAddKey(builder, k)
	fb.BlockOffsetAddOffset(builder, startOffset)
	fb.BlockOffsetAddLen(builder, uint32(bl.end))
	return fb.BlockOffsetEnd(builder)
}