Rebase and complete per-sample ha deduplication

Signed-off-by: juliette.orain <[email protected]>
Co-authored-by: Nicolas DUPEUX <[email protected]>

Author: julietteO

CHANGELOG.md

@@ -67,6 +67,7 @@
 * [ENHANCEMENT] Usage tracker: `loadSnapshot()` checks shard emptiness instead of using explicit `first` parameter. #13534
 * [ENHANCEMENT] OTLP: Add metric `cortex_distributor_otlp_requests_by_content_type_total` to track content type (json or proto) of OTLP packets. #13525
 * [ENHANCEMENT] OTLP: Add experimental metric `cortex_distributor_otlp_array_lengths` to better understand the layout of OTLP packets in practice. #13525
+* [ENHANCEMENT] HA: Deduplication per sample instead of per batch. #13665
 * [BUGFIX] Compactor: Fix potential concurrent map writes. #13053
 * [BUGFIX] Query-frontend: Fix issue where queries sometimes fail with `failed to receive query result stream message: rpc error: code = Canceled desc = context canceled` if remote execution is enabled. #13084
 * [BUGFIX] Query-frontend: Fix issue where query stats, such as series read, did not include the parameters to the `histogram_quantile` and `histogram_fraction` functions if remote execution was enabled. #13084

pkg/distributor/distributor.go

@@ -1151,24 +1151,35 @@ type haReplica struct {
 	cluster, replica string
 }
 
