Add protobuf codec for query range and instant query responses

pkg/api/api.go

@@ -73,13 +73,16 @@ type Config struct {
 	corsRegexString string `yaml:"cors_origin"`
 
 	buildInfoEnabled bool `yaml:"build_info_enabled"`
+
+	QuerierDefaultCodec string `yaml:"querier_default_codec"`
 }
 
 // RegisterFlags adds the flags required to config this to the given FlagSet.
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&cfg.ResponseCompression, "api.response-compression-enabled", false, "Use GZIP compression for API responses. Some endpoints serve large YAML or JSON blobs which can benefit from compression.")
 	f.Var(&cfg.HTTPRequestHeadersToLog, "api.http-request-headers-to-log", "Which HTTP Request headers to add to logs")
 	f.BoolVar(&cfg.buildInfoEnabled, "api.build-info-enabled", false, "If enabled, build Info API will be served by query frontend or querier.")
+	f.StringVar(&cfg.QuerierDefaultCodec, "api.querier-default-codec", "json", "Choose default codec for querier response serialization. Supports 'json' and 'protobuf'.")
 	cfg.RegisterFlagsWithPrefix("", f)
 }
 

pkg/api/handlers.go

@@ -26,6 +26,7 @@ import (
 	"github.com/weaveworks/common/middleware"
 
 	"github.com/cortexproject/cortex/pkg/querier"
+	"github.com/cortexproject/cortex/pkg/querier/codec"
 	"github.com/cortexproject/cortex/pkg/querier/stats"
 	"github.com/cortexproject/cortex/pkg/util"
 )
@@ -231,6 +232,9 @@ func NewQuerierHandler(
 		false,
 	)
 
+	// JSON codec is already installed. Install Protobuf codec to give the option for using either.
+	api.InstallCodec(codec.ProtobufCodec{})
+
 	router := mux.NewRouter()
 
 	// Use a separate metric for the querier in order to differentiate requests from the query-frontend when
@@ -294,7 +298,7 @@ type buildInfoHandler struct {
 }
 
 type buildInfoResponse struct {
-	Status string                `json:"status"`
+	Status string                   `json:"status"`
 	Data   *v1.PrometheusVersion `json:"data"`
 }
 

pkg/cortex/modules.go

