Skip to content

Commit

Permalink
Use snappy from klauspost/compress
Browse files Browse the repository at this point in the history
  • Loading branch information
fpetkovski committed Aug 8, 2022
1 parent 3c954c8 commit 394c12b
Show file tree
Hide file tree
Showing 4 changed files with 162 additions and 3 deletions.
4 changes: 2 additions & 2 deletions cmd/thanos/receive.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,11 @@ import (
"github.com/thanos-io/objstore/client"
"gopkg.in/yaml.v2"

"github.com/thanos-io/thanos/internal/cortex/util/grpcencoding/snappy"
"github.com/thanos-io/thanos/pkg/block/metadata"
"github.com/thanos-io/thanos/pkg/component"
"github.com/thanos-io/thanos/pkg/exemplars"
"github.com/thanos-io/thanos/pkg/extgrpc"
"github.com/thanos-io/thanos/pkg/extgrpc/snappy"
"github.com/thanos-io/thanos/pkg/extkingpin"
"github.com/thanos-io/thanos/pkg/extprom"
"github.com/thanos-io/thanos/pkg/info"
Expand Down Expand Up @@ -869,7 +869,7 @@ func (rc *receiveConfig) registerFlag(cmd extkingpin.FlagClause) {
cmd.Flag("receive.replica-header", "HTTP header specifying the replica number of a write request.").Default(receive.DefaultReplicaHeader).StringVar(&rc.replicaHeader)

compressionOptions := strings.Join([]string{compressionNone, snappy.Name}, ", ")
cmd.Flag("receive.grpc-compression", "Compression algorithm to use for gRPC requests to other receivers. Must be one of: "+compressionOptions).Default(compressionNone).EnumVar(&rc.compression, compressionNone, snappy.Name)
cmd.Flag("receive.grpc-compression", "Compression algorithm to use for gRPC requests to other receivers. Must be one of: "+compressionOptions).Default(snappy.Name).EnumVar(&rc.compression, compressionNone, snappy.Name)

cmd.Flag("receive.replication-factor", "How many times to replicate incoming write requests.").Default("1").Uint64Var(&rc.replicationFactor)

Expand Down
2 changes: 1 addition & 1 deletion docs/components/receive.md
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,7 @@ Flags:
--receive.default-tenant-id="default-tenant"
Default tenant ID to use when none is provided
via a header.
--receive.grpc-compression=none
--receive.grpc-compression=snappy
Compression algorithm to use for gRPC requests
to other receivers. Must be one of: none,
snappy
Expand Down
86 changes: 86 additions & 0 deletions pkg/extgrpc/snappy/snappy.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
package snappy

import (
"github.com/klauspost/compress/snappy"
"google.golang.org/grpc/encoding"
"io"
"sync"
)

// Name is the name registered for the snappy compressor.
const Name = "snappy"

func init() {
encoding.RegisterCompressor(newCompressor())
}

type compressor struct {
writersPool sync.Pool
readersPool sync.Pool
}

func newCompressor() *compressor {
c := &compressor{}
c.readersPool = sync.Pool{
New: func() interface{} {
return snappy.NewReader(nil)
},
}
c.writersPool = sync.Pool{
New: func() interface{} {
return snappy.NewBufferedWriter(nil)
},
}
return c
}

func (c *compressor) Name() string {
return Name
}

func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
wr := c.writersPool.Get().(*snappy.Writer)
wr.Reset(w)
return writeCloser{wr, &c.writersPool}, nil
}

func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
dr := c.readersPool.Get().(*snappy.Reader)
dr.Reset(r)
return reader{dr, &c.readersPool}, nil
}

type writeCloser struct {
writer *snappy.Writer
pool *sync.Pool
}

func (w writeCloser) Write(p []byte) (n int, err error) {
return w.writer.Write(p)
}

func (w writeCloser) Close() error {
defer func() {
w.writer.Reset(nil)
w.pool.Put(w.writer)
}()

if w.writer != nil {
return w.writer.Close()
}
return nil
}

type reader struct {
reader *snappy.Reader
pool *sync.Pool
}

func (r reader) Read(p []byte) (n int, err error) {
n, err = r.reader.Read(p)
if err == io.EOF {
r.reader.Reset(nil)
r.pool.Put(r.reader)
}
return n, err
}
73 changes: 73 additions & 0 deletions pkg/extgrpc/snappy/snappy_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
// Copyright (c) The Cortex Authors.
// Licensed under the Apache License 2.0.

package snappy

import (
"bytes"
"io"
"strings"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func TestSnappy(t *testing.T) {
c := newCompressor()
assert.Equal(t, "snappy", c.Name())

tests := []struct {
test string
input string
}{
{"empty", ""},
{"short", "hello world"},
{"long", strings.Repeat("123456789", 1024)},
}
for _, test := range tests {
t.Run(test.test, func(t *testing.T) {
var buf bytes.Buffer
// Compress
w, err := c.Compress(&buf)
require.NoError(t, err)
n, err := w.Write([]byte(test.input))
require.NoError(t, err)
assert.Len(t, test.input, n)
err = w.Close()
require.NoError(t, err)
// Decompress
r, err := c.Decompress(&buf)
require.NoError(t, err)
out, err := io.ReadAll(r)
require.NoError(t, err)
assert.Equal(t, test.input, string(out))
})
}
}

func BenchmarkSnappyCompress(b *testing.B) {
data := []byte(strings.Repeat("123456789", 1024))
c := newCompressor()
b.ResetTimer()
for i := 0; i < b.N; i++ {
w, _ := c.Compress(io.Discard)
_, _ = w.Write(data)
_ = w.Close()
}
}

func BenchmarkSnappyDecompress(b *testing.B) {
data := []byte(strings.Repeat("123456789", 1024))
c := newCompressor()
var buf bytes.Buffer
w, _ := c.Compress(&buf)
_, _ = w.Write(data)
reader := bytes.NewReader(buf.Bytes())
b.ResetTimer()
for i := 0; i < b.N; i++ {
r, _ := c.Decompress(reader)
_, _ = io.ReadAll(r)
_, _ = reader.Seek(0, io.SeekStart)
}
}

0 comments on commit 394c12b

Please sign in to comment.