-// TODO document what this does
-func (d *Distributor) replicaObserved(ctx context.Context, userID string, replica haReplica) (replicaState, error) {
+type replicaInfo struct {
+	state       replicaState
+	sampleCount int
+}
+
+// replicaObserved checks if a sample from a given replica should be accepted for ingestion based on HA deduplication rules.
+//
+// Returns a replicaState indicating the acceptance status and classification of the replica:
+//   - replicaIsPrimary: sample is from the elected primary replica and should be accepted
+//   - replicaNotHA: sample doesn't have both HA labels and should be accepted
+//   - replicaDeduped: sample is from a non-primary replica and should be deduplicated
+//   - replicaRejectedTooManyClusters: sample is rejected because the tenant has too many HA clusters
+//   - replicaRejectedUnknown: sample is rejected due to an unknown error
+func (d *Distributor) replicaObserved(ctx context.Context, userID string, replica haReplica, ts int64) (replicaState, error) {
 	if span := opentracing.SpanFromContext(ctx); span != nil {
-		// TODO move this if in a separate function
 		// Make a copy of these, since they may be retained as tags
 		span.LogFields(
-			opentracing_log.String("cluster", copyString(replica.cluster)),
-			opentracing_log.String("replica", copyString(replica.replica)),
+			opentracing_log.String("cluster", strings.Clone(replica.cluster)),
+			opentracing_log.String("replica", strings.Clone(replica.replica)),
 		)
 	}
 
-	isAccepted, err := d.checkSample(ctx, userID, replica.cluster, replica.replica)
+	isAccepted, err := d.checkSample(ctx, userID, replica.cluster, replica.replica, ts)
 	if err != nil {
 		switch {
-		case errors.Is(err, replicasNotMatchError{}):
+		case errors.As(err, &replicasDidNotMatchError{}):
 			// These samples have been deduped.
 			return replicaDeduped, err
-		case errors.Is(err, tooManyClustersError{}):
+		case errors.As(err, &tooManyClustersError{}):
 			return replicaRejectedTooManyClusters, err
 		default:
 			return replicaRejectedUnknown, err
@@ -1182,7 +1193,6 @@ func (d *Distributor) replicaObserved(ctx context.Context, userID string, replic
 	return replicaNotHA, nil
 }
 
-
 func (d *Distributor) prePushHaDedupeMiddleware(next PushFunc) PushFunc {
 	return WithCleanup(next, func(next PushFunc, ctx context.Context, pushReq *Request) error {
 		req, err := pushReq.WriteRequest()
@@ -1200,7 +1210,8 @@ func (d *Distributor) prePushHaDedupeMiddleware(next PushFunc) PushFunc {
 		}
 
 		haReplicaLabel := d.limits.HAReplicaLabel(userID)
-		cluster, replica := findHALabels(haReplicaLabel, d.limits.HAClusterLabel(userID), req.Timeseries[0].Labels)
+		haClusterLabel := d.limits.HAClusterLabel(userID)
+		cluster, replica := findHALabels(haReplicaLabel, haClusterLabel, req.Timeseries[0].Labels)
 		// Make a copy of these, since they may be retained as labels on our metrics, e.g. dedupedSamples.
 		cluster, replica = strings.Clone(cluster), strings.Clone(replica)
 
@@ -1213,12 +1224,11 @@ func (d *Distributor) prePushHaDedupeMiddleware(next PushFunc) PushFunc {
 		numSamples := 0
 		now := time.Now()
 
-		var getReplicaForSample func(int) haReplica
-		{
-			var replicaKey haReplica // TODO dimitarvdimitrov - do we need this closure or a regular func also works?
-			getReplicaForSample = func(i int) haReplica {
-				replicaKey.cluster, replicaKey.replica = findHALabels(haReplicaLabel, haClusterLabel, req.Timeseries[i].Labels)
-				return replicaKey
+		getReplicaForSample := func(i int) haReplica {
+			cluster, replica := findHALabels(haReplicaLabel, haClusterLabel, req.Timeseries[i].Labels)
+			return haReplica{
+				cluster: cluster,
+				replica: replica,
 			}
 		}
 
@@ -1247,61 +1257,57 @@ func (d *Distributor) prePushHaDedupeMiddleware(next PushFunc) PushFunc {
 		}
 
 		var errs multierror.MultiError
-		// TODO dimitarvdimitrov maybe unite the two maps that have the haReplicas as the key
-		samplesPerReplica := make(map[haReplica]int)
-		replicaStates := make(map[haReplica]replicaState, len(samplesPerReplica))
+		replicaInfos := make(map[haReplica]*replicaInfo)
 		samplesPerState := make(map[replicaState]int)
+		// Check if all timeseries belong to the same replica
+		firstReplica := getReplicaForSample(0)
 		isOneReplica := true
-		{
-			replica := getReplicaForSample(0)
-
-			for i := range req.Timeseries {
-				if getReplicaForSample(i) != replica {
-					isOneReplica = false
-					break
-				}
-			}
-		}
-		{
-			if isOneReplica {
-				samplesPerReplica[getReplicaForSample(0)] = numSamples
-			} else {
-				for i := range req.Timeseries {
-					samplesPerReplica[getReplicaForSample(i)]++
-				}
+		for i := 1; i < len(req.Timeseries); i++ {
+			if getReplicaForSample(i) != firstReplica {
+				isOneReplica = false
+				break
 			}
 		}
+
+		// Count samples per replica
 		if isOneReplica {
-			var r haReplica
-			for k := range samplesPerReplica {
-				r = k
-			}
+			replicaInfos[firstReplica] = &replicaInfo{sampleCount: numSamples}
+			// Optimize getReplicaForSample to return the single replica directly
 			getReplicaForSample = func(i int) haReplica {
-				return r
+				return firstReplica
+			}
+		} else {
+			for i := range req.Timeseries {
+				r := getReplicaForSample(i)
+				info := replicaInfos[r]
+				if info == nil {
+					info = &replicaInfo{}
+					replicaInfos[r] = info
+				}
+				info.sampleCount++
 			}
 		}
-		for replicaKey, numSamples := range samplesPerReplica {
-			state, ok := replicaStates[replicaKey]
-			if !ok {
-				state, err = d.replicaObserved(ctx, userID, replicaKey)
-				replicaStates[replicaKey] = state
+		for replicaKey, info := range replicaInfos {
+			if info.state == replicaRejectedUnknown {
+				state, err := d.replicaObserved(ctx, userID, replicaKey, sampleTimestamp)
+				info.state = state
 				errs.Add(err)
 			}
-			samplesPerState[state] += numSamples
+			samplesPerState[info.state] += info.sampleCount
 		}
-		lastAccepted := sortByAccepted(req, replicaStates, getReplicaForSample)
+		lastAccepted := sortByAccepted(req, replicaInfos, getReplicaForSample)
 		var getReplicaState func(haReplica) replicaState
 		if isOneReplica {
 			var s replicaState
-			for _, rs := range replicaStates {
-				s = rs
+			for _, info := range replicaInfos {
+				s = info.state
 			}
 			getReplicaState = func(replica haReplica) replicaState {
 				return s
 			}
 		} else {
 			getReplicaState = func(replica haReplica) replicaState {
-				return replicaStates[replica]
+				return replicaInfos[replica].state
 			}
 		}
 		for i := 0; i <= lastAccepted; i++ {
@@ -1317,35 +1323,15 @@ func (d *Distributor) prePushHaDedupeMiddleware(next PushFunc) PushFunc {
 		}
 		// We don't want to send samples beyond the last accepted sample - that was deduplicated
 		pushReq.AddCleanup(sliceUnacceptedRequests(req, lastAccepted))
-		{
-			// TODO move into a function
-			for replica, state := range replicaStates {
-				if state&replicaDeduped != 0 && samplesPerReplica[replica] > 0 {
-					cluster := copyString(replica.cluster) // Make a copy of this, since it may be retained as labels on our metrics
-					d.dedupedSamples.WithLabelValues(userID, cluster).Add(float64(samplesPerReplica[replica]))
-				}
-			}
-			if samplesPerState[replicaNotHA] > 0 {
-				d.nonHASamples.WithLabelValues(userID).Add(float64(samplesPerState[replicaNotHA]))
-			}
-			if samplesPerState[replicaRejectedTooManyClusters] > 0 {
-				d.discardedSamplesTooManyHaClusters.WithLabelValues(userID, group).Add(float64(samplesPerState[replicaRejectedTooManyClusters]))
-			}
-		}
-		var resp *mimirpb.WriteResponse
+		d.updateHADedupeMetrics(userID, group, replicaInfos, samplesPerState)
+
 		if len(req.Timeseries) > 0 {
-			cleanupInDefer = false
-			resp, err = next(ctx, pushReq)
+			err = next(ctx, pushReq)
 		}
 		errs.Add(err)
 
-		if samplesPerState[replicaDeduped] > 0 {
-			return nil, httpgrpc.Errorf(http.StatusAccepted, errs.Err().Error())
-		} else if len(errs) > 0 {
-			return nil, httpgrpc.Errorf(http.StatusBadRequest, errs.Err().Error()) // ermmm?
-		}
-		return resp, nil
-	}
+		return errs.Err()
+	})
 }
 
 func sliceUnacceptedRequests(req *mimirpb.WriteRequest, lastAccepted int) func() {
@@ -1358,20 +1344,36 @@ func sliceUnacceptedRequests(req *mimirpb.WriteRequest, lastAccepted int) func()
 
 }
 
+// updateHADedupeMetrics updates metrics related to HA deduplication.
+func (d *Distributor) updateHADedupeMetrics(userID, group string, replicaInfos map[haReplica]*replicaInfo, samplesPerState map[replicaState]int) {
+	for replica, info := range replicaInfos {
+		if info.state&replicaDeduped != 0 && info.sampleCount > 0 {
+			cluster := strings.Clone(replica.cluster) // Make a copy of this, since it may be retained as labels on our metrics
+			d.dedupedSamples.WithLabelValues(userID, cluster).Add(float64(info.sampleCount))
+		}
+	}
+	if samplesPerState[replicaNotHA] > 0 {
+		d.nonHASamples.WithLabelValues(userID).Add(float64(samplesPerState[replicaNotHA]))
+	}
+	if samplesPerState[replicaRejectedTooManyClusters] > 0 {
+		d.discardedSamplesTooManyHaClusters.WithLabelValues(userID, group).Add(float64(samplesPerState[replicaRejectedTooManyClusters]))
+	}
+}
+
 // sortByAccepted returns the index of the last acepted timeseries in the write request based on the ha dedup sattes of the eplicas
-func sortByAccepted(req *mimirpb.WriteRequest, replicaStates map[haReplica]replicaState, getReplicaForSample func(int) haReplica) int {
+func sortByAccepted(req *mimirpb.WriteRequest, replicaInfos map[haReplica]*replicaInfo, getReplicaForSample func(int) haReplica) int {
 	numAcceptedReplicas := 0
-	for _, state := range replicaStates {
-		if state&replicaAccepted != 0 {
+	for _, info := range replicaInfos {
+		if info.state&replicaAccepted != 0 {
 			numAcceptedReplicas++
 		}
 	}
-	if numAcceptedReplicas == len(replicaStates) {
+	if numAcceptedReplicas == len(replicaInfos) {
 		return len(req.Timeseries) - 1
 	}
 	findPreviousAccepted := func(i int) int {
 		for i > 0 {
-			state := replicaStates[getReplicaForSample(i)]
+			state := replicaInfos[getReplicaForSample(i)].state
 			if state&replicaAccepted != 0 {
 				break
 			}
@@ -1385,7 +1387,7 @@ func sortByAccepted(req *mimirpb.WriteRequest, replicaStates map[haReplica]repli
 		if i > lastAccepted {
 			break
 		}
-		state := replicaStates[getReplicaForSample(i)]
+		state := replicaInfos[getReplicaForSample(i)].state
 		if state&replicaAccepted == 0 {
 			req.Timeseries[i], req.Timeseries[lastAccepted] = req.Timeseries[lastAccepted], req.Timeseries[i]
 			lastAccepted--

pkg/distributor/distributor_test.go

@@ -959,15 +959,15 @@ func TestDistributor_PushHAInstances(t *testing.T) {
 			samples:          5,
 			expectedResponse: emptyResponse,
 		},
-		// Using very long replica label value results in not pushing the sample, hence a success.
+		// Using very long replica label value results in validation error.
 		{
-			enableTracker:    true,
-			acceptedReplica:  "instance0",
-			testReplica:      "instance1234567890123456789012345678901234567890",
-			cluster:          "cluster0",
-			samples:          5,
-			expectedResponse: emptyResponse,
-			expectedCode:     0,
+			enableTracker:   true,
+			acceptedReplica: "instance0",
+			testReplica:     "instance1234567890123456789012345678901234567890",
+			cluster:         "cluster0",
+			samples:         5,
+			expectedError:   status.New(codes.InvalidArgument, fmt.Sprintf(labelValueTooLongMsgFormat, 48, 15, "__replica__", "instance1234567890123456789012345678901234567890", mimirpb.FromLabelAdaptersToString(labelSetGenWithReplicaAndCluster("instance1234567890123456789012345678901234567890", "cluster0")(0)))),
+			expectedDetails: &mimirpb.ErrorDetails{Cause: mimirpb.ERROR_CAUSE_BAD_DATA},
 		},
 	} {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
@@ -2379,7 +2379,7 @@ func BenchmarkDistributor_Push(b *testing.B) {
 				samples := make([]mimirpb.Sample, numSeriesPerRequest)
 
 				for i := 0; i < numSeriesPerRequest; i++ {
-					lbls := labels.NewBuilder(labels.Labels{{Name: model.MetricNameLabel, Value: "foo"}})
+					lbls := labels.NewBuilder(labels.FromStrings(model.MetricNameLabel, "foo"))
 					for i := 0; i < 10; i++ {
 						lbls.Set(fmt.Sprintf("name_%d", i), fmt.Sprintf("value_%d", i))
 					}
@@ -2407,7 +2407,7 @@ func BenchmarkDistributor_Push(b *testing.B) {
 				samples := make([]mimirpb.Sample, numSeriesPerRequest)
 
 				for i := 0; i < numSeriesPerRequest; i++ {
-					lbls := labels.NewBuilder(labels.Labels{{Name: model.MetricNameLabel, Value: "foo"}})
+					lbls := labels.NewBuilder(labels.FromStrings(model.MetricNameLabel, "foo"))
 					for i := 0; i < 10; i++ {
 						lbls.Set(fmt.Sprintf("name_%d", i), fmt.Sprintf("value_%d", i))
 					}
@@ -5192,8 +5192,8 @@ func TestHaDedupeMiddleware(t *testing.T) {
 			},
 			expectedReqs:      []*mimirpb.WriteRequest{makeWriteRequestForGenerators(5, labelSetGenWithCluster(cluster1), nil, nil)},
 			expectedNextCalls: 1,
-			expectErrs:        []*status.Status{nil, status.New(codes.AlreadyExists, newReplicasDidNotMatchError(replica2, replica1).Error())},
-			expectDetails:     []*mimirpb.ErrorDetails{nil, replicasDidNotMatchDetails},
+			expectErrs: []*status.Status{nil, status.New(codes.AlreadyExists, newReplicasDidNotMatchError(replica2, replica1).Error())},
+			expectDetails: []*mimirpb.ErrorDetails{nil, replicasDidNotMatchDetails},
 		}, {
 			name:            "exceed max ha clusters limit",
 			ctx:             ctxWithUser,
@@ -5236,7 +5236,8 @@ func TestHaDedupeMiddleware(t *testing.T) {
 				return []*mimirpb.WriteRequest{c1, c2}
 			}(),
 			expectedNextCalls: 2,
-			expectErrs:        []int{0, 202},
+			expectErrs:        []*status.Status{nil, status.New(codes.AlreadyExists, newReplicasDidNotMatchError(replica2, replica1).Error())},
+			expectDetails:     []*mimirpb.ErrorDetails{nil, replicasDidNotMatchDetails},
 		},
 	}
 
@@ -5296,22 +5297,16 @@ func TestHaDedupeMiddleware(t *testing.T) {
 			for i := range gotReqs {
 				assert.ElementsMatch(t, tc.expectedReqs[i].Timeseries, gotReqs[i].Timeseries)
 				assert.ElementsMatch(t, tc.expectedReqs[i].Metadata, gotReqs[i].Metadata)
-				assert.Equal(t, tc.expectedReqs[i].SkipLabelNameValidation, gotReqs[i].SkipLabelNameValidation)
 				assert.Equal(t, tc.expectedReqs[i].Source, gotReqs[i].Source)
 			}
 			assert.Len(t, gotErrs, len(tc.expectErrs))
 			for errIdx, expectErr := range tc.expectErrs {
-				if expectErr > 0 {
-					// Expect an httpgrpc error with specific status code.
-					resp, ok := httpgrpc.HTTPResponseFromError(gotErrs[errIdx])
-					if assert.True(t, ok) {
-						assert.Equal(t, expectErr, int(resp.Code))
-					}
-				} else if expectErr == 0 {
+				if expectErr != nil {
+					// Expect a gRPC error.
+					checkGRPCError(t, expectErr, tc.expectDetails[errIdx], gotErrs[errIdx])
+				} else {
 					// Expect no error.
 					assert.Nil(t, gotErrs[errIdx])
-				} else {
-					checkGRPCError(t, expectErr, tc.expectDetails[errIdx], gotErrs[errIdx])
 				}
 			}