feat: check for failed reason in distributors

pkg/distributor/distributor.go

@@ -718,25 +718,6 @@ func (d *Distributor) PushWithResolver(ctx context.Context, req *logproto.PushRe
 		return &logproto.PushResponse{}, validationErr
 	}
 
-	if d.cfg.IngestLimitsEnabled {
-		streamsAfterLimits, reasonsForHashes, err := d.ingestLimits.enforceLimits(ctx, tenantID, streams)
-		if err != nil {
-			level.Error(d.logger).Log("msg", "failed to check if request exceeds limits, request has been accepted", "err", err)
-		} else if len(streamsAfterLimits) == 0 {
-			// All streams have been dropped.
-			level.Debug(d.logger).Log("msg", "request exceeded limits, all streams will be dropped", "tenant", tenantID)
-			if !d.cfg.IngestLimitsDryRunEnabled {
-				return nil, httpgrpc.Error(http.StatusTooManyRequests, "request exceeded limits: "+firstReasonForHashes(reasonsForHashes))
-			}
-		} else if len(streamsAfterLimits) < len(streams) {
-			// Some streams have been dropped.
-			level.Debug(d.logger).Log("msg", "request exceeded limits, some streams will be dropped", "tenant", tenantID)
-			if !d.cfg.IngestLimitsDryRunEnabled {
-				streams = streamsAfterLimits
-			}
-		}
-	}
-
 	if !d.ingestionRateLimiter.AllowN(now, tenantID, validationContext.validationMetrics.aggregatedPushStats.lineSize) {
 		d.trackDiscardedData(ctx, req, validationContext, tenantID, validationContext.validationMetrics, validation.RateLimited, streamResolver)
 
@@ -746,6 +727,19 @@ func (d *Distributor) PushWithResolver(ctx context.Context, req *logproto.PushRe
 		return nil, httpgrpc.Errorf(http.StatusTooManyRequests, "%s", err.Error())
 	}
 
+	// These limits are checked after the ingestion rate limit as this
+	// is how it works in ingesters.
+	if d.cfg.IngestLimitsEnabled {
+		accepted, err := d.ingestLimits.EnforceLimits(ctx, tenantID, streams)
+		if err == nil && !d.cfg.IngestLimitsDryRunEnabled {
+			if len(accepted) == 0 {
+				// All streams were rejected, the request should be failed.
+				return nil, httpgrpc.Error(http.StatusTooManyRequests, "request exceeded limits")
+			}
+			streams = accepted
+		}
+	}
+
 	// Nil check for performance reasons, to avoid dynamic lookup and/or no-op
 	// function calls that cannot be inlined.
 	if d.tee != nil {

pkg/distributor/distributor_test.go

@@ -2462,7 +2462,7 @@ func TestDistributor_PushIngestLimits(t *testing.T) {
 				Reason:     uint32(limits.ReasonMaxStreams),
 			}},
 		},
