Use snappy from klauspost/compress

Signed-off-by: Filip Petkovski <filip.petkovsky@gmail.com>

cmd/thanos/receive.go

@@ -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"
@@ -868,8 +868,8 @@ 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)
+	compressionOptions := strings.Join([]string{snappy.Name, compressionNone}, ", ")
+	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, snappy.Name, compressionNone)
 
 	cmd.Flag("receive.replication-factor", "How many times to replicate incoming write requests.").Default("1").Uint64Var(&rc.replicationFactor)
 

docs/components/receive.md

@@ -192,10 +192,10 @@ 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
+                                 to other receivers. Must be one of: snappy,
+                                 none
       --receive.hashrings=<content>
                                  Alternative to 'receive.hashrings-file' flag
                                  (lower priority). Content of file that contains

pkg/extgrpc/snappy/snappy.go

@@ -0,0 +1,87 @@
+package snappy
+
+import (
+	"io"
+	"sync"
+
+	"github.com/klauspost/compress/snappy"
+	"google.golang.org/grpc/encoding"
+)
+
+// 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
+}

pkg/extgrpc/snappy/snappy_test.go

@@ -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)
+	}
+}