From 394c12b0acb741a0084f7be6db8525925f9eed43 Mon Sep 17 00:00:00 2001 From: Filip Petkovski Date: Mon, 8 Aug 2022 09:07:03 +0200 Subject: [PATCH] Use snappy from klauspost/compress --- cmd/thanos/receive.go | 4 +- docs/components/receive.md | 2 +- pkg/extgrpc/snappy/snappy.go | 86 +++++++++++++++++++++++++++++++ pkg/extgrpc/snappy/snappy_test.go | 73 ++++++++++++++++++++++++++ 4 files changed, 162 insertions(+), 3 deletions(-) create mode 100644 pkg/extgrpc/snappy/snappy.go create mode 100644 pkg/extgrpc/snappy/snappy_test.go diff --git a/cmd/thanos/receive.go b/cmd/thanos/receive.go index 000c167c1a6..6f4f27afd17 100644 --- a/cmd/thanos/receive.go +++ b/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" @@ -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) diff --git a/docs/components/receive.md b/docs/components/receive.md index ee5a6005d1f..7d7e24d7501 100644 --- a/docs/components/receive.md +++ b/docs/components/receive.md @@ -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 diff --git a/pkg/extgrpc/snappy/snappy.go b/pkg/extgrpc/snappy/snappy.go new file mode 100644 index 00000000000..763176e21ba --- /dev/null +++ b/pkg/extgrpc/snappy/snappy.go @@ -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 +} diff --git a/pkg/extgrpc/snappy/snappy_test.go b/pkg/extgrpc/snappy/snappy_test.go new file mode 100644 index 00000000000..a51d5077fb6 --- /dev/null +++ b/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) + } +}