Add streaming for ingester metadata APIs

CHANGELOG.md

@@ -14,6 +14,7 @@
 * [FEATURE] Compactor: Add `-compactor.skip-blocks-with-out-of-order-chunks-enabled` configuration to mark blocks containing index with out-of-order chunks for no compact instead of halting the compaction. #4707
 * [FEATURE] Querier/Query-Frontend: Add `-querier.per-step-stats-enabled` and `-frontend.cache-queryable-samples-stats` configurations to enable query sample statistics. #4708
 * [FEATURE] Add shuffle sharding for the compactor #4433
+* [FEATURE] Querier: Use streaming for ingester metdata APIs. #4725
 * [ENHANCEMENT] Update Go version to 1.17.8. #4602 #4604 #4658
 * [ENHANCEMENT] Keep track of discarded samples due to bad relabel configuration in `cortex_discarded_samples_total`. #4503
 * [ENHANCEMENT] Ruler: Add `-ruler.disable-rule-group-label` to disable the `rule_group` label on exported metrics. #4571

docs/blocks-storage/querier.md

@@ -118,6 +118,10 @@ querier:
   # CLI flag: -querier.ingester-streaming
   [ingester_streaming: <boolean> | default = true]
 
+  # Use streaming RPCs for metadata APIs from ingester.
+  # CLI flag: -querier.ingester-metadata-streaming
+  [ingester_metadata_streaming: <boolean> | default = false]
+
   # Maximum number of samples a single query can load into memory.
   # CLI flag: -querier.max-samples
   [max_samples: <int> | default = 50000000]

docs/configuration/config-file-reference.md

@@ -871,6 +871,10 @@ The `querier_config` configures the Cortex querier.
 # CLI flag: -querier.ingester-streaming
 [ingester_streaming: <boolean> | default = true]
 
+# Use streaming RPCs for metadata APIs from ingester.
+# CLI flag: -querier.ingester-metadata-streaming
+[ingester_metadata_streaming: <boolean> | default = false]
+
 # Maximum number of samples a single query can load into memory.
 # CLI flag: -querier.max-samples
 [max_samples: <int> | default = 50000000]

pkg/distributor/distributor.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	io "io"
 	"net/http"
 	"sort"
 	"strings"
@@ -871,8 +872,7 @@ func (d *Distributor) ForReplicationSet(ctx context.Context, replicationSet ring
 	})
 }
 
