query-frontend: pass on limit parameter for label names and values requests

CHANGELOG.md

@@ -19,6 +19,7 @@
 * [ENHANCEMENT] Store-gateway: add `outcome` label to `cortex_bucket_stores_gate_duration_seconds` histogram metric. Possible values for the `outcome` label are: `rejected_canceled`, `rejected_deadline_exceeded`, `rejected_other`, and `permitted`. #7784
 * [ENHANCEMENT] Query-frontend: use zero-allocation experimental decoder for active series queries via `-query-frontend.use-active-series-decoder`. #7665
 * [ENHANCEMENT] Go: updated to 1.22.2. #7802
+* [ENHANCEMENT] Query-frontend: support `limit` parameter on `/prometheus/api/v1/label/{name}/values` and `/prometheus/api/v1/labels` endpoints. #7722
 * [BUGFIX] Rules: improve error handling when querier is local to the ruler. #7567
 * [BUGFIX] Querier, store-gateway: Protect against panics raised during snappy encoding. #7520
 * [BUGFIX] Ingester: Prevent timely compaction of empty blocks. #7624

pkg/frontend/querymiddleware/codec.go

@@ -138,6 +138,8 @@ type LabelsQueryRequest interface {
 	// to and from the http request format without needing to undo the Prometheus parser converting between formats
 	// like `up{job="prometheus"}` and `{__name__="up, job="prometheus"}`, or other idiosyncrasies.
 	GetLabelMatcherSets() []string
+	// GetLimit returns the limit of the number of items in the response.
+	GetLimit() uint64
 	// AddSpanTags writes information about this request to an OpenTracing span
 	AddSpanTags(opentracing.Span)
 }
@@ -309,12 +311,21 @@ func (prometheusCodec) DecodeLabelsQueryRequest(_ context.Context, r *http.Reque
 
 	labelMatcherSets := reqValues["match[]"]
 
+	limit := uint64(0) // 0 means unlimited
+	if limitStr := reqValues.Get("limit"); limitStr != "" {
+		limit, err = strconv.ParseUint(limitStr, 10, 64)
+		if err != nil || limit == 0 {
+			return nil, apierror.New(apierror.TypeBadData, fmt.Sprintf("limit parameter must be a positive number: %s", limitStr))
+		}
+	}
+
 	if IsLabelNamesQuery(r.URL.Path) {
 		return &PrometheusLabelNamesQueryRequest{
 			Path:             r.URL.Path,
 			Start:            start,
 			End:              end,
 			LabelMatcherSets: labelMatcherSets,
+			Limit:            limit,
 		}, nil
 	}
 	// else, must be Label Values Request due to IsLabelsQuery check at beginning of func
@@ -324,6 +335,7 @@ func (prometheusCodec) DecodeLabelsQueryRequest(_ context.Context, r *http.Reque
 		Start:            start,
 		End:              end,
 		LabelMatcherSets: labelMatcherSets,
+		Limit:            limit,
 	}, nil
 }
 
