blob: 9247c3553c72683ee2aed2884ee1c26adbeb5e2c [file] [log] [blame]
Dinesh Belwalkare63f7f92019-11-22 23:11:16 +00001package sarama
2
3import (
4 "bytes"
5 "compress/gzip"
6 "fmt"
7 "sync"
8
9 "github.com/eapache/go-xerial-snappy"
10 "github.com/pierrec/lz4"
11)
12
13var (
14 lz4WriterPool = sync.Pool{
15 New: func() interface{} {
16 return lz4.NewWriter(nil)
17 },
18 }
19
20 gzipWriterPool = sync.Pool{
21 New: func() interface{} {
22 return gzip.NewWriter(nil)
23 },
24 }
25)
26
27func compress(cc CompressionCodec, level int, data []byte) ([]byte, error) {
28 switch cc {
29 case CompressionNone:
30 return data, nil
31 case CompressionGZIP:
32 var (
33 err error
34 buf bytes.Buffer
35 writer *gzip.Writer
36 )
37 if level != CompressionLevelDefault {
38 writer, err = gzip.NewWriterLevel(&buf, level)
39 if err != nil {
40 return nil, err
41 }
42 } else {
43 writer = gzipWriterPool.Get().(*gzip.Writer)
44 defer gzipWriterPool.Put(writer)
45 writer.Reset(&buf)
46 }
47 if _, err := writer.Write(data); err != nil {
48 return nil, err
49 }
50 if err := writer.Close(); err != nil {
51 return nil, err
52 }
53 return buf.Bytes(), nil
54 case CompressionSnappy:
55 return snappy.Encode(data), nil
56 case CompressionLZ4:
57 writer := lz4WriterPool.Get().(*lz4.Writer)
58 defer lz4WriterPool.Put(writer)
59
60 var buf bytes.Buffer
61 writer.Reset(&buf)
62
63 if _, err := writer.Write(data); err != nil {
64 return nil, err
65 }
66 if err := writer.Close(); err != nil {
67 return nil, err
68 }
69 return buf.Bytes(), nil
70 case CompressionZSTD:
71 return zstdCompress(nil, data)
72 default:
73 return nil, PacketEncodingError{fmt.Sprintf("unsupported compression codec (%d)", cc)}
74 }
75}