Add support for snappy compression in receivers

CHANGELOG.md

@@ -26,9 +26,10 @@ We use *breaking :warning:* to mark changes that are not backward compatible (re
 - [#5470](https://github.com/thanos-io/thanos/pull/5470) Receive: Implement exposing TSDB stats for all tenants
 - [#5493](https://github.com/thanos-io/thanos/pull/5493) Compact: Added `--compact.blocks-fetch-concurrency` allowing to configure number of go routines for download blocks during compactions.
 - [#5527](https://github.com/thanos-io/thanos/pull/5527) Receive: Add per request limits for remote write.
-- [#5520](https://github.com/thanos-io/thanos/pull/5520) Receive: Meta-monitoring based active series limiting
+- [#5520](https://github.com/thanos-io/thanos/pull/5520) Receive: Meta-monitoring based active series limiting.
 - [#5555](https://github.com/thanos-io/thanos/pull/5555) Query: Added `--query.active-query-path` flag, allowing the user to configure the directory to create an active query tracking file, `queries.active`, for different resolution.
 - [#5566](https://github.com/thanos-io/thanos/pull/5566) Receive: Added experimental support to enable chunk write queue via `--tsdb.write-queue-size` flag.
+- [#5575](https://github.com/thanos-io/thanos/pull/5575) Receive: Add support for gRPC compression with snappy.
 
 ### Changed
 

cmd/thanos/receive.go

@@ -11,6 +11,8 @@ import (
 	"strings"
 	"time"
 
+	"google.golang.org/grpc"
+
 	extflag "github.com/efficientgo/tools/extkingpin"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -33,6 +35,7 @@ import (
 	"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"
@@ -48,6 +51,8 @@ import (
 	"github.com/thanos-io/thanos/pkg/tls"
 )
 
+const compressionNone = "none"
+
 func registerReceive(app *extkingpin.App) {
 	cmd := app.Command(component.Receive.String(), "Accept Prometheus remote write API requests and write to local tsdb.")
 
@@ -140,6 +145,9 @@ func runReceive(
 	if err != nil {
 		return err
 	}
+	if conf.compression != compressionNone {
+		dialOpts = append(dialOpts, grpc.WithDefaultCallOptions(grpc.UseCompressor(conf.compression)))
+	}
 
 	var bkt objstore.Bucket
 	confContentYaml, err := conf.objStoreConfig.Content()
@@ -783,6 +791,7 @@ type receiveConfig struct {
 	replicaHeader     string
 	replicationFactor uint64
 	forwardTimeout    *model.Duration
+	compression       string
 
 	tsdbMinBlockDuration       *model.Duration
 	tsdbMaxBlockDuration       *model.Duration
@@ -861,6 +870,9 @@ 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{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)
 
 	cmd.Flag("receive.tenant-limits.max-head-series", "The total number of active (head) series that a tenant is allowed to have within a Receive topology. For more details refer: https://thanos.io/tip/components/receive.md/#limiting").Hidden().Uint64Var(&rc.maxPerTenantLimit)

docs/components/receive.md

@@ -192,6 +192,10 @@ Flags:
       --receive.default-tenant-id="default-tenant"
                                  Default tenant ID to use when none is provided
                                  via a header.
+      --receive.grpc-compression=snappy
+                                 Compression algorithm to use for gRPC requests
+                                 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,90 @@
+// Copyright (c) The Thanos Authors.
+// Licensed under the Apache License 2.0.
+
+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 Thanos 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)
+	}
+}