sstable: add Compressor and Decompressor interfaces

add Compressor and Decompressor interfaces to refactor existing compressor and decompressor schemes.

Author: EdwardX29

sstable/block/block.go

@@ -537,13 +537,13 @@ func (r *Reader) doRead(
 		decompressed = compressed
 	} else {
 		// Decode the length of the decompressed value.
-		decodedLen, prefixLen, err := DecompressedLen(typ, compressed.BlockData())
+		decodedLen, err := DecompressedLen(typ, compressed.BlockData())
 		if err != nil {
 			compressed.Release()
 			return Value{}, err
 		}
 		decompressed = Alloc(decodedLen, env.BufferPool)
-		err = DecompressInto(typ, compressed.BlockData()[prefixLen:], decompressed.BlockData())
+		err = DecompressInto(typ, compressed.BlockData(), decompressed.BlockData())
 		compressed.Release()
 		if err != nil {
 			decompressed.Release()

sstable/block/compression.go

@@ -5,7 +5,6 @@
 package block
 
 import (
-	"encoding/binary"
 	"math/rand/v2"
 	"slices"
 	"sync"
@@ -15,7 +14,6 @@ import (
 	"github.com/cockroachdb/pebble/internal/bytealloc"
 	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/objstorage"
-	"github.com/golang/snappy"
 )
 
 // Compression is the per-block compression algorithm to use.
@@ -114,55 +112,21 @@ func (i CompressionIndicator) String() string {
 
 // DecompressedLen returns the length of the provided block once decompressed,
 // allowing the caller to allocate a buffer exactly sized to the decompressed
-// payload. For some compression algorithms, the payload is prefixed with a
-// varint encoding the length of the decompressed block. In such cases, a
-// non-zero prefixLength is returned indicating the length of this prefix.
-func DecompressedLen(
-	algo CompressionIndicator, b []byte,
-) (decompressedLen int, prefixLength int, err error) {
-	switch algo {
-	case NoCompressionIndicator:
-		return len(b), 0, nil
-	case SnappyCompressionIndicator:
-		l, err := snappy.DecodedLen(b)
-		return l, 0, err
-	case ZstdCompressionIndicator:
-		// This will also be used by zlib, bzip2 and lz4 to retrieve the decodedLen
-		// if we implement these algorithms in the future.
-		decodedLenU64, varIntLen := binary.Uvarint(b)
-		if varIntLen <= 0 {
-			return 0, 0, base.CorruptionErrorf("pebble/table: compression block has invalid length")
-		}
-		return int(decodedLenU64), varIntLen, nil
-	default:
-		return 0, 0, base.CorruptionErrorf("pebble/table: unknown block compression: %d", errors.Safe(algo))
-	}
+// payload.
+func DecompressedLen(algo CompressionIndicator, b []byte) (decompressedLen int, err error) {
+	decompressor := GetDecompressor(algo)
+	return decompressor.DecompressedLen(b)
 }
 
 // DecompressInto decompresses compressed into buf. The buf slice must have the
 // exact size as the decompressed value. Callers may use DecompressedLen to
 // determine the correct size.
 func DecompressInto(algo CompressionIndicator, compressed []byte, buf []byte) error {
-	var result []byte
-	var err error
-	switch algo {
-	case NoCompressionIndicator:
-		result = buf[:len(compressed)]
-		copy(result, compressed)
-	case SnappyCompressionIndicator:
-		result, err = snappy.Decode(buf, compressed)
-	case ZstdCompressionIndicator:
-		result, err = decodeZstd(buf, compressed)
-	default:
-		return base.CorruptionErrorf("pebble/table: unknown block compression: %d", errors.Safe(algo))
-	}
+	decompressor := GetDecompressor(algo)
+	err := decompressor.DecompressInto(buf, compressed)
 	if err != nil {
 		return base.MarkCorruptionError(err)
 	}
-	if len(result) != len(buf) || (len(result) > 0 && &result[0] != &buf[0]) {
-		return base.CorruptionErrorf("pebble/table: decompressed into unexpected buffer: %p != %p",
-			errors.Safe(result), errors.Safe(buf))
-	}
 	return nil
 }
 
@@ -249,7 +213,8 @@ func CompressAndChecksum(
 	// least 12.5%.
 	algo := NoCompressionIndicator
 	if compression != NoCompression {
-		algo, buf = compress(compression, blockData, buf)
+		compressor := GetCompressor(compression)
+		algo, buf = compressor.Compress(buf, blockData)
 		if len(buf) >= len(blockData)-len(blockData)/8 {
 			algo = NoCompressionIndicator
 		}
@@ -270,30 +235,6 @@ func CompressAndChecksum(
 	return pb
 }
 
-// compress compresses a sstable block, using dstBuf as the desired destination.
-//
-// The result is aliased to dstBuf if that buffer had enough capacity, otherwise
-// it is a newly-allocated buffer.
-func compress(
-	compression Compression, b []byte, dstBuf []byte,
-) (indicator CompressionIndicator, compressed []byte) {
-	switch compression {
-	case SnappyCompression:
-		// snappy relies on the length of the buffer, and not the capacity to
-		// determine if it needs to make an allocation.
-		dstBuf = dstBuf[:cap(dstBuf):cap(dstBuf)]
-		return SnappyCompressionIndicator, snappy.Encode(dstBuf, b)
-	case ZstdCompression:
-		if len(dstBuf) < binary.MaxVarintLen64 {
-			dstBuf = append(dstBuf, make([]byte, binary.MaxVarintLen64-len(dstBuf))...)
-		}
-		varIntLen := binary.PutUvarint(dstBuf, uint64(len(b)))
-		return ZstdCompressionIndicator, encodeZstd(dstBuf, varIntLen, b)
-	default:
-		panic("unreachable")
-	}
-}
-
 // A Buffer is a buffer for encoding a block. The caller mutates the buffer to
 // construct the uncompressed block, and calls CompressAndChecksum to produce
 // the physical, possibly-compressed PhysicalBlock. A Buffer recycles byte

sstable/block/compression_cgo.go

@@ -35,8 +35,8 @@ func TestCompressionRoundtrip(t *testing.T) {
 			// Create a randomly-sized buffer to house the compressed output. If it's
 			// not sufficient, Compress should allocate one that is.
 			compressedBuf := make([]byte, 1+rng.IntN(1<<10 /* 1 KiB */))
-
-			btyp, compressed := compress(compression, payload, compressedBuf)
+			compressor := GetCompressor(compression)
+			btyp, compressed := compressor.Compress(compressedBuf, payload)
 			v, err := decompress(btyp, compressed)
 			require.NoError(t, err)
 			got := payload
@@ -78,12 +78,12 @@ func decompress(algo CompressionIndicator, b []byte) (*cache.Value, error) {
 	if algo == NoCompressionIndicator {
 		return nil, nil
 	}
+
 	// first obtain the decoded length.
-	decodedLen, prefixLen, err := DecompressedLen(algo, b)
+	decodedLen, err := DecompressedLen(algo, b)
 	if err != nil {
 		return nil, err
 	}
-	b = b[prefixLen:]
 	// Allocate sufficient space from the cache.
 	decoded := cache.Alloc(decodedLen)
 	decodedBuf := decoded.RawBuffer()

sstable/block/compression_nocgo.go

@@ -0,0 +1,113 @@
+package block
+
+import (
+	"encoding/binary"
+
+	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/golang/snappy"
+)
+
+type Compressor interface {
+	Compress(dst, src []byte) (CompressionIndicator, []byte)
+}
+
+type noopCompressor struct{}
+type snappyCompressor struct{}
+type zstdCompressor struct{}
+
+var _ Compressor = noopCompressor{}
+var _ Compressor = snappyCompressor{}
+var _ Compressor = zstdCompressor{}
+
+func (noopCompressor) Compress(dst, src []byte) (CompressionIndicator, []byte) {
+	panic("NoCompressionCompressor.Compress() should not be called.")
+}
+
+func (snappyCompressor) Compress(dst, src []byte) (CompressionIndicator, []byte) {
+	dst = dst[:cap(dst):cap(dst)]
+	return SnappyCompressionIndicator, snappy.Encode(dst, src)
+}
+
+func GetCompressor(c Compression) Compressor {
+	switch c {
+	case NoCompression:
+		return noopCompressor{}
+	case SnappyCompression:
+		return snappyCompressor{}
+	case ZstdCompression:
+		return zstdCompressor{}
+	default:
+		panic("Invalid compression type.")
+	}
+}
+
+type Decompressor interface {
+	// DecompressInto decompresses compressed into buf. The buf slice must have the
+	// exact size as the decompressed value. Callers may use DecompressedLen to
+	// determine the correct size.
+	DecompressInto(buf, compressed []byte) error
+
+	// DecompressedLen returns the length of the provided block once decompressed,
+	// allowing the caller to allocate a buffer exactly sized to the decompressed
+	// payload.
+	DecompressedLen(b []byte) (decompressedLen int, err error)
+}
+
+type noopDecompressor struct{}
+type snappyDecompressor struct{}
+type zstdDecompressor struct{}
+
+var _ Decompressor = noopDecompressor{}
+var _ Decompressor = snappyDecompressor{}
+var _ Decompressor = zstdDecompressor{}
+
+func (noopDecompressor) DecompressInto(dst, src []byte) error {
+	dst = dst[:len(src)]
+	copy(dst, src)
+	return nil
+}
+
+func (noopDecompressor) DecompressedLen(b []byte) (decompressedLen int, err error) {
+	return len(b), nil
+}
+
+func (snappyDecompressor) DecompressInto(buf, compressed []byte) error {
+	result, err := snappy.Decode(buf, compressed)
+	if err != nil {
+		return err
+	}
+	if len(result) != len(buf) || (len(result) > 0 && &result[0] != &buf[0]) {
+		return base.CorruptionErrorf("pebble/table: decompressed into unexpected buffer: %p != %p",
+			errors.Safe(result), errors.Safe(buf))
+	}
+	return nil
+}
+
+func (snappyDecompressor) DecompressedLen(b []byte) (decompressedLen int, err error) {
+	l, err := snappy.DecodedLen(b)
+	return l, err
+}
+
+func (zstdDecompressor) DecompressedLen(b []byte) (decompressedLen int, err error) {
+	// This will also be used by zlib, bzip2 and lz4 to retrieve the decodedLen
+	// if we implement these algorithms in the future.
+	decodedLenU64, varIntLen := binary.Uvarint(b)
+	if varIntLen <= 0 {
+		return 0, base.CorruptionErrorf("pebble/table: compression block has invalid length")
+	}
+	return int(decodedLenU64), nil
+}
+
+func GetDecompressor(c CompressionIndicator) Decompressor {
+	switch c {
+	case NoCompressionIndicator:
+		return noopDecompressor{}
+	case SnappyCompressionIndicator:
+		return snappyDecompressor{}
+	case ZstdCompressionIndicator:
+		return zstdDecompressor{}
+	default:
+		panic("Invalid compression type.")
+	}
+}