Allow ruler to retrieve proto format query response

Signed-off-by: SungJin1212 <[email protected]>
cortexproject · Dec 4, 2024 · 0b4bf23 · 0b4bf23
1 parent 320e475
commit 0b4bf23
Show file tree

Hide file tree

Showing 22 changed files with 1,006 additions and 83 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@
 * [CHANGE] Change all max async concurrency default values `50` to `3` #6268
 * [CHANGE] Change default value of `-blocks-storage.bucket-store.index-cache.multilevel.max-async-concurrency` from `50` to `3` #6265
 * [CHANGE] Enable Compactor and Alertmanager in target all. #6204
+* [FEATURE] Ruler: Add an experimental flag `-ruler.query-response-format` to retrieve query response as a proto format. #6345
 * [FEATURE] Ruler: Pagination support for List Rules API. #6299
 * [FEATURE] Query Frontend/Querier: Add protobuf codec `-api.querier-default-codec` and the option to choose response compression type `-querier.response-compression`. #5527
 * [FEATURE] Ruler: Experimental: Add `ruler.frontend-address` to allow query to query frontends instead of ingesters. #6151

diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md
@@ -4298,6 +4298,11 @@ The `ruler_config` configures the Cortex ruler.
 # CLI flag: -ruler.frontend-address
 [frontend_address: <string> | default = ""]
 
+# [Experimental] Query response format to get query results from Query Frontend
+# when the rule evaluation. Supported values: json,protobuf
+# CLI flag: -ruler.query-response-format
+[query_response_format: <string> | default = "protobuf"]
+
 frontend_client:
   # gRPC client max receive message size (bytes).
   # CLI flag: -ruler.frontendClient.grpc-max-recv-msg-size

