forked from streamingfast/kvdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
compression.go
93 lines (74 loc) · 2.1 KB
/
compression.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
package store
import (
"bytes"
"fmt"
"github.com/klauspost/compress/zstd"
"go.uber.org/zap/zapcore"
)
type Compressor interface {
Compress(in []byte) []byte
Decompress(in []byte) ([]byte, error)
zapcore.ObjectMarshaler
}
func NewCompressor(mode string, thresholdInBytes int) (Compressor, error) {
switch mode {
case "zst", "zstd":
return NewZstdCompressor(thresholdInBytes), nil
case "", "none", "false", "no":
return NewNoOpCompressor(), nil
default:
return nil, fmt.Errorf("invalid compression value, use '' or zstd (for legacy support) or 'none'")
}
}
type NoOpCompressor struct{}
func NewNoOpCompressor() *NoOpCompressor {
return &NoOpCompressor{}
}
func (NoOpCompressor) Compress(in []byte) []byte {
return in
}
func (NoOpCompressor) Decompress(in []byte) ([]byte, error) {
return in, nil
}
func (NoOpCompressor) MarshalLogObject(enc zapcore.ObjectEncoder) error {
enc.AddString("compression", "none")
return nil
}
type ZstdCompressor struct {
enc *zstd.Encoder
dec *zstd.Decoder
thresholdInBytes int
}
func NewZstdCompressor(thresholdInBytes int) *ZstdCompressor {
// There can be errors only when using `opts` parameters, so it's safe to ignore them here
enc, _ := zstd.NewWriter(nil)
dec, _ := zstd.NewReader(nil)
return &ZstdCompressor{
enc: enc,
dec: dec,
thresholdInBytes: thresholdInBytes,
}
}
func (c *ZstdCompressor) Compress(in []byte) (out []byte) {
if len(in) > c.thresholdInBytes {
return c.enc.EncodeAll(in, out)
}
return in
}
var zstdMagicBytes = []byte{0x28, 0xB5, 0x2F, 0xFD}
func (c *ZstdCompressor) Decompress(in []byte) ([]byte, error) {
if bytes.HasPrefix(in, zstdMagicBytes) {
// We pre-allocate a bit the array to reduce allocation, at least the compression size is a good start
buf, err := c.dec.DecodeAll(in, make([]byte, 0, len(in)))
if err != nil {
return nil, err
}
return buf, nil
}
return in, nil
}
func (c *ZstdCompressor) MarshalLogObject(enc zapcore.ObjectEncoder) error {
enc.AddString("compression", "zstd")
enc.AddInt("compression_size_threshold", c.thresholdInBytes)
return nil
}