feat(v2): configurable write-path compression (#4006)

* feat(v2): configurable write-path compression

* feat(v2): configurable write-path compression

Author: kolesnikovae

pkg/distributor/write_path/router.go

@@ -101,9 +101,9 @@ func (m *Router) Send(ctx context.Context, req *distributormodel.PushRequest) er
 	config := m.overrides.WritePathOverrides(req.TenantID)
 	switch config.WritePath {
 	case SegmentWriterPath:
-		return m.send(m.segwriterRoute(true))(ctx, req)
+		return m.send(m.segwriterRoute(true, &config))(ctx, req)
 	case CombinedPath:
-		return m.sendToBoth(ctx, req, config)
+		return m.sendToBoth(ctx, req, &config)
 	default:
 		return m.send(m.ingesterRoute())(ctx, req)
 	}
@@ -120,17 +120,17 @@ func (m *Router) ingesterRoute() *route {
 	}
 }
 
-func (m *Router) segwriterRoute(primary bool) *route {
+func (m *Router) segwriterRoute(primary bool, config *Config) *route {
 	return &route{
 		path:    SegmentWriterPath,
 		primary: primary,
 		send: func(ctx context.Context, req *distributormodel.PushRequest) error {
-			return m.sendRequestsToSegmentWriter(ctx, convertRequest(req))
+			return m.sendRequestsToSegmentWriter(ctx, convertRequest(req, config.Compression))
 		},
 	}
 }
 
-func (m *Router) sendToBoth(ctx context.Context, req *distributormodel.PushRequest, config Config) error {
+func (m *Router) sendToBoth(ctx context.Context, req *distributormodel.PushRequest, config *Config) error {
 	r := rand.Float64() // [0.0, 1.0)
 	shouldIngester := config.IngesterWeight > 0.0 && config.IngesterWeight >= r
 	shouldSegwriter := config.SegmentWriterWeight > 0.0 && config.SegmentWriterWeight >= r
@@ -149,7 +149,7 @@ func (m *Router) sendToBoth(ctx context.Context, req *distributormodel.PushReque
 		}
 	}
 	if shouldSegwriter {
-		segwriter = m.segwriterRoute(!shouldIngester)
+		segwriter = m.segwriterRoute(!shouldIngester, config)
 		if segwriter.primary && !config.AsyncIngest {
 			// The request is sent to segment-writer exclusively, and the client
 			// must block until the response returns.
@@ -272,11 +272,11 @@ func (m *Router) sendRequestsToSegmentWriter(ctx context.Context, requests []*se
 	return g.Wait()
 }
 
-func convertRequest(req *distributormodel.PushRequest) []*segmentwriterv1.PushRequest {
+func convertRequest(req *distributormodel.PushRequest, compression Compression) []*segmentwriterv1.PushRequest {
 	r := make([]*segmentwriterv1.PushRequest, 0, len(req.Series)*2)
 	for _, s := range req.Series {
 		for _, p := range s.Samples {
-			r = append(r, convertProfile(p, s.Labels, req.TenantID))
+			r = append(r, convertProfile(p, s.Labels, req.TenantID, compression))
 		}
 	}
 	return r
@@ -286,8 +286,9 @@ func convertProfile(
 	sample *distributormodel.ProfileSample,
 	labels []*typesv1.LabelPair,
 	tenantID string,
+	compression Compression,
 ) *segmentwriterv1.PushRequest {
-	buf, err := pprof.Marshal(sample.Profile.Profile, true)
+	buf, err := pprof.Marshal(sample.Profile.Profile, compression == CompressionGzip)
 	if err != nil {
 		panic(fmt.Sprintf("failed to marshal profile: %v", err))
 	}

pkg/distributor/write_path/write_path.go

@@ -38,7 +38,7 @@ var paths = []WritePath{
 	CombinedPath,
 }
 
-const validOptionsString = "valid options: ingester, segment-writer, combined"
+const validWritePathOptionsString = "valid options: ingester, segment-writer, combined"
 
 func (m *WritePath) Set(text string) error {
 	x := WritePath(text)
@@ -48,26 +48,58 @@ func (m *WritePath) Set(text string) error {
 			return nil
 		}
 	}
-	return fmt.Errorf("%w: %s; %s", ErrInvalidWritePath, x, validOptionsString)
+	return fmt.Errorf("%w: %s; %s", ErrInvalidWritePath, x, validWritePathOptionsString)
 }
 
 func (m *WritePath) String() string { return string(*m) }
 
+type Compression string
+
+const (
+	CompressionNone Compression = "none"
+	CompressionGzip Compression = "gzip"
+)
+
+var ErrInvalidCompression = errors.New("invalid write path compression")
+
+var compressions = []Compression{
+	CompressionNone,
+	CompressionGzip,
+}
+
+const validCompressionOptionsString = "valid compression options: none, gzip"
+
+func (m *Compression) Set(text string) error {
+	x := Compression(text)
+	for _, name := range compressions {
+		if x == name {
+			*m = x
+			return nil
+		}
+	}
+	return fmt.Errorf("%w: %s; %s", ErrInvalidCompression, x, validCompressionOptionsString)
+}
+
+func (m *Compression) String() string { return string(*m) }
+
 type Config struct {
 	WritePath            WritePath     `yaml:"write_path" json:"write_path" doc:"hidden"`
 	IngesterWeight       float64       `yaml:"write_path_ingester_weight" json:"write_path_ingester_weight" doc:"hidden"`
 	SegmentWriterWeight  float64       `yaml:"write_path_segment_writer_weight" json:"write_path_segment_writer_weight" doc:"hidden"`
 	SegmentWriterTimeout time.Duration `yaml:"write_path_segment_writer_timeout" json:"write_path_segment_writer_timeout" doc:"hidden"`
+	Compression          Compression   `yaml:"write_path_compression" json:"write_path_compression" doc:"hidden"`
 	AsyncIngest          bool          `yaml:"async_ingest" json:"async_ingest" doc:"hidden"`
 }
 
 func (o *Config) RegisterFlags(f *flag.FlagSet) {
 	o.WritePath = IngesterPath
-	f.Var(&o.WritePath, "write-path", "Controls the write path route; "+validOptionsString+".")
+	o.Compression = CompressionNone
+	f.Var(&o.WritePath, "write-path", "Controls the write path route; "+validWritePathOptionsString+".")
 	f.Float64Var(&o.IngesterWeight, "write-path.ingester-weight", 1,
 		"Specifies the fraction [0:1] that should be send to ingester in combined mode. 0 means no traffics is sent to ingester. 1 means 100% of requests are sent to ingester.")
 	f.Float64Var(&o.SegmentWriterWeight, "write-path.segment-writer-weight", 0,
 		"Specifies the fraction [0:1] that should be send to segment-writer in combined mode. 0 means no traffics is sent to segment-writer. 1 means 100% of requests are sent to segment-writer.")
 	f.DurationVar(&o.SegmentWriterTimeout, "write-path.segment-writer-timeout", 5*time.Second, "Timeout for segment writer requests.")
+	f.Var(&o.Compression, "write-path.compression", "Compression algorithm to use for segment writer requests; "+validCompressionOptionsString+".")
 	f.BoolVar(&o.AsyncIngest, "async-ingest", false, "If true, the write path will not wait for the segment-writer to finish processing the request. Writes to ingester always synchronous.")
 }