Support ruler to retrieve proto format query response #6345

SungJin1212 · 2024-11-18T06:31:39Z

In #6151, we can retrieve query results from Query frontends when the rules are evaluated. The limitation was, since we only support json format, we cannot retrieve query results including the native histogram.

To address this concern, this PR adds -ruler.query-response-format flag, which specifies which format to use when retrieving query results from the Query frontend. The supported values are protobuf and json, so we can retrieve proto format query results when we set the value to protobuf.

Which issue(s) this PR fixes:
Fixes #

Checklist

Tests updated
Documentation added
CHANGELOG.md updated - the order of entries should be [CHANGE], [FEATURE], [ENHANCEMENT], [BUGFIX]

pkg/querier/tripperware/query.proto

pkg/ruler/frontend_decoder.go

pkg/querier/tripperware/instantquery/instant_query_test.go

yeya24 · 2024-11-21T06:25:30Z

pkg/querier/tripperware/instantquery/instant_query.go

@@ -176,7 +183,7 @@ func (instantQueryCodec) EncodeResponse(ctx context.Context, res tripperware.Res
 		return nil, httpgrpc.Errorf(http.StatusInternalServerError, "invalid response format")
 	}

-	b, err := json.Marshal(a)
+	contentType, b, err := marshalResponse(a, req.Header.Get("Accept"))


Should we only try to return non JSON response format if we know the request comes from Ruler?
If the request doesn't come from Ruler I feel we want to enforce JSON response even though they specify protobuf in their protocol like Accept: protobuf, json. To reduce impact as this protobuf format is Cortex specific.

But tbh Idk if we can identify ruler requests or not. We probably have to do what you are doing here.

WDYT? @alanprot @danielblando

@yeya24
One way is to define a new MIME type like application/x-protobuf-ruler (just example) and then use it to accept header.

Yeah.. maybe we should define our internal MIME type so external calls will only get that format if they really wanna.
I would not add "ruler" in the format though. Maybe something like application/x-cortex.query.v1+proto (idk the latest best practices to define those custom Mime types :P )

This sounds worth trying. Let's hear from other maintainers, too

SungJin1212 · 2024-11-25T06:22:07Z

@yeya24 @alanprot
I defined the application/x-cortex-query+proto MIME type and used it to communicate between the Ruler and Query Frontend.

SungJin1212 · 2024-11-25T11:13:25Z

pkg/querier/tripperware/instantquery/instant_query_test.go

+		{
+			name:                "proto type",
+			acceptHeader:        "application/x-protobuf",
+			expectedContentType: "application/json",
+		},


We can enforce json marshal if we get the other types except the application/x-cortex-query+proto.

SungJin1212 · 2024-12-04T08:27:29Z

@yeya24
Could you take a look when you have time :D?

pkg/querier/codec/protobuf_codec_test.go

pkg/querier/tripperware/instantquery/instant_query.go

yeya24 · 2024-12-05T05:38:30Z

pkg/querier/tripperware/instantquery/instant_query.go

@@ -154,7 +161,7 @@ func (c instantQueryCodec) EncodeRequest(ctx context.Context, r tripperware.Requ
 		}
 	}

-	tripperware.SetRequestHeaders(h, c.defaultCodecType, c.compression)
+	tripperware.SetRequestHeaders(h, c.defaultCodecType, c.compression, c.rulerQueryResponseFormat)


This will always ask for cortex internal response format in instant query even though the requests are not initiated from Ruler?

oh.. I see.
We can utilize the User-Agent header value at DecodeRequest and EncodeRequest.

yeya24 · 2024-12-10T00:01:38Z

pkg/querier/tripperware/instantquery/instant_query.go

