Add ability to create zstd compressor with compression level (#4203)

Author: DracoLi

utils/compression/compressor_test.go

@@ -9,6 +9,7 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/DataDog/zstd"
 	"github.com/stretchr/testify/require"
 
 	_ "embed"
@@ -149,6 +150,19 @@ func TestNewCompressorWithInvalidLimit(t *testing.T) {
 	}
 }
 
+func TestNewZstdCompressorWithLevel(t *testing.T) {
+	compressor, err := NewZstdCompressorWithLevel(maxMessageSize, zstd.BestSpeed)
+	require.NoError(t, err)
+
+	data := utils.RandomBytes(4096)
+	compressed, err := compressor.Compress(data)
+	require.NoError(t, err)
+
+	decompressed, err := compressor.Decompress(compressed)
+	require.NoError(t, err)
+	require.Equal(t, data, decompressed)
+}
+
 func FuzzZstdCompressor(f *testing.F) {
 	fuzzHelper(f, TypeZstd)
 }

utils/compression/zstd_compressor.go

@@ -22,28 +22,33 @@ var (
 )
 
 func NewZstdCompressor(maxSize int64) (Compressor, error) {
+	return NewZstdCompressorWithLevel(maxSize, zstd.DefaultCompression)
+}
+
+func NewZstdCompressorWithLevel(maxSize int64, level int) (Compressor, error) {
 	if maxSize == math.MaxInt64 {
 		// "Decompress" creates "io.LimitReader" with max size + 1:
 		// if the max size + 1 overflows, "io.LimitReader" reads nothing
 		// returning 0 byte for the decompress call
 		// require max size < math.MaxInt64 to prevent int64 overflows
 		return nil, ErrInvalidMaxSizeCompressor
 	}
-
 	return &zstdCompressor{
 		maxSize: maxSize,
+		level:   level,
 	}, nil
 }
 
 type zstdCompressor struct {
 	maxSize int64
+	level   int
 }
 
 func (z *zstdCompressor) Compress(msg []byte) ([]byte, error) {
 	if int64(len(msg)) > z.maxSize {
 		return nil, fmt.Errorf("%w: (%d) > (%d)", ErrMsgTooLarge, len(msg), z.maxSize)
 	}
-	return zstd.Compress(nil, msg)
+	return zstd.CompressLevel(nil, msg, z.level)
 }
 
 func (z *zstdCompressor) Decompress(msg []byte) ([]byte, error) {