Merge pull request IBM#2045 from honeycombio/lizf.zstd-compress-options

[zstd] pass level param through to compress/zstd encoder

compress.go

@@ -187,7 +187,7 @@ func compress(cc CompressionCodec, level int, data []byte) ([]byte, error) {
 		}
 		return buf.Bytes(), nil
 	case CompressionZSTD:
-		return zstdCompress(nil, data)
+		return zstdCompress(ZstdEncoderParams{level}, nil, data)
 	default:
 		return nil, PacketEncodingError{fmt.Sprintf("unsupported compression codec (%d)", cc)}
 	}

decompress.go

@@ -54,7 +54,7 @@ func decompress(cc CompressionCodec, data []byte) ([]byte, error) {
 
 		return io.ReadAll(reader)
 	case CompressionZSTD:
-		return zstdDecompress(nil, data)
+		return zstdDecompress(ZstdDecoderParams{}, nil, data)
 	default:
 		return nil, PacketDecodingError{fmt.Sprintf("invalid compression specified (%d)", cc)}
 	}

zstd.go

@@ -1,18 +1,50 @@
 package sarama
 
 import (
+	"sync"
+
 	"github.com/klauspost/compress/zstd"
 )
 
-var (
-	zstdDec, _ = zstd.NewReader(nil)
-	zstdEnc, _ = zstd.NewWriter(nil, zstd.WithZeroFrames(true))
-)
+type ZstdEncoderParams struct {
+	Level int
+}
+type ZstdDecoderParams struct {
+}
+
+var zstdEncMap, zstdDecMap sync.Map
+
+func getEncoder(params ZstdEncoderParams) *zstd.Encoder {
+	if ret, ok := zstdEncMap.Load(params); ok {
+		return ret.(*zstd.Encoder)
+	}
+	// It's possible to race and create multiple new writers.
+	// Only one will survive GC after use.
+	encoderLevel := zstd.SpeedDefault
+	if params.Level != CompressionLevelDefault {
+		encoderLevel = zstd.EncoderLevelFromZstd(params.Level)
+	}
+	zstdEnc, _ := zstd.NewWriter(nil, zstd.WithZeroFrames(true),
+		zstd.WithEncoderLevel(encoderLevel))
+	zstdEncMap.Store(params, zstdEnc)
+	return zstdEnc
+}
+
+func getDecoder(params ZstdDecoderParams) *zstd.Decoder {
+	if ret, ok := zstdDecMap.Load(params); ok {
+		return ret.(*zstd.Decoder)
+	}
+	// It's possible to race and create multiple new readers.
+	// Only one will survive GC after use.
+	zstdDec, _ := zstd.NewReader(nil)
+	zstdDecMap.Store(params, zstdDec)
+	return zstdDec
+}
 
-func zstdDecompress(dst, src []byte) ([]byte, error) {
-	return zstdDec.DecodeAll(src, dst)
+func zstdDecompress(params ZstdDecoderParams, dst, src []byte) ([]byte, error) {
+	return getDecoder(params).DecodeAll(src, dst)
 }
 
-func zstdCompress(dst, src []byte) ([]byte, error) {
-	return zstdEnc.EncodeAll(src, dst), nil
+func zstdCompress(params ZstdEncoderParams, dst, src []byte) ([]byte, error) {
+	return getEncoder(params).EncodeAll(src, dst), nil
 }