diff --git a/go.mod b/go.mod
@@ -80,6 +80,7 @@ require (
 	github.com/bboreham/go-loser v0.0.0-20230920113527-fcc2c21820a3
 	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/google/go-cmp v0.6.0
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822
 	github.com/sercand/kuberesolver/v5 v5.1.1
 	go.opentelemetry.io/collector/pdata v1.20.0
 	golang.org/x/exp v0.0.0-20240613232115-7f521ea00fb8
@@ -188,7 +189,6 @@ require (
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
-	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect
 	github.com/ncw/swift v1.0.53 // indirect
 	github.com/oklog/run v1.1.0 // indirect

diff --git a/integration/ruler_test.go b/integration/ruler_test.go
@@ -1670,42 +1670,61 @@ func TestRulerEvalWithQueryFrontend(t *testing.T) {
 	distributor := e2ecortex.NewDistributor("distributor", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
 	ingester := e2ecortex.NewIngester("ingester", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
 	require.NoError(t, s.StartAndWaitReady(distributor, ingester))
-	queryFrontend := e2ecortex.NewQueryFrontend("query-frontend", flags, "")
-	require.NoError(t, s.Start(queryFrontend))
-
-	ruler := e2ecortex.NewRuler("ruler", consul.NetworkHTTPEndpoint(), mergeFlags(flags, map[string]string{
-		"-ruler.frontend-address": queryFrontend.NetworkGRPCEndpoint(),
-	}), "")
-	querier := e2ecortex.NewQuerier("querier", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(flags, map[string]string{
-		"-querier.frontend-address": queryFrontend.NetworkGRPCEndpoint(),
-	}), "")
-	require.NoError(t, s.StartAndWaitReady(ruler, querier))
-
-	c, err := e2ecortex.NewClient("", "", "", ruler.HTTPEndpoint(), user)
-	require.NoError(t, err)
+	for _, format := range []string{"protobuf", "json"} {
+		t.Run(fmt.Sprintf("format:%s", format), func(t *testing.T) {
+			queryFrontendFlag := mergeFlags(flags, map[string]string{
+				"-ruler.query-response-format": format,
+			})
+			queryFrontend := e2ecortex.NewQueryFrontend("query-frontend", queryFrontendFlag, "")
+			require.NoError(t, s.Start(queryFrontend))
 
-	expression := "metric"
-	groupName := "rule_group"
-	ruleName := "rule_name"
-	require.NoError(t, c.SetRuleGroup(ruleGroupWithRule(groupName, ruleName, expression), namespace))
+			querier := e2ecortex.NewQuerier("querier", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(queryFrontendFlag, map[string]string{
+				"-querier.frontend-address": queryFrontend.NetworkGRPCEndpoint(),
+			}), "")
+			require.NoError(t, s.StartAndWaitReady(querier))
 
-	rgMatcher := ruleGroupMatcher(user, namespace, groupName)
-	// Wait until ruler has loaded the group.
-	require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(1), []string{"cortex_prometheus_rule_group_rules"}, e2e.WithLabelMatchers(rgMatcher), e2e.WaitMissingMetrics))
-	// Wait until rule group has tried to evaluate the rule.
-	require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.GreaterOrEqual(1), []string{"cortex_prometheus_rule_evaluations_total"}, e2e.WithLabelMatchers(rgMatcher), e2e.WaitMissingMetrics))
+			rulerFlag := mergeFlags(queryFrontendFlag, map[string]string{
+				"-ruler.frontend-address": queryFrontend.NetworkGRPCEndpoint(),
+			})
+			ruler := e2ecortex.NewRuler("ruler", consul.NetworkHTTPEndpoint(), rulerFlag, "")
+			require.NoError(t, s.StartAndWaitReady(ruler))
 
-	matcher := labels.MustNewMatcher(labels.MatchEqual, "user", user)
-	// Check that cortex_ruler_query_frontend_clients went up
-	require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(1), []string{"cortex_ruler_query_frontend_clients"}, e2e.WaitMissingMetrics))
-	// Check that cortex_ruler_queries_total went up
-	require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.GreaterOrEqual(1), []string{"cortex_ruler_queries_total"}, e2e.WithLabelMatchers(matcher), e2e.WaitMissingMetrics))
-	// Check that cortex_ruler_queries_failed_total is zero
-	require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"cortex_ruler_queries_failed_total"}, e2e.WithLabelMatchers(matcher), e2e.WaitMissingMetrics))
-	// Check that cortex_ruler_write_requests_total went up
-	require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.GreaterOrEqual(1), []string{"cortex_ruler_write_requests_total"}, e2e.WithLabelMatchers(matcher), e2e.WaitMissingMetrics))
-	// Check that cortex_ruler_write_requests_failed_total is zero
-	require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"cortex_ruler_write_requests_failed_total"}, e2e.WithLabelMatchers(matcher), e2e.WaitMissingMetrics))
+			t.Cleanup(func() {
+				_ = s.Stop(ruler)
+				_ = s.Stop(queryFrontend)
+				_ = s.Stop(querier)
+			})
+
+			c, err := e2ecortex.NewClient("", "", "", ruler.HTTPEndpoint(), user)
+			require.NoError(t, err)
+
+			expression := "metric" // vector
+			//expression := "scalar(count(up == 1)) > bool 1" // scalar
+			groupName := "rule_group"
+			ruleName := "rule_name"
+			require.NoError(t, c.SetRuleGroup(ruleGroupWithRule(groupName, ruleName, expression), namespace))
+
+			rgMatcher := ruleGroupMatcher(user, namespace, groupName)
+			// Wait until ruler has loaded the group.
+			require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(1), []string{"cortex_prometheus_rule_group_rules"}, e2e.WithLabelMatchers(rgMatcher), e2e.WaitMissingMetrics))
+			// Wait until rule group has tried to evaluate the rule.
+			require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.GreaterOrEqual(1), []string{"cortex_prometheus_rule_evaluations_total"}, e2e.WithLabelMatchers(rgMatcher), e2e.WaitMissingMetrics))
+			// Make sure not to fail
+			require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"cortex_prometheus_rule_evaluation_failures_total"}, e2e.WithLabelMatchers(rgMatcher), e2e.WaitMissingMetrics))
+
+			matcher := labels.MustNewMatcher(labels.MatchEqual, "user", user)
+			// Check that cortex_ruler_query_frontend_clients went up
+			require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(1), []string{"cortex_ruler_query_frontend_clients"}, e2e.WaitMissingMetrics))
+			// Check that cortex_ruler_queries_total went up
+			require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.GreaterOrEqual(1), []string{"cortex_ruler_queries_total"}, e2e.WithLabelMatchers(matcher), e2e.WaitMissingMetrics))
+			// Check that cortex_ruler_queries_failed_total is zero
+			require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"cortex_ruler_queries_failed_total"}, e2e.WithLabelMatchers(matcher), e2e.WaitMissingMetrics))
+			// Check that cortex_ruler_write_requests_total went up
+			require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.GreaterOrEqual(1), []string{"cortex_ruler_write_requests_total"}, e2e.WithLabelMatchers(matcher), e2e.WaitMissingMetrics))
+			// Check that cortex_ruler_write_requests_failed_total is zero
+			require.NoError(t, ruler.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"cortex_ruler_write_requests_failed_total"}, e2e.WithLabelMatchers(matcher), e2e.WaitMissingMetrics))
+		})
+	}
 }
 
 func parseAlertFromRule(t *testing.T, rules interface{}) *alertingRule {

diff --git a/pkg/cortex/modules.go b/pkg/cortex/modules.go
@@ -475,7 +475,7 @@ func (t *Cortex) initQueryFrontendTripperware() (serv services.Service, err erro
 	prometheusCodec := queryrange.NewPrometheusCodec(false, t.Cfg.Querier.ResponseCompression, 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, t.Cfg.Querier.ResponseCompression, t.Cfg.API.QuerierDefaultCodec)
-	instantQueryCodec := instantquery.NewInstantQueryCodec(t.Cfg.Querier.ResponseCompression, t.Cfg.API.QuerierDefaultCodec)
+	instantQueryCodec := instantquery.NewInstantQueryCodec(t.Cfg.Querier.ResponseCompression, t.Cfg.API.QuerierDefaultCodec, t.Cfg.Ruler.QueryResponseFormat)
 
 	queryRangeMiddlewares, cache, err := queryrange.Middlewares(
 		t.Cfg.QueryRange,

diff --git a/pkg/querier/codec/protobuf_codec.go b/pkg/querier/codec/protobuf_codec.go
@@ -25,8 +25,8 @@ func (p ProtobufCodec) ContentType() v1.MIMEType {
 	if !p.CortexInternal {
 		return v1.MIMEType{Type: "application", SubType: "x-protobuf"}
 	}
-	// TODO: switch to use constants.
-	return v1.MIMEType{Type: "application", SubType: "x-cortex-query+proto"}
+
+	return v1.MIMEType{Type: "application", SubType: tripperware.QueryResponseCortexMIMESubType}
 }
 
 func (p ProtobufCodec) CanEncode(resp *v1.Response) bool {