Fix tenant federation performance hotspot

CHANGELOG.md

@@ -42,6 +42,7 @@
 * [ENHANCEMENT] New options `-server.http-listen-network` and `-server.grpc-listen-network` allow binding as 'tcp4' or 'tcp6'. #4462
 * [ENHANCEMENT] Rulers: Using shuffle sharding subring on GetRules API. #4466
 * [ENHANCEMENT] Support memcached auto-discovery via `auto-discovery` flag, introduced by thanos in https://github.com/thanos-io/thanos/pull/4487. Both AWS and Google Cloud memcached service support auto-discovery, which returns a list of nodes of the memcached cluster. #4412
+* [ENHANCEMENT] Query federation: improve performance in MergeQueryable by memoizing labels. #4502
 * [BUGFIX] Fixes a panic in the query-tee when comparing result. #4465
 * [BUGFIX] Frontend: Fixes @ modifier functions (start/end) when splitting queries by time. #4464
 * [BUGFIX] Compactor: compactor will no longer try to compact blocks that are already marked for deletion. Previously compactor would consider blocks marked for deletion within `-compactor.deletion-delay / 2` period as eligible for compaction. #4328

pkg/querier/tenantfederation/merge_queryable.go

@@ -387,20 +387,26 @@ func filterValuesByMatchers(idLabelName string, ids []string, matchers ...*label
 }
 
 type addLabelsSeriesSet struct {
-	upstream storage.SeriesSet
-	labels   labels.Labels
+	upstream   storage.SeriesSet
+	labels     labels.Labels
+	currSeries storage.Series
 }
 
 func (m *addLabelsSeriesSet) Next() bool {
+	m.currSeries = nil
 	return m.upstream.Next()
 }
 
 // At returns full series. Returned series should be iteratable even after Next is called.
 func (m *addLabelsSeriesSet) At() storage.Series {
-	return &addLabelsSeries{
-		upstream: m.upstream.At(),
-		labels:   m.labels,
+	if m.currSeries == nil {
+		upstream := m.upstream.At()
+		m.currSeries = &addLabelsSeries{
+			upstream: upstream,
+			labels:   setLabelsRetainExisting(upstream.Labels(), m.labels...),
+		}
 	}
+	return m.currSeries
 }
 
 // The error that iteration as failed with.
@@ -441,7 +447,7 @@ type addLabelsSeries struct {
 
 // Labels returns the complete set of labels. For series it means all labels identifying the series.
 func (a *addLabelsSeries) Labels() labels.Labels {
-	return setLabelsRetainExisting(a.upstream.Labels(), a.labels...)
+	return a.labels
 }
 
 // Iterator returns a new, independent iterator of the data of the series.

pkg/querier/tenantfederation/merge_queryable_test.go

@@ -277,6 +277,8 @@ type selectTestCase struct {
 	matchers []*labels.Matcher
 	// expectedSeriesCount is the expected number of series returned by a Select filtered by the Matchers in selector.
 	expectedSeriesCount int
+	// expectedLabels is the expected label sets returned by a Select filtered by the Matchers in selector.
+	expectedLabels []labels.Labels
 	// expectedWarnings is a slice of storage.Warnings messages expected when querying.
 	expectedWarnings []string
 	// expectedQueryErr is the error expected when querying.
@@ -425,6 +427,41 @@ func TestMergeQueryable_Select(t *testing.T) {
 				{
 					name:                "should return all series when no matchers are provided",
 					expectedSeriesCount: 6,
+					expectedLabels: []labels.Labels{
+						{
+							{Name: "__tenant_id__", Value: "team-a"},
+							{Name: "instance", Value: "host1"},
+							{Name: "original___tenant_id__", Value: "original-value"},
+							{Name: "tenant-team-a", Value: "static"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-a"},
+							{Name: "instance", Value: "host2.team-a"},
+							{Name: "original___tenant_id__", Value: "original-value"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-b"},
+							{Name: "instance", Value: "host1"},
+							{Name: "original___tenant_id__", Value: "original-value"},
+							{Name: "tenant-team-b", Value: "static"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-b"},
+							{Name: "instance", Value: "host2.team-b"},
+							{Name: "original___tenant_id__", Value: "original-value"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-c"},
+							{Name: "instance", Value: "host1"},
+							{Name: "original___tenant_id__", Value: "original-value"},
+							{Name: "tenant-team-c", Value: "static"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-c"},
+							{Name: "instance", Value: "host2.team-c"},
+							{Name: "original___tenant_id__", Value: "original-value"},
+						},
+					},
 				},
 				{
 					name:                "should return only series for team-a and team-c tenants when there is with not-equals matcher for the team-b tenant",
@@ -493,9 +530,16 @@ func TestMergeQueryable_Select(t *testing.T) {
 						assertEqualWarnings(t, tc.expectedWarnings, seriesSet.Warnings())
 					}
 
+					if tc.expectedLabels != nil {
+						require.Equal(t, len(tc.expectedLabels), tc.expectedSeriesCount)
+					}
+
 					count := 0
-					for seriesSet.Next() {
+					for i := 0; seriesSet.Next(); i++ {
 						count++
+						if tc.expectedLabels != nil {
+							require.Equal(t, tc.expectedLabels[i], seriesSet.At().Labels(), fmt.Sprintf("labels index: %d", i))
+						}
 					}
 					require.Equal(t, tc.expectedSeriesCount, count)
 				})