-		expectedErr: "rpc error: code = Code(429) desc = request exceeded limits: max streams",
+		expectedErr: "rpc error: code = Code(429) desc = request exceeded limits",
 	}, {
 		name:                "one of two streams exceed max stream limit, request is accepted",
 		ingestLimitsEnabled: true,

pkg/distributor/ingest_limits.go

@@ -18,7 +18,7 @@ import (
 
 // ingestLimitsFrontendClient is used for tests.
 type ingestLimitsFrontendClient interface {
-	exceedsLimits(context.Context, *proto.ExceedsLimitsRequest) (*proto.ExceedsLimitsResponse, error)
+	ExceedsLimits(context.Context, *proto.ExceedsLimitsRequest) (*proto.ExceedsLimitsResponse, error)
 }
 
 // ingestLimitsFrontendRingClient uses the ring to query ingest-limits frontends.
@@ -35,7 +35,7 @@ func newIngestLimitsFrontendRingClient(ring ring.ReadRing, pool *ring_client.Poo
 }
 
 // Implements the ingestLimitsFrontendClient interface.
-func (c *ingestLimitsFrontendRingClient) exceedsLimits(ctx context.Context, req *proto.ExceedsLimitsRequest) (*proto.ExceedsLimitsResponse, error) {
+func (c *ingestLimitsFrontendRingClient) ExceedsLimits(ctx context.Context, req *proto.ExceedsLimitsRequest) (*proto.ExceedsLimitsResponse, error) {
 	// We use an FNV-1 of all stream hashes in the request to load balance requests
 	// to limits-frontends instances.
 	h := fnv.New32()
@@ -78,64 +78,85 @@ func (c *ingestLimitsFrontendRingClient) exceedsLimits(ctx context.Context, req
 
 type ingestLimits struct {
 	client         ingestLimitsFrontendClient
-	limitsFailures prometheus.Counter
+	requests       prometheus.Counter
+	requestsFailed prometheus.Counter
 }
 
 func newIngestLimits(client ingestLimitsFrontendClient, r prometheus.Registerer) *ingestLimits {
 	return &ingestLimits{
 		client: client,
-		limitsFailures: promauto.With(r).NewCounter(prometheus.CounterOpts{
-			Name: "loki_distributor_ingest_limits_failures_total",
-			Help: "The total number of failures checking ingest limits.",
+		requests: promauto.With(r).NewCounter(prometheus.CounterOpts{
+			Name: "loki_distributor_ingest_limits_requests_total",
+			Help: "The total number of requests.",
+		}),
+		requestsFailed: promauto.With(r).NewCounter(prometheus.CounterOpts{
+			Name: "loki_distributor_ingest_limits_requests_failed_total",
+			Help: "The total number of requests that failed.",
 		}),
 	}
 }
 
-// enforceLimits returns a slice of streams that are within the per-tenant
-// limits, and in the case where one or more streams exceed per-tenant
-// limits, the reasons those streams were not included in the result.
-// An error is returned if per-tenant limits could not be enforced.
-func (l *ingestLimits) enforceLimits(ctx context.Context, tenant string, streams []KeyedStream) ([]KeyedStream, map[uint64][]string, error) {
-	exceedsLimits, reasons, err := l.exceedsLimits(ctx, tenant, streams)
-	if !exceedsLimits || err != nil {
-		return streams, nil, err
+// EnforceLimits checks all streams against the per-tenant limits and returns
+// a slice containing the streams that are accepted (within the per-tenant
+// limits). Any streams that could not have their limits checked are also
+// accepted.
+func (l *ingestLimits) EnforceLimits(ctx context.Context, tenant string, streams []KeyedStream) ([]KeyedStream, error) {
+	results, err := l.ExceedsLimits(ctx, tenant, streams)
+	if err != nil {
+		return streams, err
+	}
+	// Fast path. No results means all streams were accepted and there were
+	// no failures, so we can return the input streams.
+	if len(results) == 0 {
+		return streams, nil
 	}
 	// We can do this without allocation if needed, but doing so will modify
 	// the original backing array. See "Filtering without allocation" from
 	// https://go.dev/wiki/SliceTricks.
-	withinLimits := make([]KeyedStream, 0, len(streams))
+	accepted := make([]KeyedStream, 0, len(streams))
 	for _, s := range streams {
-		if _, ok := reasons[s.HashKeyNoShard]; !ok {
-			withinLimits = append(withinLimits, s)
+		// Check each stream to see if it failed.
+		// TODO(grobinson): We have an O(N*M) loop here. Need to benchmark if
+		// its faster to do this or if we should create a map instead.
+		var (
+			found  bool
+			reason uint32
+		)
+		for _, res := range results {
+			if res.StreamHash == s.HashKeyNoShard {
+				found = true
+				reason = res.Reason
+				break
+			}
+		}
+		if !found || reason == uint32(limits.ReasonFailed) {
+			accepted = append(accepted, s)
 		}
 	}
-	return withinLimits, reasons, nil
+	return accepted, nil
 }
 
-// ExceedsLimits returns true if one or more streams exceeds per-tenant limits,
-// and false if no streams exceed per-tenant limits. In the case where one or
-// more streams exceeds per-tenant limits, it returns the reasons for each stream.
-// An error is returned if per-tenant limits could not be checked.
-func (l *ingestLimits) exceedsLimits(ctx context.Context, tenant string, streams []KeyedStream) (bool, map[uint64][]string, error) {
+// ExceedsLimits checks all streams against the per-tenant limits. It returns
+// an error if the client failed to send the request or receive a response
+// from the server. Any streams that could not have their limits checked
+// and returned in the results with the reason "ReasonFailed".
+func (l *ingestLimits) ExceedsLimits(
+	ctx context.Context,
+	tenant string,
+	streams []KeyedStream,
+) ([]*proto.ExceedsLimitsResult, error) {
+	l.requests.Inc()
 	req, err := newExceedsLimitsRequest(tenant, streams)
 	if err != nil {
-		return false, nil, err
+		l.requestsFailed.Inc()
+		return nil, err
 	}
-	resp, err := l.client.exceedsLimits(ctx, req)
+	resp, err := l.client.ExceedsLimits(ctx, req)
 	if err != nil {
-		return false, nil, err
-	}
-	if len(resp.Results) == 0 {
-		return false, nil, nil
+		l.requestsFailed.Inc()
+		return nil, err
 	}
-	reasonsForHashes := make(map[uint64][]string)
-	for _, result := range resp.Results {
-		reasons := reasonsForHashes[result.StreamHash]
-		humanized := limits.Reason(result.Reason).String()
-		reasons = append(reasons, humanized)
-		reasonsForHashes[result.StreamHash] = reasons
-	}
-	return true, reasonsForHashes, nil
+	return resp.Results, nil
 }
 
 func newExceedsLimitsRequest(tenant string, streams []KeyedStream) (*proto.ExceedsLimitsRequest, error) {
@@ -156,10 +177,3 @@ func newExceedsLimitsRequest(tenant string, streams []KeyedStream) (*proto.Excee
 		Streams: streamMetadata,
 	}, nil
 }
-
-func firstReasonForHashes(reasonsForHashes map[uint64][]string) string {
-	for _, reasons := range reasonsForHashes {
-		return reasons[0]
-	}
-	return "unknown reason"
-}

pkg/distributor/ingest_limits_test.go

@@ -26,7 +26,7 @@ type mockIngestLimitsFrontendClient struct {
 }
 
 // Implements the ingestLimitsFrontendClient interface.
-func (c *mockIngestLimitsFrontendClient) exceedsLimits(_ context.Context, r *proto.ExceedsLimitsRequest) (*proto.ExceedsLimitsResponse, error) {
+func (c *mockIngestLimitsFrontendClient) ExceedsLimits(_ context.Context, r *proto.ExceedsLimitsRequest) (*proto.ExceedsLimitsResponse, error) {
 	c.calls.Add(1)
 	if c.expectedRequest != nil {
 		require.Equal(c.t, c.expectedRequest, r)
@@ -49,7 +49,6 @@ func TestIngestLimits_EnforceLimits(t *testing.T) {
 		response        *proto.ExceedsLimitsResponse
 		responseErr     error
 		expectedStreams []KeyedStream
-		expectedReasons map[uint64][]string
 		expectedErr     string
 	}{{
 		// This test also asserts that streams are returned unmodified.
@@ -116,7 +115,6 @@ func TestIngestLimits_EnforceLimits(t *testing.T) {
 			}},
 		},
 		expectedStreams: []KeyedStream{},
-		expectedReasons: map[uint64][]string{1: {"max streams"}},
 	}, {
 		name:   "one of two streams exceeds limits",
 		tenant: "test",
@@ -145,7 +143,6 @@ func TestIngestLimits_EnforceLimits(t *testing.T) {
 			HashKey:        2000, // Should not be used.
 			HashKeyNoShard: 2,
 		}},
-		expectedReasons: map[uint64][]string{1: {"max streams"}},
 	}, {
 		name:   "does not exceed limits",
 		tenant: "test",
@@ -174,7 +171,6 @@ func TestIngestLimits_EnforceLimits(t *testing.T) {
 			HashKey:        2000, // Should not be used.
 			HashKeyNoShard: 2,
 		}},
-		expectedReasons: nil,
 	}}
 
 	for _, test := range tests {
@@ -188,35 +184,29 @@ func TestIngestLimits_EnforceLimits(t *testing.T) {
 			l := newIngestLimits(&mockClient, prometheus.NewRegistry())
 			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 			defer cancel()
-			streams, reasons, err := l.enforceLimits(ctx, test.tenant, test.streams)
+			accepted, err := l.EnforceLimits(ctx, test.tenant, test.streams)
 			if test.expectedErr != "" {
 				require.EqualError(t, err, test.expectedErr)
 				// The streams should be returned unmodified.
-				require.Equal(t, test.streams, streams)
-				require.Nil(t, reasons)
+				require.Equal(t, test.streams, accepted)
 			} else {
 				require.Nil(t, err)
-				require.Equal(t, test.expectedStreams, streams)
-				require.Equal(t, test.expectedReasons, reasons)
+				require.Equal(t, test.expectedStreams, accepted)
 			}
 		})
 	}
 }
 
-// This test asserts that when checking ingest limits the expected proto
-// message is sent, and that for a given response, the result contains the
-// expected streams each with their expected reasons.
 func TestIngestLimits_ExceedsLimits(t *testing.T) {
 	tests := []struct {
-		name                  string
-		tenant                string
-		streams               []KeyedStream
-		expectedRequest       *proto.ExceedsLimitsRequest
-		response              *proto.ExceedsLimitsResponse
-		responseErr           error
-		expectedExceedsLimits bool
-		expectedReasons       map[uint64][]string
-		expectedErr           string
+		name            string
+		tenant          string
+		streams         []KeyedStream
+		expectedRequest *proto.ExceedsLimitsRequest
+		response        *proto.ExceedsLimitsResponse
+		responseErr     error
+		expectedResult  []*proto.ExceedsLimitsResult
+		expectedErr     string
 	}{{
 		name:   "error should be returned if limits cannot be checked",
 		tenant: "test",
@@ -249,8 +239,10 @@ func TestIngestLimits_ExceedsLimits(t *testing.T) {
 				Reason:     uint32(limits.ReasonMaxStreams),
 			}},
 		},
-		expectedExceedsLimits: true,
-		expectedReasons:       map[uint64][]string{1: {"max streams"}},
+		expectedResult: []*proto.ExceedsLimitsResult{{
+			StreamHash: 1,
+			Reason:     uint32(limits.ReasonMaxStreams),
+		}},
 	}, {
 		name:   "does not exceed limits",
 		tenant: "test",
@@ -266,7 +258,7 @@ func TestIngestLimits_ExceedsLimits(t *testing.T) {
 		response: &proto.ExceedsLimitsResponse{
 			Results: []*proto.ExceedsLimitsResult{},
 		},
-		expectedReasons: nil,
+		expectedResult: []*proto.ExceedsLimitsResult{},
 	}}
 
 	for _, test := range tests {
@@ -280,15 +272,13 @@ func TestIngestLimits_ExceedsLimits(t *testing.T) {
 			l := newIngestLimits(&mockClient, prometheus.NewRegistry())
 			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 			defer cancel()
-			exceedsLimits, reasons, err := l.exceedsLimits(ctx, test.tenant, test.streams)
+			res, err := l.ExceedsLimits(ctx, test.tenant, test.streams)
 			if test.expectedErr != "" {
 				require.EqualError(t, err, test.expectedErr)
-				require.False(t, exceedsLimits)
-				require.Nil(t, reasons)
+				require.Nil(t, res)
 			} else {
 				require.Nil(t, err)
-				require.Equal(t, test.expectedExceedsLimits, exceedsLimits)
-				require.Equal(t, test.expectedReasons, reasons)
+				require.Equal(t, test.expectedResult, res)
 			}
 		})
 	}