block: rework physical block code

We change `PhysicalBlock` to always be backed by a `TempBuffer`; this
simplifies the code related to creating the blocks. The trailer is now
also stored directly in this buffer, and we no longer have to issue
two separate `Write` calls.

Note that we were already using a `TempBuffer` in the most important
write paths (columnar sstable writer and blob value writer).

We add a `PhysicalBlockMaker` type which encompasses the compressor
and the checksummer and a single `Make()` method replaces all
`CompressAndChecksum/CopyAndChecksum` variations.

Author: RaduBerinde

sstable/blob/blob.go

@@ -102,13 +102,12 @@ type FileWriter struct {
 	// index block encoding the offsets at which each block is written.
 	// Additionally, when rewriting a blob file, the index block's virtualBlocks
 	// column is also populated to remap blockIDs to the physical block indexes.
-	indexEncoder indexBlockEncoder
-	stats        FileWriterStats
-	flushGov     block.FlushGovernor
-	checksummer  block.Checksummer
-	compressor   block.Compressor
-	cpuMeasurer  base.CPUMeasurer
-	writeQueue   struct {
+	indexEncoder   indexBlockEncoder
+	stats          FileWriterStats
+	flushGov       block.FlushGovernor
+	physBlockMaker block.PhysicalBlockMaker
+	cpuMeasurer    base.CPUMeasurer
+	writeQueue     struct {
 		wg  sync.WaitGroup
 		ch  chan compressedBlock
 		err error
@@ -117,7 +116,6 @@ type FileWriter struct {
 
 type compressedBlock struct {
 	pb  block.PhysicalBlock
-	bh  *block.TempBuffer
 	off uint64
 }
 
@@ -130,8 +128,7 @@ func NewFileWriter(fn base.DiskFileNum, w objstorage.Writable, opts FileWriterOp
 	fw.valuesEncoder.Init()
 	fw.flushGov = opts.FlushGovernor
 	fw.indexEncoder.Init()
-	fw.checksummer = block.Checksummer{Type: opts.ChecksumType}
-	fw.compressor = block.MakeCompressor(opts.Compression)
+	fw.physBlockMaker.Init(opts.Compression, opts.ChecksumType)
 	fw.cpuMeasurer = opts.CpuMeasurer
 	fw.writeQueue.ch = make(chan compressedBlock)
 	fw.writeQueue.wg.Add(1)
@@ -215,12 +212,12 @@ func (w *FileWriter) flush() {
 	if w.valuesEncoder.Count() == 0 {
 		panic(errors.AssertionFailedf("no values to flush"))
 	}
-	pb, bh := block.CompressAndChecksumToTempBuffer(w.valuesEncoder.Finish(), blockkind.BlobValue, &w.compressor, &w.checksummer)
+	pb := w.physBlockMaker.Make(w.valuesEncoder.Finish(), blockkind.BlobValue, block.NoFlags)
 	compressedLen := uint64(pb.LengthWithoutTrailer())
 	w.stats.BlockCount++
 	off := w.stats.FileLen
 	w.stats.FileLen += compressedLen + block.TrailerLen
-	w.writeQueue.ch <- compressedBlock{pb: pb, bh: bh, off: off}
+	w.writeQueue.ch <- compressedBlock{pb: pb, off: off}
 	w.valuesEncoder.Reset()
 }
 
@@ -244,9 +241,8 @@ func (w *FileWriter) drainWriteQueue() {
 			Offset: cb.off,
 			Length: uint64(cb.pb.LengthWithoutTrailer()),
 		})
-		// We're done with the buffer associated with this physical block.
-		// Release it back to its pool.
-		cb.bh.Release()
+		// We're done with this physical block.
+		cb.pb.Release()
 	}
 }
 
@@ -284,14 +280,13 @@ func (w *FileWriter) Close() (FileWriterStats, error) {
 	var indexBlockHandle block.Handle
 	{
 		indexBlock := w.indexEncoder.Finish()
-		var compressedBuf []byte
-		pb := block.CopyAndChecksum(&compressedBuf, indexBlock, blockkind.Metadata, &w.compressor, &w.checksummer)
+		pb := w.physBlockMaker.Make(indexBlock, blockkind.Metadata, block.DontCompress)
+		defer pb.Release()
 		if _, w.err = pb.WriteTo(w.w); w.err != nil {
-			err = w.err
 			if w.w != nil {
 				w.w.Abort()
 			}
-			return FileWriterStats{}, err
+			return FileWriterStats{}, w.err
 		}
 		indexBlockHandle.Offset = stats.FileLen
 		indexBlockHandle.Length = uint64(pb.LengthWithoutTrailer())
@@ -301,7 +296,7 @@ func (w *FileWriter) Close() (FileWriterStats, error) {
 	// Write the footer.
 	footer := fileFooter{
 		format:          FileFormatV1,
-		checksum:        w.checksummer.Type,
+		checksum:        w.physBlockMaker.Checksummer.Type,
 		indexHandle:     indexBlockHandle,
 		originalFileNum: w.fileNum,
 	}

sstable/block/block.go

@@ -101,20 +101,6 @@ func DecodeHandleWithProperties(src []byte) (HandleWithProperties, error) {
 	}, nil
 }
 
-// TrailerLen is the length of the trailer at the end of a block.
-const TrailerLen = 5
-
-// Trailer is the trailer at the end of a block, encoding the block type
-// (compression) and a checksum.
-type Trailer = [TrailerLen]byte
-
-// MakeTrailer constructs a trailer from a block type and a checksum.
-func MakeTrailer(blockType byte, checksum uint32) (t Trailer) {
-	t[0] = blockType
-	binary.LittleEndian.PutUint32(t[1:5], checksum)
-	return t
-}
-
 // ChecksumType specifies the checksum used for blocks.
 type ChecksumType byte
 

sstable/block/compression.go

@@ -6,16 +6,11 @@ package block
 
 import (
 	"runtime"
-	"slices"
 	"strings"
-	"sync"
 
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
-	"github.com/cockroachdb/pebble/internal/bytealloc"
 	"github.com/cockroachdb/pebble/internal/compression"
-	"github.com/cockroachdb/pebble/internal/invariants"
-	"github.com/cockroachdb/pebble/objstorage"
 )
 
 // CompressionProfile contains the parameters for compressing blocks in an
@@ -243,203 +238,3 @@ func DecompressInto(ci CompressionIndicator, compressed []byte, buf []byte) erro
 	}
 	return nil
 }
-
-// PhysicalBlock represents a block (possibly compressed) as it is stored
-// physically on disk, including its trailer.
-type PhysicalBlock struct {
-	// data contains the possibly compressed block data.
-	data    []byte
-	trailer Trailer
-}
-
-// NewPhysicalBlock returns a new PhysicalBlock with the provided block
-// data. The trailer is set from the last TrailerLen bytes of the
-// block. The data could be compressed.
-func NewPhysicalBlock(data []byte) PhysicalBlock {
-	trailer := Trailer(data[len(data)-TrailerLen:])
-	data = data[:len(data)-TrailerLen]
-	return PhysicalBlock{data: data, trailer: trailer}
-}
-
-// LengthWithTrailer returns the length of the data block, including the trailer.
-func (b *PhysicalBlock) LengthWithTrailer() int {
-	return len(b.data) + TrailerLen
-}
-
-// LengthWithoutTrailer returns the length of the data block, excluding the trailer.
-func (b *PhysicalBlock) LengthWithoutTrailer() int {
-	return len(b.data)
-}
-
-// CloneWithByteAlloc returns a deep copy of the block, using the provided
-// bytealloc.A to allocate memory for the new copy.
-func (b *PhysicalBlock) CloneWithByteAlloc(a *bytealloc.A) PhysicalBlock {
-	var data []byte
-	*a, data = (*a).Alloc(len(b.data))
-	copy(data, b.data)
-	return PhysicalBlock{
-		data:    data,
-		trailer: b.trailer,
-	}
-}
-
-// Clone returns a deep copy of the block.
-func (b PhysicalBlock) Clone() PhysicalBlock {
-	data := make([]byte, len(b.data))
-	copy(data, b.data)
-	return PhysicalBlock{data: data, trailer: b.trailer}
-}
-
-// WriteTo writes the block (including its trailer) to the provided Writable. If
-// err == nil, n is the number of bytes successfully written to the Writable.
-//
-// WriteTo might mangle the block data.
-func (b *PhysicalBlock) WriteTo(w objstorage.Writable) (n int, err error) {
-	if err := w.Write(b.data); err != nil {
-		return 0, err
-	}
-	if err := w.Write(b.trailer[:]); err != nil {
-		return 0, err
-	}
-
-	// WriteTo is allowed to mangle the data. Mangle it ourselves some of the time
-	// in invariant builds to catch callers that don't handle this.
-	if invariants.Enabled && invariants.Sometimes(1) {
-		invariants.Mangle(b.data)
-	}
-	return len(b.data) + len(b.trailer), nil
-}
-
-// CompressAndChecksum compresses and checksums the provided block, returning
-// the compressed block and its trailer. The result is appended to the dst
-// argument.
-func CompressAndChecksum(
-	dst *[]byte, blockData []byte, blockKind Kind, compressor *Compressor, checksummer *Checksummer,
-) PhysicalBlock {
-	buf := (*dst)[:0]
-	ci, buf := compressor.Compress(buf, blockData, blockKind)
-	*dst = buf
-
-	// Calculate the checksum.
-	pb := PhysicalBlock{data: buf}
-	checksum := checksummer.Checksum(buf, byte(ci))
-	pb.trailer = MakeTrailer(byte(ci), checksum)
-	return pb
-}
-
-// CopyAndChecksum copies the provided block (without compressing it) and
-// checksums it, returning the physical block. The result is appended to the dst
-// argument.
-//
-// Note that we still need to provide a Compressor so we can inform it of the
-// uncompressed block (for statistics).
-func CopyAndChecksum(
-	dst *[]byte, blockData []byte, blockKind Kind, compressor *Compressor, checksummer *Checksummer,
-) PhysicalBlock {
-	buf := *dst
-	buf = append(buf[:0], blockData...)
-	*dst = buf
-
-	// Calculate the checksum.
-	pb := PhysicalBlock{data: buf}
-	checksum := checksummer.Checksum(buf, byte(NoCompressionIndicator))
-	pb.trailer = MakeTrailer(byte(NoCompressionIndicator), checksum)
-	compressor.UncompressedBlock(len(blockData), blockKind)
-	return pb
-}
-
-// CompressAndChecksumToTempBuffer compresses and checksums the provided block
-// into a TempBuffer. The caller should Release() the TempBuffer once it is no
-// longer necessary.
-func CompressAndChecksumToTempBuffer(
-	blockData []byte, blockKind Kind, compressor *Compressor, checksummer *Checksummer,
-) (PhysicalBlock, *TempBuffer) {
-	// Grab a buffer to use as the destination for compression.
-	compressedBuf := NewTempBuffer()
-	pb := CompressAndChecksum(&compressedBuf.b, blockData, blockKind, compressor, checksummer)
-	return pb, compressedBuf
-}
-
-// CopyAndChecksumToTempBuffer copies (without compressing) and checksums
-// the provided block into a TempBuffer. The caller should Release() the
-// TempBuffer once it is no longer necessary.
-func CopyAndChecksumToTempBuffer(
-	blockData []byte, blockKind Kind, compressor *Compressor, checksummer *Checksummer,
-) (PhysicalBlock, *TempBuffer) {
-	// Grab a buffer to use as the destination for compression.
-	compressedBuf := NewTempBuffer()
-	pb := CopyAndChecksum(&compressedBuf.b, blockData, blockKind, compressor, checksummer)
-	return pb, compressedBuf
-}
-
-// TempBuffer is a buffer that is used temporarily and is released back to a
-// pool for reuse.
-type TempBuffer struct {
-	b []byte
-}
-
-// NewTempBuffer returns a TempBuffer from the pool. The buffer will have zero
-// size and length and arbitrary capacity.
-func NewTempBuffer() *TempBuffer {
-	tb := tempBufferPool.Get().(*TempBuffer)
-	if invariants.Enabled && len(tb.b) > 0 {
-		panic("NewTempBuffer length not 0")
-	}
-	return tb
-}
-
-// Data returns the byte slice currently backing the Buffer.
-func (tb *TempBuffer) Data() []byte {
-	return tb.b
-}
-
-// Size returns the current size of the buffer.
-func (tb *TempBuffer) Size() int {
-	return len(tb.b)
-}
-
-// Append appends the contents of v to the buffer, growing the buffer if
-// necessary. Returns the offset at which it was appended.
-func (tb *TempBuffer) Append(v []byte) (startOffset int) {
-	startOffset = len(tb.b)
-	tb.b = append(tb.b, v...)
-	return startOffset
-}
-
-// Resize resizes the buffer to the specified length, allocating if necessary.
-// If the length is longer than the current length, the values of the new bytes
-// are arbitrary.
-func (tb *TempBuffer) Resize(length int) {
-	if length > cap(tb.b) {
-		tb.b = slices.Grow(tb.b, length-len(tb.b))
-	}
-	tb.b = tb.b[:length]
-}
-
-// Reset is equivalent to Resize(0).
-func (tb *TempBuffer) Reset() {
-	tb.b = tb.b[:0]
-}
-
-// Release releases the buffer back to the pool for reuse.
-func (tb *TempBuffer) Release() {
-	// Note we avoid releasing buffers that are larger than the configured
-	// maximum to the pool. This avoids holding on to occasional large buffers
-	// necessary for e.g. singular large values.
-	if tb.b != nil && len(tb.b) < tempBufferMaxReusedSize {
-		invariants.MaybeMangle(tb.b)
-		tb.b = tb.b[:0]
-		tempBufferPool.Put(tb)
-	}
-}
-
-// tempBufferPool is a pool of buffers that are used to temporarily hold either
-// compressed or uncompressed block data.
-var tempBufferPool = sync.Pool{
-	New: func() any {
-		return &TempBuffer{b: make([]byte, 0, tempBufferInitialSize)}
-	},
-}
-
-const tempBufferInitialSize = 32 * 1024
-const tempBufferMaxReusedSize = 256 * 1024

sstable/block/compression_test.go

@@ -19,10 +19,9 @@ func TestBufferRandomized(t *testing.T) {
 	t.Logf("seed %d", seed)
 	rng := rand.New(rand.NewPCG(0, seed))
 
-	compressor := MakeCompressor(SnappyCompression)
-	defer compressor.Close()
-	var checksummer Checksummer
-	checksummer.Init(ChecksumTypeCRC32c)
+	var physBlockMaker PhysicalBlockMaker
+	physBlockMaker.Init(SnappyCompression, ChecksumTypeCRC32c)
+	defer physBlockMaker.Close()
 	b := NewTempBuffer()
 	defer b.Release()
 	vbuf := make([]byte, 0, 1<<10) // 1 KiB
@@ -53,7 +52,7 @@ func TestBufferRandomized(t *testing.T) {
 				s := b.Data()
 				require.Equal(t, vbuf, s[len(s)-len(vbuf):])
 			}
-			_, bh := CompressAndChecksumToTempBuffer(b.Data(), blockkind.SSTableData, &compressor, &checksummer)
+			bh := physBlockMaker.Make(b.Data(), blockkind.SSTableData, NoFlags)
 			b.Reset()
 			bh.Release()
 		})