Merge pull request #1185 from retailnext/reuse-compression-objects

Reuse compression objects
IBM · Dec 14, 2018 · 94536b3 · 94536b3
2 parents 97315fe + f352e5c
commit 94536b3
Show file tree

Hide file tree

Showing 4 changed files with 151 additions and 165 deletions.
diff --git a/compress.go b/compress.go
@@ -0,0 +1,75 @@
+package sarama
+
+import (
+ "bytes"
+ "compress/gzip"
+ "fmt"
+ "sync"
+
+ "github.com/eapache/go-xerial-snappy"
+ "github.com/pierrec/lz4"
+)
+
+var (
+ lz4WriterPool = sync.Pool{
+ New: func() interface{} {
+ return lz4.NewWriter(nil)
+ },
+ }
+
+ gzipWriterPool = sync.Pool{
+ New: func() interface{} {
+ return gzip.NewWriter(nil)
+ },
+ }
+)
+
+func compress(cc CompressionCodec, level int, data []byte) ([]byte, error) {
+ switch cc {
+ case CompressionNone:
+ return data, nil
+ case CompressionGZIP:
+ var (
+ err error
+ buf bytes.Buffer
+ writer *gzip.Writer
+ )
+ if level != CompressionLevelDefault {
+ writer, err = gzip.NewWriterLevel(&buf, level)
+ if err != nil {
+ return nil, err
+ }
+ } else {
+ writer = gzipWriterPool.Get().(*gzip.Writer)
+ defer gzipWriterPool.Put(writer)
+ writer.Reset(&buf)
+ }
+ if _, err := writer.Write(data); err != nil {
+ return nil, err
+ }
+ if err := writer.Close(); err != nil {
+ return nil, err
+ }
+ return buf.Bytes(), nil
+ case CompressionSnappy:
+ return snappy.Encode(data), nil
+ case CompressionLZ4:
+ writer := lz4WriterPool.Get().(*lz4.Writer)
+ defer lz4WriterPool.Put(writer)
+
+ var buf bytes.Buffer
+ writer.Reset(&buf)
+
+ if _, err := writer.Write(data); err != nil {
+ return nil, err
+ }
+ if err := writer.Close(); err != nil {
+ return nil, err
+ }
+ return buf.Bytes(), nil
+ case CompressionZSTD:
+ return zstdCompressLevel(nil, data, level)
+ default:
+ return nil, PacketEncodingError{fmt.Sprintf("unsupported compression codec (%d)", cc)}
+ }
+}
diff --git a/decompress.go b/decompress.go
@@ -0,0 +1,63 @@
+package sarama
+
+import (
+ "bytes"
+ "compress/gzip"
+ "fmt"
+ "io/ioutil"
+ "sync"
+
+ "github.com/eapache/go-xerial-snappy"
+ "github.com/pierrec/lz4"
+)
+
+var (
+ lz4ReaderPool = sync.Pool{
+ New: func() interface{} {
+ return lz4.NewReader(nil)
+ },
+ }
+
+ gzipReaderPool sync.Pool
+)
+
+func decompress(cc CompressionCodec, data []byte) ([]byte, error) {
+ switch cc {
+ case CompressionNone:
+ return data, nil
+ case CompressionGZIP:
+ var (
+ err error
+ reader *gzip.Reader
+ readerIntf = gzipReaderPool.Get()
+ )
+ if readerIntf != nil {
+ reader = readerIntf.(*gzip.Reader)
+ } else {
+ reader, err = gzip.NewReader(bytes.NewReader(data))
+ if err != nil {
+ return nil, err
+ }
+ }
+
+ defer gzipReaderPool.Put(reader)
+
+ if err := reader.Reset(bytes.NewReader(data)); err != nil {
+ return nil, err
+ }
+
+ return ioutil.ReadAll(reader)
+ case CompressionSnappy:
+ return snappy.Decode(data)
+ case CompressionLZ4:
+ reader := lz4ReaderPool.Get().(*lz4.Reader)
+ defer lz4ReaderPool.Put(reader)
+
+ reader.Reset(bytes.NewReader(data))
+ return ioutil.ReadAll(reader)
+ case CompressionZSTD:
+ return zstdDecompress(nil, data)
+ default:
+ return nil, PacketDecodingError{fmt.Sprintf("invalid compression specified (%d)", cc)}
+ }
+}
diff --git a/message.go b/message.go
@@ -1,14 +1,8 @@
 package sarama
 
 import (
- "bytes"
- "compress/gzip"
  "fmt"
- "io/ioutil"
  "time"
-
- "github.com/eapache/go-xerial-snappy"
- "github.com/pierrec/lz4"
 )
 
 // CompressionCodec represents the various compression codecs recognized by Kafka in messages.
