cubefs

1
package sarama
2

3
import (
4
	"sync"
5

6
	"github.com/klauspost/compress/zstd"
7
)
8

9
type ZstdEncoderParams struct {
10
	Level int
11
}
12
type ZstdDecoderParams struct {
13
}
14

15
var zstdEncMap, zstdDecMap sync.Map
16

17
func getEncoder(params ZstdEncoderParams) *zstd.Encoder {
18
	if ret, ok := zstdEncMap.Load(params); ok {
19
		return ret.(*zstd.Encoder)
20
	}
21
	// It's possible to race and create multiple new writers.
22
	// Only one will survive GC after use.
23
	encoderLevel := zstd.SpeedDefault
24
	if params.Level != CompressionLevelDefault {
25
		encoderLevel = zstd.EncoderLevelFromZstd(params.Level)
26
	}
27
	zstdEnc, _ := zstd.NewWriter(nil, zstd.WithZeroFrames(true),
28
		zstd.WithEncoderLevel(encoderLevel))
29
	zstdEncMap.Store(params, zstdEnc)
30
	return zstdEnc
31
}
32

33
func getDecoder(params ZstdDecoderParams) *zstd.Decoder {
34
	if ret, ok := zstdDecMap.Load(params); ok {
35
		return ret.(*zstd.Decoder)
36
	}
37
	// It's possible to race and create multiple new readers.
38
	// Only one will survive GC after use.
39
	zstdDec, _ := zstd.NewReader(nil)
40
	zstdDecMap.Store(params, zstdDec)
41
	return zstdDec
42
}
43

44
func zstdDecompress(params ZstdDecoderParams, dst, src []byte) ([]byte, error) {
45
	return getDecoder(params).DecodeAll(src, dst)
46
}
47

48
func zstdCompress(params ZstdEncoderParams, dst, src []byte) ([]byte, error) {
49
	return getEncoder(params).EncodeAll(src, dst), nil
50
}
51