Distributor accept multiple HA Tracker pairs in the same request (find first)

CHANGELOG.md

@@ -15,6 +15,7 @@
 * [FEATURE] Ruler: Minimize chances of missed rule group evaluations that can occur due to OOM kills, bad underlying nodes, or due to an unhealthy ruler that appears in the ring as healthy. This feature is enabled via `-ruler.enable-ha-evaluation` flag. #6129
 * [FEATURE] Store Gateway: Add an in-memory chunk cache. #6245
 * [FEATURE] Chunk Cache: Support multi level cache and add metrics. #6249
+* [FEATURE] Distributor: Accept multiple HA Tracker pairs in the same request. #6256
 * [ENHANCEMENT] OTLP: Add `-distributor.otlp-max-recv-msg-size` flag to limit OTLP request size in bytes. #6333
 * [ENHANCEMENT] S3 Bucket Client: Add a list objects version configs to configure list api object version. #6280
 * [ENHANCEMENT] OpenStack Swift: Add application credential configs for Openstack swift object storage backend. #6255

docs/configuration/config-file-reference.md

@@ -3250,6 +3250,12 @@ The `limits_config` configures default and per-tenant limits imposed by Cortex s
 # CLI flag: -distributor.ha-tracker.enable-for-all-users
 [accept_ha_samples: <boolean> | default = false]
 
+# [Experimental] Flag to enable handling of samples with mixed external labels
+# identifying replicas in an HA Prometheus setup. Supported only if
+# -distributor.ha-tracker.enable-for-all-users is true.
+# CLI flag: -experimental.distributor.ha-tracker.mixed-ha-samples
+[accept_mixed_ha_samples: <boolean> | default = false]
+
 # Prometheus label to look for in samples to identify a Prometheus HA cluster.
 # CLI flag: -distributor.ha-tracker.cluster
 [ha_cluster_label: <string> | default = "cluster"]

docs/configuration/v1-guarantees.md

@@ -54,7 +54,9 @@ Currently experimental features are:
 - Metric relabeling in the distributor.
 - Scalable query-frontend (when using query-scheduler)
 - Ingester: do not unregister from ring on shutdown (`-ingester.unregister-on-shutdown=false`)
-- Distributor: do not extend writes on unhealthy ingesters (`-distributor.extend-writes=false`)
+- Distributor:
+  - Do not extend writes on unhealthy ingesters (`-distributor.extend-writes=false`)
+  - Accept multiple HA pairs in the same request (enabled via `-experimental.distributor.ha-tracker.mixed-ha-samples=true`)
 - Tenant Deletion in Purger, for blocks storage.
 - Query-frontend: query stats tracking (`-frontend.query-stats-enabled`)
 - Blocks storage bucket index

integration/distributor_mixed_ha_samples_test.go

