Support Zstd compression for gRPC

docs/configuration/config-file-reference.md

@@ -198,7 +198,7 @@ query_scheduler:
     [max_send_msg_size: <int> | default = 16777216]
 
     # Use compression when sending messages. Supported values are: 'gzip',
-    # 'snappy' and '' (disable compression)
+    # 'snappy', 'zstd' and '' (disable compression)
     # CLI flag: -query-scheduler.grpc-client-config.grpc-compression
     [grpc_compression: <string> | default = ""]
 
@@ -963,7 +963,7 @@ grpc_client_config:
   [max_send_msg_size: <int> | default = 16777216]
 
   # Use compression when sending messages. Supported values are: 'gzip',
-  # 'snappy' and '' (disable compression)
+  # 'snappy', 'zstd' and '' (disable compression)
   # CLI flag: -frontend.grpc-client-config.grpc-compression
   [grpc_compression: <string> | default = ""]
 
@@ -1126,7 +1126,7 @@ ruler_client:
   [max_send_msg_size: <int> | default = 16777216]
 
   # Use compression when sending messages. Supported values are: 'gzip',
-  # 'snappy' and '' (disable compression)
+  # 'snappy', 'zstd' and '' (disable compression)
   # CLI flag: -ruler.client.grpc-compression
   [grpc_compression: <string> | default = ""]
 
@@ -2118,7 +2118,7 @@ grpc_client_config:
   [max_send_msg_size: <int> | default = 16777216]
 
   # Use compression when sending messages. Supported values are: 'gzip',
-  # 'snappy' and '' (disable compression)
+  # 'snappy', 'zstd' and '' (disable compression)
   # CLI flag: -ingester.client.grpc-compression
   [grpc_compression: <string> | default = ""]
 
@@ -2225,7 +2225,7 @@ grpc_client_config:
   [max_send_msg_size: <int> | default = 16777216]
 
   # Use compression when sending messages. Supported values are: 'gzip',
-  # 'snappy' and '' (disable compression)
+  # 'snappy', 'zstd' and '' (disable compression)
   # CLI flag: -querier.frontend-client.grpc-compression
   [grpc_compression: <string> | default = ""]
 

go.mod