-			if strings.EqualFold(h, header) {
-				result.Headers[h] = hv
-				break
+	if isSourceRuler(&r.Header) {


Nit: We can initialize isSourceRuler(&r.Header) as a var.

yeya24 · 2024-12-10T00:08:42Z

pkg/querier/tripperware/instantquery/instant_query.go

+	return tripperware.ApplicationJson, b, err
+}
+
+func isSourceRuler(h *http.Header) bool {


I like that we identified query traffic and ruler traffic now.
Can we also add this to frontend/transport/handler.go? We can have additional logging about the query source and update metrics (like queries from rule or API, etc).

Do you mean to add a label (source) to query stats metrics?

h.queryScannedSamples = promauto.With(reg).NewCounterVec(prometheus.CounterOpts{ Name: "cortex_query_samples_scanned_total", Help: "Number of samples scanned to execute a query.", }, []string{"user", "source"})

And add a log field to ServeHTTP. right?

Do you mean to add a label (source) to query stats metrics?

We don't need to add this for all metrics. Only the counter for number of requests should be good enough

We don't have a counter metric in the handler.
Should I add it?
FYI. We have a cortex_ruler_query_frontend_request_duration_seconds histogram.

I add a User-Agent field to logger.
logger := util_log.WithUserAgent(userAgent, util_log.WithContext(r.Context(), f.log))

Umm, we actually already have one called queriesPerTenant but it is in roundtrip.go. Sorry for the wrong info.

thanks!
I add a source label to cortex_query_frontend_queries_total.
The label value is ruler or others.
WDYT?

Can we call it API and Rule? This aligns with https://github.com/cortexproject/cortex/blob/master/pkg/cortexpb/cortex.proto#L14

pkg/frontend/transport/handler.go

yeya24

It is in good shape now... Some nites only.
Would be good to have more people to review this change.

pkg/ruler/frontend_decoder.go

yeya24 · 2024-12-11T08:41:06Z

pkg/ruler/frontend_decoder.go

@@ -18,8 +20,18 @@ const (
 )

 type JsonDecoder struct{}
+type ProtobufDecoder struct{}
 type Warning []string


What do you mean is we should return nil instead of the resp.Warnings at the ProtobufDecoder decode func since the error is only decoded by the json codec?

No... Just the naming thing. Better to call it Warnings instead of Warning

Thanks, fixed it.

yeya24 · 2024-12-12T05:12:25Z

@rajagopalanand @rapphil Maybe you can help review this PR as well. Thanks

SungJin1212 · 2024-12-18T11:50:34Z

@alanprot
Could you take a look?

CharlieTLe · 2024-12-19T17:02:52Z

pkg/querier/codec/protobuf_codec_test.go

@@ -147,6 +149,7 @@ func TestProtobufCodec_Encode(t *testing.T) {
 			},
 		},
 		{
+			name: "matrix",


nit: missing name for this matrix test https://github.com/cortexproject/cortex/pull/6345/files#diff-9c2062cda65af13cf43cc9371cae0424744499ad5ca83597ea2bf4df2437f4d8R106

seems like it is the same as https://github.com/cortexproject/cortex/pull/6345/files#diff-9c2062cda65af13cf43cc9371cae0424744499ad5ca83597ea2bf4df2437f4d8R186?

Yeah.. let's add a name and delete the duplicate one. Thanks!

CharlieTLe · 2024-12-19T17:05:46Z

pkg/querier/tripperware/instantquery/instant_query.go

@@ -29,6 +31,9 @@ var (
 		SortMapKeys:            true,
 		ValidateJsonRawMessage: false,
 	}.Froze()
+
+	rulerMIMEType = v1.MIMEType{Type: "application", SubType: tripperware.QueryResponseCortexMIMESubType}


Can we do rulerMIMEType = tripperware.QueryResponseCortexMIMEType here?

The tripperware.QueryResponseCortexMIMEType is a string, so it is hard to use it to MIMEType directly.

CharlieTLe · 2024-12-19T17:09:58Z

pkg/querier/tripperware/instantquery/instant_query_test.go

+		{
+			name:                "empty accept header",
+			acceptHeader:        "",
+			expectedContentType: "application/json",


nit: can we do expectedContentType: tripperware.ApplicationJson instead?

CharlieTLe · 2024-12-19T17:21:40Z

pkg/querier/tripperware/query.go

@@ -763,7 +769,8 @@ func SetRequestHeaders(h http.Header, defaultCodecType CodecType, compression Co
 }

 func UnmarshalResponse(r *http.Response, buf []byte, resp *PrometheusResponse) error {
-	if r.Header != nil && r.Header.Get("Content-Type") == ApplicationProtobuf {
+	contentType := r.Header.Get("Content-Type")


Should we check that r.Header != nil before getting the contentType?

Thanks for catching!

CharlieTLe · 2024-12-19T17:26:45Z

pkg/querier/tripperware/queryrange/query_range.go

@@ -230,7 +230,7 @@ func (c prometheusCodec) DecodeResponse(ctx context.Context, r *http.Response, _
 	return &resp, nil
 }

-func (prometheusCodec) EncodeResponse(ctx context.Context, res tripperware.Response) (*http.Response, error) {
+func (prometheusCodec) EncodeResponse(ctx context.Context, _ *http.Request, res tripperware.Response) (*http.Response, error) {


Seems weird that we need to pass in a request here if it's not being used.

It is just for the interface.
It is not used in the range query, since the ruler only does instant query.

CharlieTLe · 2024-12-19T18:44:39Z

pkg/querier/tripperware/roundtrip.go

@@ -156,7 +156,8 @@ func NewQueryTripperware(
 				now := time.Now()
 				userStr := tenant.JoinTenantIDs(tenantIDs)
 				activeUsers.UpdateUserTimestamp(userStr, now)
-				queriesPerTenant.WithLabelValues(op, userStr).Inc()
+				source := getSource(r.Header.Get("User-Agent"))
+				queriesPerTenant.WithLabelValues(op, userStr, source).Inc()


We probably shouldn't have the user-agent as a possible label value.

If it can be set by the user, then there's a possibility of cardinality explosion.

Good call out. I think source is good enough and we don't need the user agent here

The label values would be ruler or api. Don't we?

Nvm ignore my comment. We are not using user agent value directly as label value.

Can we add a test here to make sure the user agent is either ruler or api to ensure we don't allow values other than those?

Signed-off-by: SungJin1212 <[email protected]>

rajagopalanand · 2024-12-24T16:21:30Z

pkg/querier/tripperware/instantquery/instant_query.go

+	isSourceRuler := strings.Contains(r.Header.Get("User-Agent"), tripperware.RulerUserAgent)
+	if isSourceRuler {
+		// When the source is the Ruler, then forward whole headers
+		result.Headers = r.Header


Nit: Consider returning here vs else block

pull-request-size bot added the size/XL label Nov 18, 2024

dosubot bot added the component/rules Bits & bobs todo with rules and alerts: the ruler, config service etc. label Nov 18, 2024

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch 4 times, most recently from 872120e to 5f93620 Compare November 18, 2024 06:56

SungJin1212 changed the title ~~Allow ruler to retrieve proto format query response~~ Support ruler to retrieve proto format query response Nov 19, 2024

yeya24 reviewed Nov 19, 2024

View reviewed changes

pkg/querier/tripperware/query.proto Outdated Show resolved Hide resolved

yeya24 reviewed Nov 19, 2024

View reviewed changes

pkg/ruler/frontend_decoder.go Outdated Show resolved Hide resolved

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch 2 times, most recently from 6915258 to fb61f5a Compare November 19, 2024 08:16

SungJin1212 requested a review from yeya24 November 19, 2024 08:36

SungJin1212 mentioned this pull request Nov 19, 2024

[Chore] Fix instant query test data type #6344

Closed

3 tasks

yeya24 reviewed Nov 19, 2024

View reviewed changes

pkg/querier/tripperware/instantquery/instant_query_test.go Show resolved Hide resolved

yeya24 reviewed Nov 21, 2024

View reviewed changes

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch from fb61f5a to 80d5e87 Compare November 25, 2024 06:20

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch from 80d5e87 to ec03c22 Compare November 25, 2024 11:11

SungJin1212 commented Nov 25, 2024

View reviewed changes

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch from ec03c22 to 6c4fe46 Compare November 26, 2024 00:09

SungJin1212 requested review from alanprot and yeya24 November 27, 2024 01:10

yeya24 mentioned this pull request Dec 3, 2024

Query Frontend: add new field for dense native histogram format #6199

Merged

3 tasks

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch from 6c4fe46 to 0b4bf23 Compare December 4, 2024 07:15

pull-request-size bot added size/XXL and removed size/XL labels Dec 4, 2024

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch 2 times, most recently from 86a0a52 to 27d40bd Compare December 4, 2024 07:25

yeya24 reviewed Dec 5, 2024

View reviewed changes

pull-request-size bot added the size/XL label Dec 5, 2024

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch 4 times, most recently from 194aaa3 to 5a369e5 Compare December 6, 2024 01:29

SungJin1212 requested a review from yeya24 December 6, 2024 01:35

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch from 5a369e5 to 26ee074 Compare December 6, 2024 02:11

yeya24 reviewed Dec 10, 2024

View reviewed changes

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch 3 times, most recently from 7c686bf to d988c9e Compare December 10, 2024 07:52

yeya24 reviewed Dec 11, 2024

View reviewed changes

pkg/frontend/transport/handler.go Outdated Show resolved Hide resolved

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch from d988c9e to d3f99c0 Compare December 11, 2024 01:30

yeya24 approved these changes Dec 11, 2024

View reviewed changes

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch from d3f99c0 to 437a332 Compare December 12, 2024 01:17

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch 4 times, most recently from 31f2cce to af3e05b Compare December 18, 2024 09:18

CharlieTLe reviewed Dec 19, 2024

View reviewed changes

Allow ruler to retrieve proto format query response

da4a4da

Signed-off-by: SungJin1212 <[email protected]>

SungJin1212 force-pushed the Add-protobuf-codec-to-ruler branch from af3e05b to da4a4da Compare December 20, 2024 02:36

rajagopalanand reviewed Dec 24, 2024

View reviewed changes

Support ruler to retrieve proto format query response #6345

Are you sure you want to change the base?

Support ruler to retrieve proto format query response #6345

Conversation

SungJin1212 commented Nov 18, 2024 • edited Loading

yeya24 Nov 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yeya24 Nov 21, 2024 • edited Loading

Choose a reason for hiding this comment

SungJin1212 commented Nov 25, 2024

SungJin1212 Nov 25, 2024 • edited Loading

Choose a reason for hiding this comment

SungJin1212 commented Dec 4, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SungJin1212 Dec 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SungJin1212 Dec 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SungJin1212 Dec 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yeya24 left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SungJin1212 Dec 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yeya24 commented Dec 12, 2024

SungJin1212 commented Dec 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SungJin1212 Dec 20, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rajagopalanand Dec 24, 2024 • edited Loading

Choose a reason for hiding this comment

SungJin1212 commented Nov 18, 2024 •

edited

Loading

yeya24 Nov 21, 2024 •

edited

Loading

yeya24 Nov 21, 2024 •

edited

Loading

SungJin1212 Nov 25, 2024 •

edited

Loading

SungJin1212 commented Dec 4, 2024 •

edited

Loading

SungJin1212 Dec 10, 2024 •

edited

Loading

SungJin1212 Dec 10, 2024 •

edited

Loading

SungJin1212 Dec 10, 2024 •

edited

Loading

yeya24 left a comment •

edited

Loading

SungJin1212 Dec 12, 2024 •

edited

Loading

SungJin1212 Dec 20, 2024 •

edited

Loading

rajagopalanand Dec 24, 2024 •

edited

Loading