@@ -524,6 +536,9 @@ func (c prometheusCodec) EncodeLabelsQueryRequest(ctx context.Context, req Label
 		if len(req.GetLabelMatcherSets()) > 0 {
 			urlValues["match[]"] = req.GetLabelMatcherSets()
 		}
+		if req.GetLimit() > 0 {
+			urlValues["limit"] = []string{strconv.FormatUint(req.GetLimit(), 10)}
+		}
 		u = &url.URL{
 			Path:     req.Path,
 			RawQuery: urlValues.Encode(),
@@ -541,6 +556,9 @@ func (c prometheusCodec) EncodeLabelsQueryRequest(ctx context.Context, req Label
 		if len(req.GetLabelMatcherSets()) > 0 {
 			urlValues["match[]"] = req.GetLabelMatcherSets()
 		}
+		if req.GetLimit() > 0 {
+			urlValues["limit"] = []string{strconv.FormatUint(req.GetLimit(), 10)}
+		}
 		u = &url.URL{
 			Path:     req.Path, // path still contains label name
 			RawQuery: urlValues.Encode(),

pkg/frontend/querymiddleware/codec_test.go

@@ -121,7 +121,8 @@ func TestLabelsQueryRequest(t *testing.T) {
 		expectedGetLabelName      string
 		expectedGetStartOrDefault int64
 		expectedGetEndOrDefault   int64
-		expectedErr               error
+		expectedErr               string
+		expectedLimit             uint64
 	}{
 		{
 			name: "label names with start and end timestamps, no matcher sets",
@@ -205,7 +206,7 @@ func TestLabelsQueryRequest(t *testing.T) {
 			expectedGetEndOrDefault:   1708588800 * 1e3,
 		},
 		{
-			name: "label names with start timestamp, no end timestamp, multiple matcher sets",
+			name: "label names with start and end timestamp, multiple matcher sets",
 			url:  "/api/v1/labels?end=1708588800&match%5B%5D=go_goroutines%7Bcontainer%3D~%22quer.%2A%22%7D&match%5B%5D=go_goroutines%7Bcontainer%21%3D%22query-scheduler%22%7D&start=1708502400",
 			expectedStruct: &PrometheusLabelNamesQueryRequest{
 				Path:  "/api/v1/labels",
@@ -221,7 +222,7 @@ func TestLabelsQueryRequest(t *testing.T) {
 			expectedGetEndOrDefault:   1708588800 * 1e3,
 		},
 		{
-			name: "label values with start timestamp, no end timestamp, multiple matcher sets",
+			name: "label values with start and end timestamp, multiple matcher sets",
 			url:  "/api/v1/label/job/values?end=1708588800&match%5B%5D=go_goroutines%7Bcontainer%3D~%22quer.%2A%22%7D&match%5B%5D=go_goroutines%7Bcontainer%21%3D%22query-scheduler%22%7D&start=1708502400",
 			expectedStruct: &PrometheusLabelValuesQueryRequest{
 				Path:      "/api/v1/label/job/values",
@@ -237,6 +238,53 @@ func TestLabelsQueryRequest(t *testing.T) {
 			expectedGetStartOrDefault: 1708502400 * 1e3,
 			expectedGetEndOrDefault:   1708588800 * 1e3,
 		},
+		{
+			name: "label names with start and end timestamp, multiple matcher sets, limit",
+			url:  "/api/v1/labels?end=1708588800&limit=10&match%5B%5D=go_goroutines%7Bcontainer%3D~%22quer.%2A%22%7D&match%5B%5D=go_goroutines%7Bcontainer%21%3D%22query-scheduler%22%7D&start=1708502400",
+			expectedStruct: &PrometheusLabelNamesQueryRequest{
+				Path:  "/api/v1/labels",
+				Start: 1708502400 * 1e3,
+				End:   1708588800 * 1e3,
+				Limit: 10,
+				LabelMatcherSets: []string{
+					"go_goroutines{container=~\"quer.*\"}",
+					"go_goroutines{container!=\"query-scheduler\"}",
+				},
+			},
+			expectedGetLabelName:      "",
+			expectedLimit:             10,
+			expectedGetStartOrDefault: 1708502400 * 1e3,
+			expectedGetEndOrDefault:   1708588800 * 1e3,
+		},
+		{
+			name: "label values with start and end timestamp, multiple matcher sets, limit",
+			url:  "/api/v1/label/job/values?end=1708588800&limit=10&match%5B%5D=go_goroutines%7Bcontainer%3D~%22quer.%2A%22%7D&match%5B%5D=go_goroutines%7Bcontainer%21%3D%22query-scheduler%22%7D&start=1708502400",
+			expectedStruct: &PrometheusLabelValuesQueryRequest{
+				Path:      "/api/v1/label/job/values",
+				LabelName: "job",
+				Start:     1708502400 * 1e3,
+				End:       1708588800 * 1e3,
+				Limit:     10,
+				LabelMatcherSets: []string{
+					"go_goroutines{container=~\"quer.*\"}",
+					"go_goroutines{container!=\"query-scheduler\"}",
+				},
+			},
+			expectedGetLabelName:      "job",
+			expectedLimit:             10,
+			expectedGetStartOrDefault: 1708502400 * 1e3,
+			expectedGetEndOrDefault:   1708588800 * 1e3,
+		},
+		{
+			name:        "zero limit is not allowed",
+			url:         "/api/v1/label/job/values?limit=0",
+			expectedErr: "limit parameter must be a positive number: 0",
+		},
+		{
+			name:        "negative limit is not allowed",
+			url:         "/api/v1/label/job/values?limit=-1",
+			expectedErr: "limit parameter must be a positive number: -1",
+		},
 	} {
 		t.Run(testCase.name, func(t *testing.T) {
 			for _, reqMethod := range []string{http.MethodGet, http.MethodPost} {
@@ -261,13 +309,14 @@ func TestLabelsQueryRequest(t *testing.T) {
 				r = r.WithContext(ctx)
 
 				reqDecoded, err := codec.DecodeLabelsQueryRequest(ctx, r)
-				if err != nil || testCase.expectedErr != nil {
-					require.EqualValues(t, testCase.expectedErr, err)
+				if err != nil || testCase.expectedErr != "" {
+					require.EqualError(t, err, testCase.expectedErr)
 					return
 				}
 				require.EqualValues(t, testCase.expectedStruct, reqDecoded)
 				require.EqualValues(t, testCase.expectedGetStartOrDefault, reqDecoded.GetStartOrDefault())
 				require.EqualValues(t, testCase.expectedGetEndOrDefault, reqDecoded.GetEndOrDefault())
+				require.EqualValues(t, testCase.expectedLimit, reqDecoded.GetLimit())
 
 				reqEncoded, err := codec.EncodeLabelsQueryRequest(context.Background(), reqDecoded)
 				require.NoError(t, err)

pkg/frontend/querymiddleware/labels_query_cache.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strconv"
 	"strings"
 	"time"
 
@@ -78,6 +79,7 @@ func (g DefaultCacheKeyGenerator) LabelValues(r *http.Request) (*GenericQueryCac
 		labelValuesReq.GetEndOrDefault(),
 		labelValuesReq.GetLabelName(),
 		labelMatcherSets,
+		labelValuesReq.GetLimit(),
 	)
 
 	return &GenericQueryCacheKey{
@@ -86,7 +88,7 @@ func (g DefaultCacheKeyGenerator) LabelValues(r *http.Request) (*GenericQueryCac
 	}, nil
 }
 
-func generateLabelsQueryRequestCacheKey(startTime, endTime int64, labelName string, matcherSets [][]*labels.Matcher) string {
+func generateLabelsQueryRequestCacheKey(startTime, endTime int64, labelName string, matcherSets [][]*labels.Matcher, limit uint64) string {
 	var (
 		twoHoursMillis = (2 * time.Hour).Milliseconds()
 		b              = strings.Builder{}
@@ -121,6 +123,12 @@ func generateLabelsQueryRequestCacheKey(startTime, endTime int64, labelName stri
 	b.WriteRune(stringParamSeparator)
 	b.WriteString(util.MultiMatchersStringer(matcherSets).String())
 
+	// if limit is set, then it will be a positive number
+	if limit > 0 {
+		b.WriteRune(stringParamSeparator)
+		b.WriteString(strconv.Itoa(int(limit)))
+	}
+
 	return b.String()
 }
 

pkg/frontend/querymiddleware/labels_query_cache_test.go

@@ -177,6 +177,7 @@ func TestGenerateLabelsQueryRequestCacheKey(t *testing.T) {
 		labelName        string
 		matcherSets      [][]*labels.Matcher
 		expectedCacheKey string
+		limit            uint64
 	}{
 		"start and end time are aligned to 2h boundaries": {
 			startTime: mustParseTime("2023-07-05T00:00:00Z"),
@@ -256,11 +257,32 @@ func TestGenerateLabelsQueryRequestCacheKey(t *testing.T) {
 				`{first="1",second!="2"},{first!="0"}`,
 			}, string(stringParamSeparator)),
 		},
+		"multiple label matcher sets, label name, and limit": {
+			startTime: mustParseTime("2023-07-05T00:00:00Z"),
+			endTime:   mustParseTime("2023-07-05T06:00:00Z"),
+			labelName: "test",
+			matcherSets: [][]*labels.Matcher{
+				{
+					labels.MustNewMatcher(labels.MatchEqual, "first", "1"),
+					labels.MustNewMatcher(labels.MatchNotEqual, "second", "2"),
+				}, {
+					labels.MustNewMatcher(labels.MatchNotEqual, "first", "0"),
+				},
+			},
+			limit: 10,
+			expectedCacheKey: strings.Join([]string{
+				fmt.Sprintf("%d", mustParseTime("2023-07-05T00:00:00Z")),
+				fmt.Sprintf("%d", mustParseTime("2023-07-05T06:00:00Z")),
+				"test",
+				`{first="1",second!="2"},{first!="0"}`,
+				"10",
+			}, string(stringParamSeparator)),
+		},
 	}
 
 	for testName, testData := range tests {
 		t.Run(testName, func(t *testing.T) {
-			assert.Equal(t, testData.expectedCacheKey, generateLabelsQueryRequestCacheKey(testData.startTime, testData.endTime, testData.labelName, testData.matcherSets))
+			assert.Equal(t, testData.expectedCacheKey, generateLabelsQueryRequestCacheKey(testData.startTime, testData.endTime, testData.labelName, testData.matcherSets, testData.limit))
 		})
 	}
 }

pkg/frontend/querymiddleware/model_extra.go

@@ -336,6 +336,8 @@ type PrometheusLabelNamesQueryRequest struct {
 	LabelMatcherSets []string
 	// ID of the request used to correlate downstream requests and responses.
 	ID int64
+	// Limit the number of label names returned. A value of 0 means no limit
+	Limit uint64
 }
 
 func (r *PrometheusLabelNamesQueryRequest) GetPath() string {
@@ -358,6 +360,10 @@ func (r *PrometheusLabelNamesQueryRequest) GetID() int64 {
 	return r.ID
 }
 
+func (r *PrometheusLabelNamesQueryRequest) GetLimit() uint64 {
+	return r.Limit
+}
+
 type PrometheusLabelValuesQueryRequest struct {
 	Path      string
 	LabelName string
@@ -370,6 +376,8 @@ type PrometheusLabelValuesQueryRequest struct {
 	LabelMatcherSets []string
 	// ID of the request used to correlate downstream requests and responses.
 	ID int64
+	// Limit the number of label values returned. A value of 0 means no limit.
+	Limit uint64
 }
 
 func (r *PrometheusLabelValuesQueryRequest) GetLabelName() string {
@@ -393,6 +401,10 @@ func (r *PrometheusLabelValuesQueryRequest) GetID() int64 {
 	return r.ID
 }
 
+func (r *PrometheusLabelValuesQueryRequest) GetLimit() uint64 {
+	return r.Limit
+}
+
 func (d *PrometheusData) UnmarshalJSON(b []byte) error {
 	v := struct {
 		Type   model.ValueType    `json:"resultType"`