@@ -31,6 +31,7 @@ require (
 	github.com/hashicorp/go-sockaddr v1.0.2
 	github.com/hashicorp/memberlist v0.5.0
 	github.com/json-iterator/go v1.1.12
+	github.com/klauspost/compress v1.15.14
 	github.com/lib/pq v1.10.7
 	github.com/minio/minio-go/v7 v7.0.46
 	github.com/mitchellh/go-wordwrap v1.0.1
@@ -150,7 +151,6 @@ require (
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/jpillora/backoff v1.0.0 // indirect
 	github.com/julienschmidt/httprouter v1.3.0 // indirect
-	github.com/klauspost/compress v1.15.9 // indirect
 	github.com/klauspost/cpuid/v2 v2.1.0 // indirect
 	github.com/kylelemons/godebug v1.1.0 // indirect
 	github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect

go.sum

@@ -1022,8 +1022,8 @@ github.com/klauspost/compress v1.11.13/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdY
 github.com/klauspost/compress v1.13.1/go.mod h1:8dP1Hq4DHOhN9w426knH3Rhby4rFm6D8eO+e+Dq5Gzg=
 github.com/klauspost/compress v1.13.4/go.mod h1:8dP1Hq4DHOhN9w426knH3Rhby4rFm6D8eO+e+Dq5Gzg=
 github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
-github.com/klauspost/compress v1.15.9 h1:wKRjX6JRtDdrE9qwa4b/Cip7ACOshUI4smpCQanqjSY=
-github.com/klauspost/compress v1.15.9/go.mod h1:PhcZ0MbTNciWF3rruxRgKxI5NkcHHrHUDtV4Yw2GlzU=
+github.com/klauspost/compress v1.15.14 h1:i7WCKDToww0wA+9qrUZ1xOjp218vfFo3nTU6UHp+gOc=
+github.com/klauspost/compress v1.15.14/go.mod h1:QPwzmACJjUTFsnSHH934V6woptycfrDDJnH7hvFVbGM=
 github.com/klauspost/cpuid/v2 v2.0.1/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.0.4/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.1.0 h1:eyi1Ad2aNJMW95zcSbmGg7Cg6cq3ADwLpMAP96d8rF0=

pkg/util/grpcclient/grpcclient.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/util/backoff"
 	"github.com/cortexproject/cortex/pkg/util/grpcencoding/snappy"
+	"github.com/cortexproject/cortex/pkg/util/grpcencoding/zstd"
 	"github.com/cortexproject/cortex/pkg/util/tls"
 )
 
@@ -40,7 +41,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 func (cfg *Config) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
 	f.IntVar(&cfg.MaxRecvMsgSize, prefix+".grpc-max-recv-msg-size", 100<<20, "gRPC client max receive message size (bytes).")
 	f.IntVar(&cfg.MaxSendMsgSize, prefix+".grpc-max-send-msg-size", 16<<20, "gRPC client max send message size (bytes).")
-	f.StringVar(&cfg.GRPCCompression, prefix+".grpc-compression", "", "Use compression when sending messages. Supported values are: 'gzip', 'snappy' and '' (disable compression)")
+	f.StringVar(&cfg.GRPCCompression, prefix+".grpc-compression", "", "Use compression when sending messages. Supported values are: 'gzip', 'snappy', 'zstd' and '' (disable compression)")
 	f.Float64Var(&cfg.RateLimit, prefix+".grpc-client-rate-limit", 0., "Rate limit for gRPC client; 0 means disabled.")
 	f.IntVar(&cfg.RateLimitBurst, prefix+".grpc-client-rate-limit-burst", 0, "Rate limit burst for gRPC client.")
 	f.BoolVar(&cfg.BackoffOnRatelimits, prefix+".backoff-on-ratelimits", false, "Enable backoff and retry when we hit ratelimits.")
@@ -53,7 +54,7 @@ func (cfg *Config) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
 
 func (cfg *Config) Validate(log log.Logger) error {
 	switch cfg.GRPCCompression {
-	case gzip.Name, snappy.Name, "":
+	case gzip.Name, snappy.Name, zstd.Name, "":
 		// valid
 	default:
 		return errors.Errorf("unsupported compression type: %s", cfg.GRPCCompression)

pkg/util/grpcencoding/zstd/zstd.go

@@ -0,0 +1,82 @@
+package zstd
+
+import (
+	"bytes"
+	"io"
+
+	"github.com/klauspost/compress/zstd"
+	"google.golang.org/grpc/encoding"
+)
+
+const Name = "zstd"
+
+type compressor struct {
+	encoder *zstd.Encoder
+	decoder *zstd.Decoder
+}
+
+func init() {
+	enc, _ := zstd.NewWriter(nil)
+	dec, _ := zstd.NewReader(nil)
+	c := &compressor{
+		encoder: enc,
+		decoder: dec,
+	}
+	encoding.RegisterCompressor(c)
+}
+
+// SetLevel updates the registered compressor to use a particular compression
+// level. NOTE: this function must only be called from an init function, and
+// is not threadsafe.
+func SetLevel(level zstd.EncoderLevel) error {
+	c := encoding.GetCompressor(Name).(*compressor)
+
+	enc, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(level))
+	if err != nil {
+		return err
+	}
+
+	c.encoder = enc
+	return nil
+}
+
+func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
+	return &zstdWriteCloser{
+		enc:    c.encoder,
+		writer: w,
+	}, nil
+}
+
+type zstdWriteCloser struct {
+	enc    *zstd.Encoder
+	writer io.Writer    // Compressed data will be written here.
+	buf    bytes.Buffer // Buffer uncompressed data here, compress on Close.
+}
+
+func (z *zstdWriteCloser) Write(p []byte) (int, error) {
+	return z.buf.Write(p)
+}
+
+func (z *zstdWriteCloser) Close() error {
+	compressed := z.enc.EncodeAll(z.buf.Bytes(), nil)
+	_, err := io.Copy(z.writer, bytes.NewReader(compressed))
+	return err
+}
+
+func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
+	compressed, err := io.ReadAll(r)
+	if err != nil {
+		return nil, err
+	}
+
+	uncompressed, err := c.decoder.DecodeAll(compressed, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return bytes.NewReader(uncompressed), nil
+}
+
+func (c *compressor) Name() string {
+	return Name
+}