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 <steve.simpson@grafana.com>

Author: stevesg

CHANGELOG.md

@@ -50,6 +50,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) {
+	rateReservation := d.ingestionRateLimiter.AllowN(now, userID, totalN)
+	if !rateReservation.OK() {
 		// 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 ingestors: 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,14 @@ type RateLimiter struct {
 	tenants     map[string]*tenantLimiter
 }
 
+// Reservation encapsulates rate.Reservation to exclude 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 struct {
+	ok          bool
+	reservation *rate.Reservation
+}
+
 type tenantLimiter struct {
 	limiter   *rate.Limiter
 	recheckAt time.Time
@@ -42,9 +50,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) *Reservation {
+
+	// Using ReserveN allows cancalation of the reservation, but
+	// the semantics are subtly different to AllowN.
+	r := l.getTenantLimiter(now, tenantID).ReserveN(now, n)
+	if !r.OK() {
+		return &Reservation{false, r}
+	}
+
+	// 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) > time.Duration(0) {
+		// Having decided not to use the reservation, return the
+		// tokens to the rate limiter.
+		r.CancelAt(now)
+		return &Reservation{false, r}
+	}
+
+	return &Reservation{true, r}
 }
 
 // Limit returns the currently configured maximum overall tokens rate.
@@ -120,3 +148,15 @@ func (l *RateLimiter) recheckTenantLimiter(now time.Time, tenantID string) *rate
 
 	return entry.limiter
 }
+
+// OK indicates whether the request to rate limiter was allowed.
+func (a *Reservation) OK() bool {
+	return a.ok
+}
+
+// Cancel 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.
+func (a *Reservation) CancelAt(now time.Time) {
+	a.reservation.CancelAt(now)
+}

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, limiter.AllowN(now, "tenant-1", 8).OK())
+	assert.Equal(t, true, limiter.AllowN(now, "tenant-1", 10).OK())
+	assert.Equal(t, false, limiter.AllowN(now, "tenant-1", 3).OK())
+	assert.Equal(t, true, limiter.AllowN(now, "tenant-1", 2).OK())
 
-	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, limiter.AllowN(now.Add(time.Second), "tenant-1", 8).OK())
+	assert.Equal(t, false, limiter.AllowN(now.Add(time.Second), "tenant-1", 3).OK())
+	assert.Equal(t, true, limiter.AllowN(now.Add(time.Second), "tenant-1", 2).OK())
 
 	// 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, limiter.AllowN(now, "tenant-2", 18).OK())
+	assert.Equal(t, true, limiter.AllowN(now, "tenant-2", 20).OK())
+	assert.Equal(t, false, limiter.AllowN(now, "tenant-2", 3).OK())
+	assert.Equal(t, true, limiter.AllowN(now, "tenant-2", 2).OK())
+
+	assert.Equal(t, true, limiter.AllowN(now.Add(time.Second), "tenant-2", 18).OK())
+	assert.Equal(t, false, limiter.AllowN(now.Add(time.Second), "tenant-2", 3).OK())
+	assert.Equal(t, true, limiter.AllowN(now.Add(time.Second), "tenant-2", 2).OK())
+}
+
+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, limiter.AllowN(now, "tenant-1", 12).OK())
+	assert.Equal(t, false, limiter.AllowN(now, "tenant-1", 9).OK())
+
+	r1 := limiter.AllowN(now, "tenant-1", 8)
+	assert.Equal(t, true, r1.OK())
+	r1.CancelAt(now)
+
+	assert.Equal(t, true, limiter.AllowN(now, "tenant-1", 8).OK())
+
+	// +10 tokens (1s)
+	nowPlus := now.Add(time.Second)
+
+	assert.Equal(t, true, limiter.AllowN(nowPlus, "tenant-1", 6).OK())
+	assert.Equal(t, false, limiter.AllowN(nowPlus, "tenant-1", 5).OK())
+
+	r2 := limiter.AllowN(nowPlus, "tenant-1", 4)
+	assert.Equal(t, true, r2.OK())
+	r2.CancelAt(nowPlus)
+
+	assert.Equal(t, true, limiter.AllowN(nowPlus, "tenant-1", 2).OK())
+	assert.Equal(t, false, limiter.AllowN(nowPlus, "tenant-1", 3).OK())
+	assert.Equal(t, true, limiter.AllowN(nowPlus, "tenant-1", 2).OK())
+	assert.Equal(t, false, limiter.AllowN(nowPlus, "tenant-1", 1).OK())
 }
 
 func BenchmarkRateLimiter_CustomMultiTenant(b *testing.B) {