@@ -0,0 +1,140 @@
+//go:build requires_docker
+// +build requires_docker
+
+package integration
+
+import (
+	"testing"
+	"time"
+
+	"github.com/prometheus/common/model"
+	"github.com/prometheus/prometheus/prompb"
+	"github.com/stretchr/testify/require"
+
+	"github.com/cortexproject/cortex/integration/e2e"
+	e2edb "github.com/cortexproject/cortex/integration/e2e/db"
+	"github.com/cortexproject/cortex/integration/e2ecortex"
+)
+
+func TestDistributorAcceptMixedHASamplesRunningInMicroservicesMode(t *testing.T) {
+	const blockRangePeriod = 5 * time.Minute
+
+	t.Run("Distriubtor accept mixed HA samples in the same request", func(t *testing.T) {
+		s, err := e2e.NewScenario(networkName)
+		require.NoError(t, err)
+		defer s.Close()
+
+		// Configure the querier to only look in ingester
+		// and enbale distributor ha tracker with mixed samples.
+		distributorFlags := map[string]string{
+			"-distributor.ha-tracker.enable":                        "true",
+			"-distributor.ha-tracker.enable-for-all-users":          "true",
+			"-experimental.distributor.ha-tracker.mixed-ha-samples": "true",
+			"-distributor.ha-tracker.cluster":                       "cluster",
+			"-distributor.ha-tracker.replica":                       "__replica__",
+			"-distributor.ha-tracker.store":                         "etcd",
+			"-distributor.ha-tracker.etcd.endpoints":                "etcd:2379",
+		}
+		querierFlags := mergeFlags(BlocksStorageFlags(), map[string]string{
+			"-querier.query-store-after": (1 * time.Hour).String(),
+		})
+		flags := mergeFlags(BlocksStorageFlags(), map[string]string{
+			"-blocks-storage.tsdb.block-ranges-period":          blockRangePeriod.String(),
+			"-blocks-storage.tsdb.ship-interval":                "5s",
+			"-blocks-storage.tsdb.retention-period":             ((blockRangePeriod * 2) - 1).String(),
+			"-blocks-storage.bucket-store.max-chunk-pool-bytes": "1",
+		})
+
+		// Start dependencies.
+		consul := e2edb.NewConsul()
+		etcd := e2edb.NewETCD()
+		minio := e2edb.NewMinio(9000, flags["-blocks-storage.s3.bucket-name"])
+		require.NoError(t, s.StartAndWaitReady(consul, etcd, minio))
+
+		// Start Cortex components.
+		distributor := e2ecortex.NewDistributor("distributor", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), distributorFlags, "")
+		ingester := e2ecortex.NewIngester("ingester", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+		require.NoError(t, s.StartAndWaitReady(distributor, ingester))
+
+		// Wait until both the distributor and ingester have updated the ring.
+		require.NoError(t, distributor.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+
+		c, err := e2ecortex.NewClient(distributor.HTTPEndpoint(), "", "", "", "user-1")
+		require.NoError(t, err)
+
+		// Push some series to Cortex.
+		series1Timestamp := time.Now()
+		series2Timestamp := series1Timestamp.Add(-2 * time.Second)
+		series3Timestamp := series1Timestamp.Add(-4 * time.Second)
+		series4Timestamp := series1Timestamp.Add(-6 * time.Second)
+		series5Timestamp := series1Timestamp.Add(-8 * time.Second)
+		series6Timestamp := series1Timestamp.Add(-10 * time.Second)
+		series7Timestamp := series1Timestamp.Add(-12 * time.Second)
+		series1, _ := generateSeries("foo", series1Timestamp, prompb.Label{Name: "__replica__", Value: "replica0"}, prompb.Label{Name: "cluster", Value: "cluster0"})
+		series2, _ := generateSeries("foo", series2Timestamp, prompb.Label{Name: "__replica__", Value: "replica1"}, prompb.Label{Name: "cluster", Value: "cluster0"})
+		series3, _ := generateSeries("foo", series3Timestamp, prompb.Label{Name: "__replica__", Value: "replica0"}, prompb.Label{Name: "cluster", Value: "cluster1"})
+		series4, _ := generateSeries("foo", series4Timestamp, prompb.Label{Name: "__replica__", Value: "replica1"}, prompb.Label{Name: "cluster", Value: "cluster1"})
+		series5, _ := generateSeries("foo", series5Timestamp, prompb.Label{Name: "__replica__", Value: "replicaNoCluster"})
+		series6, _ := generateSeries("foo", series6Timestamp, prompb.Label{Name: "cluster", Value: "clusterNoReplica"})
+		series7, _ := generateSeries("foo", series7Timestamp, prompb.Label{Name: "other", Value: "label"})
+
+		res, err := c.Push([]prompb.TimeSeries{series1[0], series2[0], series3[0], series4[0], series5[0], series6[0], series7[0]})
+		require.NoError(t, err)
+		require.Equal(t, 200, res.StatusCode)
+
+		// Wait until the samples have been deduped.
+		require.NoError(t, distributor.WaitSumMetrics(e2e.Equals(2), "cortex_distributor_deduped_samples_total"))
+		require.NoError(t, distributor.WaitSumMetrics(e2e.Equals(3), "cortex_distributor_non_ha_samples_received_total"))
+
+		// Start the querier and store-gateway, and configure them to frequently sync blocks fast enough to trigger consistency check.
+		storeGateway := e2ecortex.NewStoreGateway("store-gateway", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+		querier := e2ecortex.NewQuerier("querier", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(querierFlags, flags), "")
+		require.NoError(t, s.StartAndWaitReady(querier, storeGateway))
+
+		// Wait until the querier and store-gateway have updated the ring, and wait until the blocks are old enough for consistency check
+		require.NoError(t, querier.WaitSumMetrics(e2e.Equals(512*2), "cortex_ring_tokens_total"))
+		require.NoError(t, storeGateway.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+
+		// Query back the series.
+		c, err = e2ecortex.NewClient("", querier.HTTPEndpoint(), "", "", "user-1")
+		require.NoError(t, err)
+
+		// Query back the series (only in the ingesters).
+		result, err := c.Query("foo[5m]", series1Timestamp)
+		require.NoError(t, err)
+
+		require.Equal(t, model.ValMatrix, result.Type())
+		m := result.(model.Matrix)
+		require.Equal(t, 5, m.Len())
+		numValidHA := 0
+		numNonHA := 0
+		for _, ss := range m {
+			replicaLabel, okReplica := ss.Metric["__replica__"]
+			if okReplica {
+				require.Equal(t, string(replicaLabel), "replicaNoCluster")
+			}
+			clusterLabel, okCluster := ss.Metric["cluster"]
+			if okCluster {
+				require.Equal(t, string(clusterLabel) == "cluster1" || string(clusterLabel) == "cluster0" || string(clusterLabel) == "clusterNoReplica", true)
+				if clusterLabel == "cluster1" || clusterLabel == "cluster0" {
+					numValidHA++
+				}
+			}
+			if (okReplica && !okCluster && replicaLabel == "replicaNoCluster") || (okCluster && !okReplica && clusterLabel == "clusterNoReplica") || (!okCluster && !okReplica) {
+				numNonHA++
+			}
+			require.NotEmpty(t, ss.Values)
+			for _, v := range ss.Values {
+				require.NotEmpty(t, v)
+			}
+		}
+		require.Equal(t, numNonHA, 3)
+		require.Equal(t, numValidHA, 2)
+
+		// Ensure no service-specific metrics prefix is used by the wrong service.
+		assertServiceMetricsPrefixes(t, Distributor, distributor)
+		assertServiceMetricsPrefixes(t, Ingester, ingester)
+		assertServiceMetricsPrefixes(t, StoreGateway, storeGateway)
+		assertServiceMetricsPrefixes(t, Querier, querier)
+	})
+}

pkg/distributor/distributor.go

@@ -175,6 +175,11 @@ type InstanceLimits struct {
 	MaxInflightPushRequests int     `yaml:"max_inflight_push_requests"`
 }
 
+type HAPair struct {
+	Cluster string
+	Replica string
+}
+
 type OTLPConfig struct {
 	ConvertAllAttributes bool `yaml:"convert_all_attributes"`
 	DisableTargetInfo    bool `yaml:"disable_target_info"`
@@ -665,7 +670,7 @@ func (d *Distributor) Push(ctx context.Context, req *cortexpb.WriteRequest) (*co
 	// Cache user limit with overrides so we spend less CPU doing locking. See issue #4904
 	limits := d.limits.GetOverridesForUser(userID)
 
-	if limits.AcceptHASamples && len(req.Timeseries) > 0 {
+	if limits.AcceptHASamples && len(req.Timeseries) > 0 && !limits.AcceptMixedHASamples {
 		cluster, replica := findHALabels(limits.HAReplicaLabel, limits.HAClusterLabel, req.Timeseries[0].Labels)
 		removeReplica, err = d.checkSample(ctx, userID, cluster, replica, limits)
 		if err != nil {
@@ -868,10 +873,39 @@ func (d *Distributor) prepareSeriesKeys(ctx context.Context, req *cortexpb.Write
 		}
 	}()
 
+	var haPairs map[HAPair]error
+	if limits.AcceptHASamples && limits.AcceptMixedHASamples && len(req.Timeseries) > 0 {
+		haPairs = d.computeHAPairs(req.Timeseries, ctx, userID, limits)
+	}
+
 	// For each timeseries, compute a hash to distribute across ingesters;
 	// check each sample and discard if outside limits.
 	skipLabelNameValidation := d.cfg.SkipLabelNameValidation || req.GetSkipLabelNameValidation()
 	for _, ts := range req.Timeseries {
+		if len(haPairs) > 0 {
+			cluster, replica := findHALabels(limits.HAReplicaLabel, limits.HAClusterLabel, ts.Labels)
+
+			if cluster != "" && replica != "" {
+				err := haPairs[HAPair{Cluster: cluster, Replica: replica}]
+				if err != nil {
+					// discard sample (non valid HA sample)
+					if errors.Is(err, ha.ReplicasNotMatchError{}) {
+						// These samples have been deduped.
+						d.dedupedSamples.WithLabelValues(userID, cluster).Add(float64(len(ts.Samples) + len(ts.Histograms)))
+					}
+					if errors.Is(err, ha.TooManyReplicaGroupsError{}) {
+						d.validateMetrics.DiscardedSamples.WithLabelValues(validation.TooManyHAClusters, userID).Add(float64(len(ts.Samples) + len(ts.Histograms)))
+					}
+
+					continue
+				}
+				removeReplica = true // valid HA sample
+			} else {
+				// non HA sample
+				removeReplica = false
+				d.nonHASamples.WithLabelValues(userID).Add(float64(len(ts.Samples) + len(ts.Histograms)))
+			}
+		}
 		// Use timestamp of latest sample in the series. If samples for series are not ordered, metric for user may be wrong.
 		if len(ts.Samples) > 0 {
 			latestSampleTimestampMs = max(latestSampleTimestampMs, ts.Samples[len(ts.Samples)-1].TimestampMs)
@@ -1480,6 +1514,30 @@ func findHALabels(replicaLabel, clusterLabel string, labels []cortexpb.LabelAdap
 	return cluster, replica
 }
 
+func (d *Distributor) computeHAPairs(timeseries []cortexpb.PreallocTimeseries, ctx context.Context, userID string, limits *validation.Limits) map[HAPair]error {
+	var haPairs = make(map[HAPair]error)
+	clusterLabel := limits.HAClusterLabel
+	replicaLabel := limits.HAReplicaLabel
+
+	for _, ts := range timeseries {
+		cluster, replica := findHALabels(replicaLabel, clusterLabel, ts.Labels)
+		if cluster != "" && replica != "" {
+			haPair := HAPair{Cluster: cluster, Replica: replica}
+			_, ok := haPairs[haPair]
+			if !ok {
+				removeReplica, err := d.checkSample(ctx, userID, cluster, replica, limits)
+				if err == nil && !removeReplica {
+					continue
+				}
+
+				haPairs[haPair] = err
+			}
+		}
+	}
+
+	return haPairs
+}
+
 func getLimitFromLabelHints(hints *storage.LabelHints) int {
 	if hints != nil {
 		return hints.Limit