-// LabelValuesForLabelName returns all of the label values that are associated with a given label name.
-func (d *Distributor) LabelValuesForLabelName(ctx context.Context, from, to model.Time, labelName model.LabelName, matchers ...*labels.Matcher) ([]string, error) {
+func (d *Distributor) LabelValuesForLabelNameCommon(ctx context.Context, from, to model.Time, labelName model.LabelName, f func(ctx context.Context, rs ring.ReplicationSet, req *ingester_client.LabelValuesRequest) ([]interface{}, error), matchers ...*labels.Matcher) ([]string, error) {
 	replicationSet, err := d.GetIngestersForMetadata(ctx)
 	if err != nil {
 		return nil, err
@@ -883,16 +883,14 @@ func (d *Distributor) LabelValuesForLabelName(ctx context.Context, from, to mode
 		return nil, err
 	}
 
-	resps, err := d.ForReplicationSet(ctx, replicationSet, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
-		return client.LabelValues(ctx, req)
-	})
+	resps, err := f(ctx, replicationSet, req)
 	if err != nil {
 		return nil, err
 	}
 
 	valueSet := map[string]struct{}{}
 	for _, resp := range resps {
-		for _, v := range resp.(*ingester_client.LabelValuesResponse).LabelValues {
+		for _, v := range resp.([]string) {
 			valueSet[v] = struct{}{}
 		}
 	}
@@ -908,8 +906,46 @@ func (d *Distributor) LabelValuesForLabelName(ctx context.Context, from, to mode
 	return values, nil
 }
 
-// LabelNames returns all of the label names.
-func (d *Distributor) LabelNames(ctx context.Context, from, to model.Time) ([]string, error) {
+// LabelValuesForLabelName returns all of the label values that are associated with a given label name.
+func (d *Distributor) LabelValuesForLabelName(ctx context.Context, from, to model.Time, labelName model.LabelName, matchers ...*labels.Matcher) ([]string, error) {
+	return d.LabelValuesForLabelNameCommon(ctx, from, to, labelName, func(ctx context.Context, rs ring.ReplicationSet, req *ingester_client.LabelValuesRequest) ([]interface{}, error) {
+		return d.ForReplicationSet(ctx, rs, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
+			resp, err := client.LabelValues(ctx, req)
+			if err != nil {
+				return nil, err
+			}
+			return resp.LabelValues, nil
+		})
+	}, matchers...)
+}
+
+// LabelValuesForLabelName returns all of the label values that are associated with a given label name.
+func (d *Distributor) LabelValuesForLabelNameStream(ctx context.Context, from, to model.Time, labelName model.LabelName, matchers ...*labels.Matcher) ([]string, error) {
+	return d.LabelValuesForLabelNameCommon(ctx, from, to, labelName, func(ctx context.Context, rs ring.ReplicationSet, req *ingester_client.LabelValuesRequest) ([]interface{}, error) {
+		return d.ForReplicationSet(ctx, rs, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
+			stream, err := client.LabelValuesStream(ctx, req)
+			if err != nil {
+				return nil, err
+			}
+			defer stream.CloseSend() //nolint:errcheck
+			allLabelValues := []string{}
+			for {
+				resp, err := stream.Recv()
+
+				if err == io.EOF {
+					break
+				} else if err != nil {
+					return nil, err
+				}
+				allLabelValues = append(allLabelValues, resp.LabelValues...)
+			}
+
+			return allLabelValues, nil
+		})
+	}, matchers...)
+}
+
+func (d *Distributor) LabelNamesCommon(ctx context.Context, from, to model.Time, f func(ctx context.Context, rs ring.ReplicationSet, req *ingester_client.LabelNamesRequest) ([]interface{}, error)) ([]string, error) {
 	replicationSet, err := d.GetIngestersForMetadata(ctx)
 	if err != nil {
 		return nil, err
@@ -919,16 +955,14 @@ func (d *Distributor) LabelNames(ctx context.Context, from, to model.Time) ([]st
 		StartTimestampMs: int64(from),
 		EndTimestampMs:   int64(to),
 	}
-	resps, err := d.ForReplicationSet(ctx, replicationSet, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
-		return client.LabelNames(ctx, req)
-	})
+	resps, err := f(ctx, replicationSet, req)
 	if err != nil {
 		return nil, err
 	}
 
 	valueSet := map[string]struct{}{}
 	for _, resp := range resps {
-		for _, v := range resp.(*ingester_client.LabelNamesResponse).LabelNames {
+		for _, v := range resp.([]string) {
 			valueSet[v] = struct{}{}
 		}
 	}
@@ -943,8 +977,106 @@ func (d *Distributor) LabelNames(ctx context.Context, from, to model.Time) ([]st
 	return values, nil
 }
 
+func (d *Distributor) LabelNamesStream(ctx context.Context, from, to model.Time) ([]string, error) {
+	return d.LabelNamesCommon(ctx, from, to, func(ctx context.Context, rs ring.ReplicationSet, req *ingester_client.LabelNamesRequest) ([]interface{}, error) {
+		return d.ForReplicationSet(ctx, rs, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
+			stream, err := client.LabelNamesStream(ctx, req)
+			if err != nil {
+				return nil, err
+			}
+			defer stream.CloseSend() //nolint:errcheck
+			allLabelNames := []string{}
+			for {
+				resp, err := stream.Recv()
+
+				if err == io.EOF {
+					break
+				} else if err != nil {
+					return nil, err
+				}
+				allLabelNames = append(allLabelNames, resp.LabelNames...)
+			}
+
+			return allLabelNames, nil
+		})
+	})
+}
+
+// LabelNames returns all of the label names.
+func (d *Distributor) LabelNames(ctx context.Context, from, to model.Time) ([]string, error) {
+	return d.LabelNamesCommon(ctx, from, to, func(ctx context.Context, rs ring.ReplicationSet, req *ingester_client.LabelNamesRequest) ([]interface{}, error) {
+		return d.ForReplicationSet(ctx, rs, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
+			resp, err := client.LabelNames(ctx, req)
+			if err != nil {
+				return nil, err
+			}
+			return resp.LabelNames, nil
+		})
+	})
+}
+
 // MetricsForLabelMatchers gets the metrics that match said matchers
 func (d *Distributor) MetricsForLabelMatchers(ctx context.Context, from, through model.Time, matchers ...*labels.Matcher) ([]metric.Metric, error) {
+	return d.metricsForLabelMatchersCommon(ctx, from, through, func(ctx context.Context, rs ring.ReplicationSet, req *ingester_client.MetricsForLabelMatchersRequest, metrics *map[model.Fingerprint]model.Metric, mutex *sync.Mutex, queryLimiter *limiter.QueryLimiter) error {
+		_, err := d.ForReplicationSet(ctx, rs, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
+			resp, err := client.MetricsForLabelMatchers(ctx, req)
+			if err != nil {
+				return nil, err
+			}
+			ms := ingester_client.FromMetricsForLabelMatchersResponse(resp)
+			for _, m := range ms {
+				if err := queryLimiter.AddSeries(cortexpb.FromMetricsToLabelAdapters(m)); err != nil {
+					return nil, err
+				}
+				mutex.Lock()
+				(*metrics)[m.Fingerprint()] = m
+				mutex.Unlock()
+			}
+
+			return nil, nil
+		})
+
+		return err
+	}, matchers...)
+}
+
+func (d *Distributor) MetricsForLabelMatchersStream(ctx context.Context, from, through model.Time, matchers ...*labels.Matcher) ([]metric.Metric, error) {
+	return d.metricsForLabelMatchersCommon(ctx, from, through, func(ctx context.Context, rs ring.ReplicationSet, req *ingester_client.MetricsForLabelMatchersRequest, metrics *map[model.Fingerprint]model.Metric, mutex *sync.Mutex, queryLimiter *limiter.QueryLimiter) error {
+		_, err := d.ForReplicationSet(ctx, rs, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
+			stream, err := client.MetricsForLabelMatchersStream(ctx, req)
+			if err != nil {
+				return nil, err
+			}
+			defer stream.CloseSend() //nolint:errcheck
+			for {
+				resp, err := stream.Recv()
+
+				if err == io.EOF {
+					break
+				} else if err != nil {
+					return nil, err
+				}
+
+				for _, metric := range resp.Metric {
+					m := cortexpb.FromLabelAdaptersToMetric(metric.Labels)
+
+					if err := queryLimiter.AddSeries(cortexpb.FromMetricsToLabelAdapters(m)); err != nil {
+						return nil, err
+					}
+					mutex.Lock()
+					(*metrics)[m.Fingerprint()] = m
+					mutex.Unlock()
+				}
+			}
+
+			return nil, nil
+		})
+
+		return err
+	}, matchers...)
+}
+
+func (d *Distributor) metricsForLabelMatchersCommon(ctx context.Context, from, through model.Time, f func(context.Context, ring.ReplicationSet, *ingester_client.MetricsForLabelMatchersRequest, *map[model.Fingerprint]model.Metric, *sync.Mutex, *limiter.QueryLimiter) error, matchers ...*labels.Matcher) ([]metric.Metric, error) {
 	replicationSet, err := d.GetIngestersForMetadata(ctx)
 	queryLimiter := limiter.QueryLimiterFromContextWithFallback(ctx)
 	if err != nil {
@@ -958,23 +1090,7 @@ func (d *Distributor) MetricsForLabelMatchers(ctx context.Context, from, through
 	mutex := sync.Mutex{}
 	metrics := map[model.Fingerprint]model.Metric{}
 
-	_, err = d.ForReplicationSet(ctx, replicationSet, func(ctx context.Context, client ingester_client.IngesterClient) (interface{}, error) {
-		resp, err := client.MetricsForLabelMatchers(ctx, req)
-		if err != nil {
-			return nil, err
-		}
-		ms := ingester_client.FromMetricsForLabelMatchersResponse(resp)
-		for _, m := range ms {
-			if err := queryLimiter.AddSeries(cortexpb.FromMetricsToLabelAdapters(m)); err != nil {
-				return nil, err
-			}
-			mutex.Lock()
-			metrics[m.Fingerprint()] = m
-			mutex.Unlock()
-		}
-
-		return resp, nil
-	})
+	err = f(ctx, replicationSet, req, &metrics, &mutex, queryLimiter)
 
 	if err != nil {
 		return nil, err

pkg/distributor/distributor_test.go

@@ -2017,21 +2017,36 @@ func TestDistributor_MetricsForLabelMatchers(t *testing.T) {
 				require.NoError(t, err)
 			}
 
-			metrics, err := ds[0].MetricsForLabelMatchers(ctx, now, now, testData.matchers...)
+			{
+				metrics, err := ds[0].MetricsForLabelMatchers(ctx, now, now, testData.matchers...)
 
-			if testData.expectedErr != nil {
-				assert.EqualError(t, err, testData.expectedErr.Error())
-				return
+				if testData.expectedErr != nil {
+					assert.EqualError(t, err, testData.expectedErr.Error())
+					return
+				}
+
+				require.NoError(t, err)
+				assert.ElementsMatch(t, testData.expectedResult, metrics)
+
+				// Check how many ingesters have been queried.
+				// Due to the quorum the distributor could cancel the last request towards ingesters
+				// if all other ones are successful, so we're good either has been queried X or X-1
+				// ingesters.
+				assert.Contains(t, []int{testData.expectedIngesters, testData.expectedIngesters - 1}, countMockIngestersCalls(ingesters, "MetricsForLabelMatchers"))
 			}
 
-			require.NoError(t, err)
-			assert.ElementsMatch(t, testData.expectedResult, metrics)
+			{
+				metrics, err := ds[0].MetricsForLabelMatchersStream(ctx, now, now, testData.matchers...)
+				if testData.expectedErr != nil {
+					assert.EqualError(t, err, testData.expectedErr.Error())
+					return
+				}
 
-			// Check how many ingesters have been queried.
-			// Due to the quorum the distributor could cancel the last request towards ingesters
-			// if all other ones are successful, so we're good either has been queried X or X-1
-			// ingesters.
-			assert.Contains(t, []int{testData.expectedIngesters, testData.expectedIngesters - 1}, countMockIngestersCalls(ingesters, "MetricsForLabelMatchers"))
+				require.NoError(t, err)
+				assert.ElementsMatch(t, testData.expectedResult, metrics)
+
+				assert.Contains(t, []int{testData.expectedIngesters, testData.expectedIngesters - 1}, countMockIngestersCalls(ingesters, "MetricsForLabelMatchersStream"))
+			}
 		})
 	}
 }
@@ -2661,7 +2676,39 @@ func (i *mockIngester) QueryStream(ctx context.Context, req *client.QueryRequest
 			},
 		})
 	}