@@ -77,53 +71,12 @@ func (m *Message) encode(pe packetEncoder) error {
  payload = m.compressedCache
  m.compressedCache = nil
  } else if m.Value != nil {
- switch m.Codec {
- case CompressionNone:
- payload = m.Value
- case CompressionGZIP:
- var buf bytes.Buffer
- var writer *gzip.Writer
- if m.CompressionLevel != CompressionLevelDefault {
- writer, err = gzip.NewWriterLevel(&buf, m.CompressionLevel)
- if err != nil {
- return err
- }
- } else {
- writer = gzip.NewWriter(&buf)
- }
- if _, err = writer.Write(m.Value); err != nil {
- return err
- }
- if err = writer.Close(); err != nil {
- return err
- }
- m.compressedCache = buf.Bytes()
- payload = m.compressedCache
- case CompressionSnappy:
- tmp := snappy.Encode(m.Value)
- m.compressedCache = tmp
- payload = m.compressedCache
- case CompressionLZ4:
- var buf bytes.Buffer
- writer := lz4.NewWriter(&buf)
- if _, err = writer.Write(m.Value); err != nil {
- return err
- }
- if err = writer.Close(); err != nil {
- return err
- }
- m.compressedCache = buf.Bytes()
- payload = m.compressedCache
- case CompressionZSTD:
- c, err := zstdCompressLevel(nil, m.Value, m.CompressionLevel)
- if err != nil {
- return err
- }
- m.compressedCache = c
- payload = m.compressedCache
- default:
- return PacketEncodingError{fmt.Sprintf("unsupported compression codec (%d)", m.Codec)}
+
+ payload, err = compress(m.Codec, m.CompressionLevel, m.Value)
+ if err != nil {
+ return err
  }
+ m.compressedCache = payload
  // Keep in mind the compressed payload size for metric gathering
  m.compressedSize = len(payload)
  }
@@ -179,53 +132,18 @@ func (m *Message) decode(pd packetDecoder) (err error) {
  switch m.Codec {
  case CompressionNone:
  // nothing to do
- case CompressionGZIP:
+ default:
  if m.Value == nil {
  break
  }
- reader, err := gzip.NewReader(bytes.NewReader(m.Value))
+
+ m.Value, err = decompress(m.Codec, m.Value)
  if err != nil {
  return err
  }
- if m.Value, err = ioutil.ReadAll(reader); err != nil {
- return err
- }
  if err := m.decodeSet(); err != nil {
  return err
  }
- case CompressionSnappy:
- if m.Value == nil {
- break
- }
- if m.Value, err = snappy.Decode(m.Value); err != nil {
- return err
- }
- if err := m.decodeSet(); err != nil {
- return err
- }
- case CompressionLZ4:
- if m.Value == nil {
- break
- }
- reader := lz4.NewReader(bytes.NewReader(m.Value))
- if m.Value, err = ioutil.ReadAll(reader); err != nil {
- return err
- }
- if err := m.decodeSet(); err != nil {
- return err
- }
- case CompressionZSTD:
- if m.Value == nil {
- break
- }
- if m.Value, err = zstdDecompress(nil, m.Value); err != nil {
- return err
- }
- if err := m.decodeSet(); err != nil {
- return err
- }
- default:
- return PacketDecodingError{fmt.Sprintf("invalid compression specified (%d)", m.Codec)}
  }
 
  return pd.pop()

diff --git a/record_batch.go b/record_batch.go
@@ -1,14 +1,8 @@
 package sarama
 
 import (
- "bytes"
- "compress/gzip"
  "fmt"
- "io/ioutil"
  "time"
-
- "github.com/eapache/go-xerial-snappy"
- "github.com/pierrec/lz4"
 )
 
 const recordBatchOverhead = 49
@@ -174,31 +168,9 @@ func (b *RecordBatch) decode(pd packetDecoder) (err error) {
  return err
  }
 
- switch b.Codec {
- case CompressionNone:
- case CompressionGZIP:
- reader, err := gzip.NewReader(bytes.NewReader(recBuffer))
- if err != nil {
- return err
- }
- if recBuffer, err = ioutil.ReadAll(reader); err != nil {
- return err
- }
- case CompressionSnappy:
- if recBuffer, err = snappy.Decode(recBuffer); err != nil {
- return err
- }
- case CompressionLZ4:
- reader := lz4.NewReader(bytes.NewReader(recBuffer))
- if recBuffer, err = ioutil.ReadAll(reader); err != nil {
- return err
- }
- case CompressionZSTD:
- if recBuffer, err = zstdDecompress(nil, recBuffer); err != nil {
- return err
- }
- default:
- return PacketDecodingError{fmt.Sprintf("invalid compression specified (%d)", b.Codec)}
+ recBuffer, err = decompress(b.Codec, recBuffer)
+ if err != nil {
+ return err
  }
 
  b.recordsLen = len(recBuffer)
@@ -219,50 +191,8 @@ func (b *RecordBatch) encodeRecords(pe packetEncoder) error {
  }
  b.recordsLen = len(raw)
 
- switch b.Codec {
- case CompressionNone:
- b.compressedRecords = raw
- case CompressionGZIP:
- var buf bytes.Buffer
- var writer *gzip.Writer
- if b.CompressionLevel != CompressionLevelDefault {
- writer, err = gzip.NewWriterLevel(&buf, b.CompressionLevel)
- if err != nil {
- return err
- }
- } else {
- writer = gzip.NewWriter(&buf)
- }
- if _, err := writer.Write(raw); err != nil {
- return err
- }
- if err := writer.Close(); err != nil {
- return err
- }
- b.compressedRecords = buf.Bytes()
- case CompressionSnappy:
- b.compressedRecords = snappy.Encode(raw)
- case CompressionLZ4:
- var buf bytes.Buffer
- writer := lz4.NewWriter(&buf)
- if _, err := writer.Write(raw); err != nil {
- return err
- }
- if err := writer.Close(); err != nil {
- return err
- }
- b.compressedRecords = buf.Bytes()
- case CompressionZSTD:
- c, err := zstdCompressLevel(nil, raw, b.CompressionLevel)
- if err != nil {
- return err
- }
- b.compressedRecords = c
- default:
- return PacketEncodingError{fmt.Sprintf("unsupported compression codec (%d)", b.Codec)}
- }
-
- return nil
+ b.compressedRecords, err = compress(b.Codec, b.CompressionLevel, raw)
+ return err
 }
 
 func (b *RecordBatch) computeAttributes() int16 {