Prevent failed ingestion from affecting rate limiting in distributor. (cortexproject#3825)

* Prevent failed ingestion from affecting rate limiting in distributor.

Change the handling of the rate limiting in the distributor such that when the ingestion fails, the reservation in the rate limiter is canceled, returning the tokens for future use in the proceeding requests.

Signed-off-by: Steve Simpson <[email protected]>

* Spelling and typos.

Signed-off-by: Steve Simpson <[email protected]>

* Change Reservation to an interface to avoid extra allocation.

Signed-off-by: Steve Simpson <[email protected]>

* Return nil on failure in AllowN.

Signed-off-by: Steve Simpson <[email protected]>

* Rename unit test helper OK to isOK.

Signed-off-by: Steve Simpson <[email protected]>

* Review suggestions.

Signed-off-by: Steve Simpson <[email protected]>

Author: stevesg

CHANGELOG.md

@@ -52,6 +52,7 @@
   * `cortex_ha_tracker_replicas_cleanup_deleted_total`
   * `cortex_ha_tracker_replicas_cleanup_delete_failed_total`
 * [ENHANCEMENT] Tenant deletion endpoints now support deletion of ruler groups. This only works when using rule store that supports deletion. #3750
+* [ENHANCEMENT] Distributor: Prevent failed ingestion from affecting rate limiting. #3825
 * [BUGFIX] Cortex: Fixed issue where fatal errors and various log messages where not logged. #3778
 * [BUGFIX] HA Tracker: don't track as error in the `cortex_kv_request_duration_seconds` metric a CAS operation intentionally aborted. #3745
 * [BUGFIX] Querier / ruler: do not log "error removing stale clients" if the ring is empty. #3761

pkg/distributor/distributor.go

@@ -580,7 +580,8 @@ func (d *Distributor) Push(ctx context.Context, req *ingester_client.WriteReques
 	}
 
 	totalN := validatedSamples + len(validatedMetadata)
-	if !d.ingestionRateLimiter.AllowN(now, userID, totalN) {
+	rateOK, rateReservation := d.ingestionRateLimiter.AllowN(now, userID, totalN)
+	if !rateOK {
 		// Ensure the request slice is reused if the request is rate limited.
 		ingester_client.ReuseSlice(req.Timeseries)
 
@@ -632,6 +633,8 @@ func (d *Distributor) Push(ctx context.Context, req *ingester_client.WriteReques
 		return d.send(localCtx, ingester, timeseries, metadata, req.Source)
 	}, func() { ingester_client.ReuseSlice(req.Timeseries) })
 	if err != nil {
+		// Ingestion failed, so roll-back the reservation from the rate limiter.
+		rateReservation.CancelAt(now)
 		return nil, err
 	}
 	return &ingester_client.WriteResponse{}, firstPartialErr

pkg/distributor/distributor_test.go

@@ -369,6 +369,7 @@ func TestDistributor_PushIngestionRateLimiter(t *testing.T) {
 		ingestionRateStrategy string
 		ingestionRate         float64
 		ingestionBurstSize    int
+		ingestionFailing      bool
 		pushes                []testPush
 	}{
 		"local strategy: limit should be set to each distributor": {
@@ -413,6 +414,17 @@ func TestDistributor_PushIngestionRateLimiter(t *testing.T) {
 				{metadata: 1, expectedError: httpgrpc.Errorf(http.StatusTooManyRequests, "ingestion rate limit (5) exceeded while adding 0 samples and 1 metadata")},
 			},
 		},
+		"unhappy ingesters: rate limit should be unaffected when ingestion fails": {
+			distributors:          1,
+			ingestionRateStrategy: validation.LocalIngestionRateStrategy,
+			ingestionRate:         10,
+			ingestionBurstSize:    10,
+			ingestionFailing:      true,
+			pushes: []testPush{
+				{samples: 10, expectedError: errFail},
+				{samples: 10, expectedError: errFail},
+			},
+		},
 	}
 
 	for testName, testData := range tests {
@@ -425,10 +437,15 @@ func TestDistributor_PushIngestionRateLimiter(t *testing.T) {
 			limits.IngestionRate = testData.ingestionRate
 			limits.IngestionBurstSize = testData.ingestionBurstSize
 
+			happyIngesters := 3
+			if testData.ingestionFailing {
+				happyIngesters = 0
+			}
+
 			// Start all expected distributors
 			distributors, _, r := prepare(t, prepConfig{
 				numIngesters:     3,
-				happyIngesters:   3,
+				happyIngesters:   happyIngesters,
 				numDistributors:  testData.distributors,
 				shardByAllLabels: true,
 				limits:           limits,

pkg/util/limiter/rate_limiter.go

@@ -26,6 +26,17 @@ type RateLimiter struct {
 	tenants     map[string]*tenantLimiter
 }
 
+// Reservation is similar to rate.Reservation but excludes interfaces which do
+// not make sense to expose, because we are following the semantics of AllowN,
+// being an immediate reservation, i.e. not delayed into the future.
+type Reservation interface {
+	// CancelAt returns the reservation to the rate limiter for use by other
+	// requests. Note that typically the reservation should be canceled with
+	// the same timestamp it was requested with, or not all the tokens
+	// consumed will be returned.
+	CancelAt(now time.Time)
+}
+
 type tenantLimiter struct {
 	limiter   *rate.Limiter
 	recheckAt time.Time
@@ -42,9 +53,29 @@ func NewRateLimiter(strategy RateLimiterStrategy, recheckPeriod time.Duration) *
 	}
 }
 
-// AllowN reports whether n tokens may be consumed happen at time now.
-func (l *RateLimiter) AllowN(now time.Time, tenantID string, n int) bool {
-	return l.getTenantLimiter(now, tenantID).AllowN(now, n)
+// AllowN reports whether n tokens may be consumed happen at time now. The
+// reservation of tokens can be canceled using CancelAt on the returned object.
+func (l *RateLimiter) AllowN(now time.Time, tenantID string, n int) (bool, Reservation) {
+
+	// Using ReserveN allows cancellation of the reservation, but
+	// the semantics are subtly different to AllowN.
+	r := l.getTenantLimiter(now, tenantID).ReserveN(now, n)
+	if !r.OK() {
+		return false, nil
+	}
+
+	// ReserveN will still return OK if the necessary tokens are
+	// available in the future, and tells us this time delay. In
+	// order to mimic the semantics of AllowN, we must check that
+	// there is no delay before we can use them.
+	if r.DelayFrom(now) > 0 {
+		// Having decided not to use the reservation, return the
+		// tokens to the rate limiter.
+		r.CancelAt(now)
+		return false, nil
+	}
+
+	return true, r
 }
 
 // Limit returns the currently configured maximum overall tokens rate.

pkg/util/limiter/rate_limiter_test.go

@@ -45,24 +45,60 @@ func TestRateLimiter_AllowN(t *testing.T) {
 	now := time.Now()
 
 	// Tenant #1
-	assert.Equal(t, true, limiter.AllowN(now, "tenant-1", 8))
-	assert.Equal(t, true, limiter.AllowN(now, "tenant-1", 10))
-	assert.Equal(t, false, limiter.AllowN(now, "tenant-1", 3))
-	assert.Equal(t, true, limiter.AllowN(now, "tenant-1", 2))
+	assert.Equal(t, true, isOK(limiter.AllowN(now, "tenant-1", 8)))
+	assert.Equal(t, true, isOK(limiter.AllowN(now, "tenant-1", 10)))
+	assert.Equal(t, false, isOK(limiter.AllowN(now, "tenant-1", 3)))
+	assert.Equal(t, true, isOK(limiter.AllowN(now, "tenant-1", 2)))
 
-	assert.Equal(t, true, limiter.AllowN(now.Add(time.Second), "tenant-1", 8))
-	assert.Equal(t, false, limiter.AllowN(now.Add(time.Second), "tenant-1", 3))
-	assert.Equal(t, true, limiter.AllowN(now.Add(time.Second), "tenant-1", 2))
+	assert.Equal(t, true, isOK(limiter.AllowN(now.Add(time.Second), "tenant-1", 8)))
+	assert.Equal(t, false, isOK(limiter.AllowN(now.Add(time.Second), "tenant-1", 3)))
+	assert.Equal(t, true, isOK(limiter.AllowN(now.Add(time.Second), "tenant-1", 2)))
 
 	// Tenant #2
-	assert.Equal(t, true, limiter.AllowN(now, "tenant-2", 18))
-	assert.Equal(t, true, limiter.AllowN(now, "tenant-2", 20))
-	assert.Equal(t, false, limiter.AllowN(now, "tenant-2", 3))
-	assert.Equal(t, true, limiter.AllowN(now, "tenant-2", 2))
-
-	assert.Equal(t, true, limiter.AllowN(now.Add(time.Second), "tenant-2", 18))
-	assert.Equal(t, false, limiter.AllowN(now.Add(time.Second), "tenant-2", 3))
-	assert.Equal(t, true, limiter.AllowN(now.Add(time.Second), "tenant-2", 2))
+	assert.Equal(t, true, isOK(limiter.AllowN(now, "tenant-2", 18)))
+	assert.Equal(t, true, isOK(limiter.AllowN(now, "tenant-2", 20)))
+	assert.Equal(t, false, isOK(limiter.AllowN(now, "tenant-2", 3)))
+	assert.Equal(t, true, isOK(limiter.AllowN(now, "tenant-2", 2)))
+
+	assert.Equal(t, true, isOK(limiter.AllowN(now.Add(time.Second), "tenant-2", 18)))
+	assert.Equal(t, false, isOK(limiter.AllowN(now.Add(time.Second), "tenant-2", 3)))
+	assert.Equal(t, true, isOK(limiter.AllowN(now.Add(time.Second), "tenant-2", 2)))
+}
+
+func TestRateLimiter_AllowNCancelation(t *testing.T) {
+	strategy := &staticLimitStrategy{tenants: map[string]struct {
+		limit float64
+		burst int
+	}{
+		"tenant-1": {limit: 10, burst: 20},
+	}}
+
+	limiter := NewRateLimiter(strategy, 10*time.Second)
+	now := time.Now()
+
+	assert.Equal(t, true, isOK(limiter.AllowN(now, "tenant-1", 12)))
+	assert.Equal(t, false, isOK(limiter.AllowN(now, "tenant-1", 9)))
+
+	ok1, r1 := limiter.AllowN(now, "tenant-1", 8)
+	assert.Equal(t, true, ok1)
+	r1.CancelAt(now)
+
+	assert.Equal(t, true, isOK(limiter.AllowN(now, "tenant-1", 8)))
+
+	// +10 tokens (1s)
+	nowPlus := now.Add(time.Second)
+
+	assert.Equal(t, true, isOK(limiter.AllowN(nowPlus, "tenant-1", 6)))
+	assert.Equal(t, false, isOK(limiter.AllowN(nowPlus, "tenant-1", 5)))
+
+	ok2, r2 := limiter.AllowN(nowPlus, "tenant-1", 4)
+	assert.Equal(t, true, ok2)
+	r2.CancelAt(nowPlus)
+
+	assert.Equal(t, true, isOK(limiter.AllowN(nowPlus, "tenant-1", 2)))
+	assert.Equal(t, false, isOK(limiter.AllowN(nowPlus, "tenant-1", 3)))
+	assert.Equal(t, true, isOK(limiter.AllowN(nowPlus, "tenant-1", 2)))
+	assert.Equal(t, false, isOK(limiter.AllowN(nowPlus, "tenant-1", 1)))
 }
 
 func BenchmarkRateLimiter_CustomMultiTenant(b *testing.B) {
@@ -127,3 +163,7 @@ func (s *staticLimitStrategy) Burst(tenantID string) int {
 
 	return tenant.burst
 }
+
+func isOK(ok bool, r Reservation) bool {
+	return ok
+}