-	return &stream{
+	return &queryStream{
+		results: results,
+	}, nil
+}
+
+func (i *mockIngester) MetricsForLabelMatchersStream(ctx context.Context, req *client.MetricsForLabelMatchersRequest, opts ...grpc.CallOption) (client.Ingester_MetricsForLabelMatchersStreamClient, error) {
+	time.Sleep(i.queryDelay)
+	i.Lock()
+	defer i.Unlock()
+
+	i.trackCall("MetricsForLabelMatchersStream")
+
+	if !i.happy.Load() {
+		return nil, errFail
+	}
+
+	_, _, multiMatchers, err := client.FromMetricsForLabelMatchersRequest(req)
+	if err != nil {
+		return nil, err
+	}
+
+	results := []*client.MetricsForLabelMatchersStreamResponse{}
+	for _, matchers := range multiMatchers {
+		for _, ts := range i.timeseries {
+			if match(ts.Labels, matchers) {
+				results = append(results, &client.MetricsForLabelMatchersStreamResponse{
+					Metric: []*cortexpb.Metric{{Labels: ts.Labels}},
+				})
+			}
+		}
+	}
+
+	return &metricsForLabelMatchersStream{
 		results: results,
 	}, nil
 }
@@ -2742,17 +2789,36 @@ func (i *noopIngester) Push(ctx context.Context, req *cortexpb.WriteRequest, opt
 	return nil, nil
 }
 
-type stream struct {
+type queryStream struct {
 	grpc.ClientStream
 	i       int
 	results []*client.QueryStreamResponse
 }
 
-func (*stream) CloseSend() error {
+func (*queryStream) CloseSend() error {
+	return nil
+}
+
+func (s *queryStream) Recv() (*client.QueryStreamResponse, error) {
+	if s.i >= len(s.results) {
+		return nil, io.EOF
+	}
+	result := s.results[s.i]
+	s.i++
+	return result, nil
+}
+
+type metricsForLabelMatchersStream struct {
+	grpc.ClientStream
+	i       int
+	results []*client.MetricsForLabelMatchersStreamResponse
+}
+
+func (*metricsForLabelMatchersStream) CloseSend() error {
 	return nil
 }
 
-func (s *stream) Recv() (*client.QueryStreamResponse, error) {
+func (s *metricsForLabelMatchersStream) Recv() (*client.MetricsForLabelMatchersStreamResponse, error) {
 	if s.i >= len(s.results) {
 		return nil, io.EOF
 	}