sstable: use zstd context to reduce compression allocations

Use zstd context to reduce zstd allocations.

Author: EdwardX29

go.mod

@@ -1,7 +1,7 @@
 module github.com/cockroachdb/pebble
 
 require (
-	github.com/DataDog/zstd v1.5.6-0.20230824185856-869dae002e5e
+	github.com/DataDog/zstd v1.5.7
 	github.com/HdrHistogram/hdrhistogram-go v1.1.2
 	github.com/cespare/xxhash/v2 v2.2.0
 	github.com/cockroachdb/crlib v0.0.0-20241112164430-1264a2edc35b

go.sum

@@ -4,8 +4,8 @@ dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7
 gioui.org v0.0.0-20210308172011-57750fc8a0a6/go.mod h1:RSH6KIUZ0p2xy5zHDxgAM4zumjgTw83q2ge/PI+yyw8=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
-github.com/DataDog/zstd v1.5.6-0.20230824185856-869dae002e5e h1:ZIWapoIRN1VqT8GR8jAwb1Ie9GyehWjVcGh32Y2MznE=
-github.com/DataDog/zstd v1.5.6-0.20230824185856-869dae002e5e/go.mod h1:g4AWEaM3yOg3HYfnJ3YIawPnVdXJh9QME85blwSAmyw=
+github.com/DataDog/zstd v1.5.7 h1:ybO8RBeh29qrxIhCA9E8gKY6xfONU9T6G6aP9DTKfLE=
+github.com/DataDog/zstd v1.5.7/go.mod h1:g4AWEaM3yOg3HYfnJ3YIawPnVdXJh9QME85blwSAmyw=
 github.com/GoogleCloudPlatform/cloudsql-proxy v0.0.0-20190129172621-c8b1d7a94ddf/go.mod h1:aJ4qN3TfrelA6NZ6AXsXRfmEVaYin3EDbSPJrKS8OXo=
 github.com/HdrHistogram/hdrhistogram-go v1.1.2 h1:5IcZpTvzydCQeHzK4Ef/D5rrSqwxob0t8PQPMybUNFM=
 github.com/HdrHistogram/hdrhistogram-go v1.1.2/go.mod h1:yDgFjdqOqDEKOvasDdhWNXYg9BVp4O+o5f6V/ehm6Oo=

sstable/block/compression.go

@@ -115,6 +115,7 @@ func (i CompressionIndicator) String() string {
 // payload.
 func DecompressedLen(algo CompressionIndicator, b []byte) (decompressedLen int, err error) {
 	decompressor := GetDecompressor(algo)
+	defer decompressor.Close()
 	return decompressor.DecompressedLen(b)
 }
 
@@ -123,6 +124,7 @@ func DecompressedLen(algo CompressionIndicator, b []byte) (decompressedLen int,
 // determine the correct size.
 func DecompressInto(algo CompressionIndicator, compressed []byte, buf []byte) error {
 	decompressor := GetDecompressor(algo)
+	defer decompressor.Close()
 	err := decompressor.DecompressInto(buf, compressed)
 	if err != nil {
 		return base.MarkCorruptionError(err)
@@ -214,6 +216,7 @@ func CompressAndChecksum(
 	algo := NoCompressionIndicator
 	if compression != NoCompression {
 		compressor := GetCompressor(compression)
+		defer compressor.Close()
 		algo, buf = compressor.Compress(buf, blockData)
 		if len(buf) >= len(blockData)-len(blockData)/8 {
 			algo = NoCompressionIndicator

sstable/block/compression_test.go

@@ -36,6 +36,7 @@ func TestCompressionRoundtrip(t *testing.T) {
 			// not sufficient, Compress should allocate one that is.
 			compressedBuf := make([]byte, 1+rng.IntN(1<<10 /* 1 KiB */))
 			compressor := GetCompressor(compression)
+			defer compressor.Close()
 			btyp, compressed := compressor.Compress(compressedBuf, payload)
 			v, err := decompress(btyp, compressed)
 			require.NoError(t, err)

sstable/block/compressor.go

@@ -10,33 +10,38 @@ import (
 
 type Compressor interface {
 	Compress(dst, src []byte) (CompressionIndicator, []byte)
+
+	// Close must be called when the Compressor is no longer needed.
+	// After Close is called, the Compressor must not be used again.
+	Close()
 }
 
 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 (noopCompressor) Close() {}
 
 func (snappyCompressor) Compress(dst, src []byte) (CompressionIndicator, []byte) {
 	dst = dst[:cap(dst):cap(dst)]
 	return SnappyCompressionIndicator, snappy.Encode(dst, src)
 }
 
+func (snappyCompressor) Close() {}
+
 func GetCompressor(c Compression) Compressor {
 	switch c {
 	case NoCompression:
 		return noopCompressor{}
 	case SnappyCompression:
 		return snappyCompressor{}
 	case ZstdCompression:
-		return zstdCompressor{}
+		return getZstdCompressor()
 	default:
 		panic("Invalid compression type.")
 	}
@@ -52,15 +57,17 @@ type Decompressor interface {
 	// allowing the caller to allocate a buffer exactly sized to the decompressed
 	// payload.
 	DecompressedLen(b []byte) (decompressedLen int, err error)
+
+	// Close must be called when the Decompressor is no longer needed.
+	// After Close is called, the Decompressor must not be used again.
+	Close()
 }
 
 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)]
@@ -72,6 +79,8 @@ func (noopDecompressor) DecompressedLen(b []byte) (decompressedLen int, err erro
 	return len(b), nil
 }
 
+func (noopDecompressor) Close() {}
+
 func (snappyDecompressor) DecompressInto(buf, compressed []byte) error {
 	result, err := snappy.Decode(buf, compressed)
 	if err != nil {
@@ -89,6 +98,8 @@ func (snappyDecompressor) DecompressedLen(b []byte) (decompressedLen int, err er
 	return l, err
 }
 
+func (snappyDecompressor) Close() {}
+
 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.
@@ -106,7 +117,7 @@ func GetDecompressor(c CompressionIndicator) Decompressor {
 	case SnappyCompressionIndicator:
 		return snappyDecompressor{}
 	case ZstdCompressionIndicator:
-		return zstdDecompressor{}
+		return getZstdDecompressor()
 	default:
 		panic("Invalid compression type.")
 	}

sstable/block/compressor_cgo.go

@@ -3,13 +3,25 @@
 package block
 
 import (
-	"bytes"
 	"encoding/binary"
+	"sync"
 
 	"github.com/DataDog/zstd"
 	"github.com/cockroachdb/errors"
 )
 
+type zstdCompressor struct {
+	ctx zstd.Ctx
+}
+
+var _ Compressor = (*zstdCompressor)(nil)
+
+var zstdCompressorPool = sync.Pool{
+	New: func() any {
+		return &zstdCompressor{ctx: zstd.NewCtx()}
+	},
+}
+
 // UseStandardZstdLib indicates whether the zstd implementation is a port of the
 // official one in the facebook/zstd repository.
 //
@@ -26,21 +38,48 @@ const UseStandardZstdLib = true
 // is sufficient. The subslice `compressedBuf[:varIntLen]` should already encode
 // the length of `b` before calling Compress. It returns the encoded byte
 // slice, including the `compressedBuf[:varIntLen]` prefix.
-func (zstdCompressor) Compress(compressedBuf []byte, b []byte) (CompressionIndicator, []byte) {
+func (z *zstdCompressor) Compress(compressedBuf []byte, b []byte) (CompressionIndicator, []byte) {
 	if len(compressedBuf) < binary.MaxVarintLen64 {
 		compressedBuf = append(compressedBuf, make([]byte, binary.MaxVarintLen64-len(compressedBuf))...)
 	}
+
+	// Get the bound and allocate the proper amount of memory instead of relying on
+	// Datadog/zstd to do it for us. This allows us to avoid memcopying data around
+	// for the varIntLen prefix.
+	bound := zstd.CompressBound(len(b))
+	if cap(compressedBuf) < binary.MaxVarintLen64+bound {
+		compressedBuf = make([]byte, binary.MaxVarintLen64, binary.MaxVarintLen64+bound)
+	}
+
 	varIntLen := binary.PutUvarint(compressedBuf, uint64(len(b)))
-	buf := bytes.NewBuffer(compressedBuf[:varIntLen])
-	writer := zstd.NewWriterLevel(buf, 3)
-	writer.Write(b)
-	writer.Close()
-	return ZstdCompressionIndicator, buf.Bytes()
+	result, err := z.ctx.CompressLevel(compressedBuf[varIntLen:varIntLen+bound], b, 3)
+	if err != nil {
+		panic("Error while compressing using Zstd.")
+	}
+	if &result[0] != &compressedBuf[varIntLen] {
+		panic("Allocated a new buffer despite checking CompressBound.")
+	}
+
+	return ZstdCompressionIndicator, compressedBuf[:varIntLen+len(result)]
 }
 
+func (z *zstdCompressor) Close() {
+	zstdCompressorPool.Put(z)
+}
+
+func getZstdCompressor() *zstdCompressor {
+	return zstdCompressorPool.Get().(*zstdCompressor)
+}
+
+type zstdDecompressor struct {
+	ctx zstd.Ctx
+}
+
+var _ Decompressor = (*zstdDecompressor)(nil)
+
 // DecompressInto decompresses src with the Zstandard algorithm. The destination
-// buffer must already be sufficiently sized, otherwise Decompress may error.
-func (zstdDecompressor) DecompressInto(dst, src []byte) error {
+// buffer must already be sufficiently sized, otherwise DecompressInto may error.
+func (z *zstdDecompressor) DecompressInto(dst, src []byte) error {
 	// The payload is prefixed with a varint encoding the length of
 	// the decompressed block.
 	_, prefixLen := binary.Uvarint(src)
@@ -51,9 +90,23 @@ func (zstdDecompressor) DecompressInto(dst, src []byte) error {
 	if len(dst) == 0 {
 		return errors.Errorf("decodeZstd: empty dst buffer")
 	}
-	_, err := zstd.DecompressInto(dst, src)
+	_, err := z.ctx.DecompressInto(dst, src)
 	if err != nil {
 		return err
 	}
 	return nil
 }
+
+func (z *zstdDecompressor) Close() {
+	zstdDecompressorPool.Put(z)
+}
+
+var zstdDecompressorPool = sync.Pool{
+	New: func() any {
+		return &zstdDecompressor{ctx: zstd.NewCtx()}
+	},
+}
+
+func getZstdDecompressor() *zstdDecompressor {
+	return zstdDecompressorPool.Get().(*zstdDecompressor)
+}

sstable/block/compressor_nocgo.go

@@ -14,6 +14,10 @@ import (
 	"github.com/klauspost/compress/zstd"
 )
 
+type zstdCompressor struct{}
+
+var _ Compressor = zstdCompressor{}
+
 // UseStandardZstdLib indicates whether the zstd implementation is a port of the
 // official one in the facebook/zstd repository.
 //
@@ -40,8 +44,18 @@ func (zstdCompressor) Compress(compressedBuf, b []byte) (CompressionIndicator, [
 	return ZstdCompressionIndicator, encoder.EncodeAll(b, compressedBuf[:varIntLen])
 }
 
+func (zstdCompressor) Close() {}
+
+func getZstdCompressor() zstdCompressor {
+	return zstdCompressor{}
+}
+
+type zstdDecompressor struct{}
+
+var _ Decompressor = zstdDecompressor{}
+
 // Decompress decompresses src with the Zstandard algorithm. The destination
-// buffer must already be sufficiently sized, otherwise Decompress may error.
+// buffer must already be sufficiently sized, otherwise DecompressInto may error.
 func (zstdDecompressor) DecompressInto(dst, src []byte) error {
 	// The payload is prefixed with a varint encoding the length of
 	// the decompressed block.
@@ -59,3 +73,9 @@ func (zstdDecompressor) DecompressInto(dst, src []byte) error {
 	}
 	return nil
 }
+
+func (zstdDecompressor) Close() {}
+
+func getZstdDecompressor() zstdDecompressor {
+	return zstdDecompressor{}
+}