@@ -442,9 +442,9 @@ func (t *Cortex) initFlusher() (serv services.Service, err error) {
 func (t *Cortex) initQueryFrontendTripperware() (serv services.Service, err error) {
 	queryAnalyzer := querysharding.NewQueryAnalyzer()
 	// PrometheusCodec is a codec to encode and decode Prometheus query range requests and responses.
-	prometheusCodec := queryrange.NewPrometheusCodec(false)
+	prometheusCodec := queryrange.NewPrometheusCodec(false, t.Cfg.Querier.PrometheusCodecCompression, t.Cfg.API.QuerierDefaultCodec)
 	// ShardedPrometheusCodec is same as PrometheusCodec but to be used on the sharded queries (it sum up the stats)
-	shardedPrometheusCodec := queryrange.NewPrometheusCodec(true)
+	shardedPrometheusCodec := queryrange.NewPrometheusCodec(true, t.Cfg.Querier.PrometheusCodecCompression, t.Cfg.API.QuerierDefaultCodec)
 
 	queryRangeMiddlewares, cache, err := queryrange.Middlewares(
 		t.Cfg.QueryRange,
@@ -472,7 +472,7 @@ func (t *Cortex) initQueryFrontendTripperware() (serv services.Service, err erro
 		queryRangeMiddlewares,
 		instantQueryMiddlewares,
 		prometheusCodec,
-		instantquery.InstantQueryCodec,
+		instantquery.NewInstantQueryCodec(t.Cfg.Querier.PrometheusCodecCompression, t.Cfg.API.QuerierDefaultCodec),
 		t.Overrides,
 		queryAnalyzer,
 		t.Cfg.Querier.DefaultEvaluationInterval,

pkg/querier/codec/protobuf_codec.go

@@ -0,0 +1,159 @@
+package codec
+
+import (
+	"github.com/cortexproject/cortex/pkg/querier/tripperware"
+	"github.com/prometheus/prometheus/promql"
+	"github.com/cortexproject/cortex/pkg/cortexpb"
+	"github.com/prometheus/prometheus/util/stats"
+	jsoniter "github.com/json-iterator/go"
+	"github.com/prometheus/common/model"
+	"github.com/prometheus/prometheus/web/api/v1"
+	"github.com/gogo/protobuf/proto"
+)
+type ProtobufCodec struct{}
+
+func (p ProtobufCodec) ContentType() v1.MIMEType {
+	return v1.MIMEType{Type: "application", SubType: "x-protobuf"}
+}
+
+func (p ProtobufCodec) CanEncode(resp *v1.Response) bool {
+	// Errors are parsed by default json codec
+	if resp.Error != "" || resp.Data == nil {
+		return false
+	}
+	return true
+}
+
+func (p ProtobufCodec) Encode(resp *v1.Response) ([]byte, error) {
+	prometheusQueryResponse, err := createPrometheusQueryResponse(resp)
+	if err != nil {
+		return []byte{}, err
+	}
+	b, err := proto.Marshal(prometheusQueryResponse)
+	return b, err
+}
+
+func createPrometheusQueryResponse(resp *v1.Response) (*tripperware.PrometheusResponse, error) {
+	var data = resp.Data.(*v1.QueryData)
+
+	var queryResult tripperware.PrometheusQueryResult
+	switch string(data.ResultType) {
+	case model.ValMatrix.String():
+		queryResult.Result = &tripperware.PrometheusQueryResult_Matrix{
+			Matrix: &tripperware.Matrix{
+				SampleStreams: *getMatrixSampleStreams(data),
+			},
+		}
+	case model.ValVector.String():
+		queryResult.Result = &tripperware.PrometheusQueryResult_Vector{
+			Vector: &tripperware.Vector{
+				Samples: *getVectorSamples(data),
+			},
+		}
+	default:
+		json := jsoniter.ConfigCompatibleWithStandardLibrary
+		rawBytes, err := json.Marshal(data)
+		if err != nil {
+			return nil, err
+		}
+		queryResult.Result = &tripperware.PrometheusQueryResult_RawBytes{RawBytes: rawBytes}
+	}
+
+	var stats *tripperware.PrometheusResponseStats
+	if data.Stats != nil {
+		builtin := data.Stats.Builtin()
+		stats = &tripperware.PrometheusResponseStats{Samples: getStats(&builtin)}
+	}
+
+	return &tripperware.PrometheusResponse{
+		Status: string(resp.Status),
+		Data: tripperware.PrometheusData{
+			ResultType: string(data.ResultType),
+			Result:     queryResult,
+			Stats:      stats,
+		},
+		ErrorType: string(resp.ErrorType),
+		Error:     resp.Error,
+		Warnings: resp.Warnings,
+	}, nil
+}
+
+func getMatrixSampleStreams(data *v1.QueryData) *[]tripperware.SampleStream {
+	sampleStreamsLen := len(data.Result.(promql.Matrix))
+	sampleStreams := make([]tripperware.SampleStream, sampleStreamsLen)
+
+	for i := 0; i < sampleStreamsLen; i++ {
+		labelsLen := len(data.Result.(promql.Matrix)[i].Metric)
+		var labels []cortexpb.LabelAdapter
+		if labelsLen > 0 {
+			labels = make([]cortexpb.LabelAdapter, labelsLen)
+			for j := 0; j < labelsLen; j++ {
+				labels[j] = cortexpb.LabelAdapter{
+					Name:  data.Result.(promql.Matrix)[i].Metric[j].Name,
+					Value: data.Result.(promql.Matrix)[i].Metric[j].Value,
+				}
+			}
+		}
+
+		samplesLen := len(data.Result.(promql.Matrix)[i].Floats)
+		var samples []cortexpb.Sample
+		if samplesLen > 0 {
+			samples = make([]cortexpb.Sample, samplesLen)
+			for j := 0; j < samplesLen; j++ {
+				samples[j] = cortexpb.Sample{
+					Value:       data.Result.(promql.Matrix)[i].Floats[j].F,
+					TimestampMs: data.Result.(promql.Matrix)[i].Floats[j].T,
+				}
+			}
+		}
+		sampleStreams[i] = tripperware.SampleStream{Labels: labels, Samples: samples}
+	}
+	return &sampleStreams
+}
+
+func getVectorSamples(data *v1.QueryData) *[]tripperware.Sample {
+	vectorSamplesLen := len(data.Result.(promql.Vector))
+	vectorSamples := make([]tripperware.Sample, vectorSamplesLen)
+
+	for i := 0; i < vectorSamplesLen; i++ {
+		labelsLen := len(data.Result.(promql.Vector)[i].Metric)
+		var labels []cortexpb.LabelAdapter
+		if labelsLen > 0 {
+			labels = make([]cortexpb.LabelAdapter, labelsLen)
+			for j := 0; j < labelsLen; j++ {
+				labels[j] = cortexpb.LabelAdapter{
+					Name:  data.Result.(promql.Vector)[i].Metric[j].Name,
+					Value: data.Result.(promql.Vector)[i].Metric[j].Value,
+				}
+			}
+		}
+
+		vectorSamples[i] = tripperware.Sample{
+			Labels: labels,
+			Sample: &cortexpb.Sample{
+				TimestampMs: data.Result.(promql.Vector)[i].T,
+				Value:       data.Result.(promql.Vector)[i].F,
+			},
+		}
+	}
+	return &vectorSamples
+}
+
+func getStats(builtin *stats.BuiltinStats) *tripperware.PrometheusResponseSamplesStats {
+	queryableSamplesStatsPerStepLen := len(builtin.Samples.TotalQueryableSamplesPerStep)
+	queryableSamplesStatsPerStep := make([]*tripperware.PrometheusResponseQueryableSamplesStatsPerStep, queryableSamplesStatsPerStepLen)
+	for i := 0; i < queryableSamplesStatsPerStepLen; i++ {
+		queryableSamplesStatsPerStep[i] = &tripperware.PrometheusResponseQueryableSamplesStatsPerStep{
+			Value:       builtin.Samples.TotalQueryableSamplesPerStep[i].V,
+			TimestampMs: builtin.Samples.TotalQueryableSamplesPerStep[i].T,
+		}
+	}
+
+	statSamples := tripperware.PrometheusResponseSamplesStats{
+		TotalQueryableSamples:        builtin.Samples.TotalQueryableSamples,
+		TotalQueryableSamplesPerStep: queryableSamplesStatsPerStep,
+		PeakSamples: 			      int64(builtin.Samples.PeakSamples),
+	}
+
+	return &statSamples
+}

pkg/querier/querier.go

@@ -49,6 +49,9 @@ type Config struct {
 	QueryIngestersWithin      time.Duration `yaml:"query_ingesters_within"`
 	EnablePerStepStats        bool          `yaml:"per_step_stats_enabled"`
 
+	// Use compression when returning promql response. Supported values are 'gzip' or '' (disable compression)
+	PrometheusCodecCompression string `yaml:"prometheus_codec_compression"`
+
 	// QueryStoreAfter the time after which queries should also be sent to the store and not just ingesters.
 	QueryStoreAfter    time.Duration `yaml:"query_store_after"`
 	MaxQueryIntoFuture time.Duration `yaml:"max_query_into_future"`
@@ -109,6 +112,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.IntVar(&cfg.MaxSamples, "querier.max-samples", 50e6, "Maximum number of samples a single query can load into memory.")
 	f.DurationVar(&cfg.QueryIngestersWithin, "querier.query-ingesters-within", 0, "Maximum lookback beyond which queries are not sent to ingester. 0 means all queries are sent to ingester.")
 	f.BoolVar(&cfg.EnablePerStepStats, "querier.per-step-stats-enabled", false, "Enable returning samples stats per steps in query response.")
+	f.StringVar(&cfg.PrometheusCodecCompression, "querier.prometheus_codec_compression", "", "Use compression when returning promql response. Supported values 'gzip', 'snappy', and '' (disable compression)")
 	f.DurationVar(&cfg.MaxQueryIntoFuture, "querier.max-query-into-future", 10*time.Minute, "Maximum duration into the future you can query. 0 to disable.")
 	f.DurationVar(&cfg.DefaultEvaluationInterval, "querier.default-evaluation-interval", time.Minute, "The default evaluation interval or step size for subqueries.")
 	f.DurationVar(&cfg.QueryStoreAfter, "querier.query-store-after", 0, "The time after which a metric should be queried from storage and not just ingesters. 0 means all queries are sent to store. When running the blocks storage, if this option is enabled, the time range of the query sent to the store will be manipulated to ensure the query end is not more recent than 'now - query-store-after'.")

pkg/querier/tripperware/instantquery/instant_query.go

@@ -22,7 +22,7 @@ import (
 )
 
 var (
-	InstantQueryCodec tripperware.Codec = newInstantQueryCodec()
+	InstantQueryCodec tripperware.Codec = NewInstantQueryCodec("", "protobuf")
 
 	json = jsoniter.Config{
 		EscapeHTML:             false, // No HTML in our responses.
@@ -33,11 +33,17 @@ var (
 
 type instantQueryCodec struct {
 	tripperware.Codec
-	now func() time.Time
+	compression    string
+	defaultCodec   string
+	now            func() time.Time
 }
 
-func newInstantQueryCodec() instantQueryCodec {
-	return instantQueryCodec{now: time.Now}
+func NewInstantQueryCodec(compression string, defaultCodec string) instantQueryCodec {
+	return instantQueryCodec{
+		compression:    compression,
+		defaultCodec:   defaultCodec,
+		now:            time.Now,
+	}
 }
 
 func (c instantQueryCodec) DecodeRequest(_ context.Context, r *http.Request, forwardHeaders []string) (tripperware.Request, error) {
@@ -83,17 +89,20 @@ func (instantQueryCodec) DecodeResponse(ctx context.Context, r *http.Response, _
 	}
 
 	var resp tripperware.PrometheusResponse
-	if err := json.Unmarshal(buf, &resp); err != nil {
+	err = tripperware.UnmarshalResponse(r, buf, &resp)
+
+	if err != nil {
 		return nil, httpgrpc.Errorf(http.StatusInternalServerError, "error decoding response: %v", err)
 	}
 
 	for h, hv := range r.Header {
 		resp.Headers = append(resp.Headers, &tripperware.PrometheusResponseHeader{Name: h, Values: hv})
 	}
+
 	return &resp, nil
 }
 
-func (instantQueryCodec) EncodeRequest(ctx context.Context, r tripperware.Request) (*http.Request, error) {
+func (c instantQueryCodec) EncodeRequest(ctx context.Context, r tripperware.Request) (*http.Request, error) {
 	promReq, ok := r.(*tripperware.PrometheusRequest)
 	if !ok {
 		return nil, httpgrpc.Errorf(http.StatusBadRequest, "invalid request format")
@@ -120,8 +129,7 @@ func (instantQueryCodec) EncodeRequest(ctx context.Context, r tripperware.Reques
 		}
 	}
 
-	// Always ask gzip to the querier
-	h.Set("Accept-Encoding", "gzip")
+	tripperware.SetRequestHeaders(h, c.defaultCodec, c.compression)
 
 	req := &http.Request{
 		Method:     "GET",
@@ -152,7 +160,7 @@ func (instantQueryCodec) EncodeResponse(ctx context.Context, res tripperware.Res
 
 	resp := http.Response{
 		Header: http.Header{
-			"Content-Type": []string{"application/json"},
+			"Content-Type": []string{tripperware.ApplicationJson},
 		},
 		Body:          io.NopCloser(bytes.NewBuffer(b)),
 		StatusCode:    http.StatusOK,