feat: add support for UpdateRates RPC to distributors

Author: grobinson-grafana

pkg/distributor/ingest_limits.go

@@ -15,9 +15,10 @@ import (
 	"github.com/grafana/loki/v3/pkg/limits/proto"
 )
 
-// ingestLimitsFrontendClient is used for tests.
+// The ingestLimitsFrontendClient interface is used to mock calls in tests.
 type ingestLimitsFrontendClient interface {
 	ExceedsLimits(context.Context, *proto.ExceedsLimitsRequest) (*proto.ExceedsLimitsResponse, error)
+	UpdateRates(context.Context, *proto.UpdateRatesRequest) (*proto.UpdateRatesResponse, error)
 }
 
 // ingestLimitsFrontendRingClient uses the ring to discover ingest-limits-frontend
@@ -34,7 +35,7 @@ func newIngestLimitsFrontendRingClient(ring ring.ReadRing, pool *ring_client.Poo
 	}
 }
 
-// Implements the ingestLimitsFrontendClient interface.
+// Implements the [ingestLimitsFrontendClient] interface.
 func (c *ingestLimitsFrontendRingClient) ExceedsLimits(ctx context.Context, req *proto.ExceedsLimitsRequest) (*proto.ExceedsLimitsResponse, error) {
 	var resp *proto.ExceedsLimitsResponse
 	err := c.withRandomShuffle(ctx, func(ctx context.Context, client proto.IngestLimitsFrontendClient) error {
@@ -45,6 +46,17 @@ func (c *ingestLimitsFrontendRingClient) ExceedsLimits(ctx context.Context, req
 	return resp, err
 }
 
+// Implements the [ingestLimitsFrontendClient] interface.
+func (c *ingestLimitsFrontendRingClient) UpdateRates(ctx context.Context, req *proto.UpdateRatesRequest) (*proto.UpdateRatesResponse, error) {
+	var resp *proto.UpdateRatesResponse
+	err := c.withRandomShuffle(ctx, func(ctx context.Context, client proto.IngestLimitsFrontendClient) error {
+		var clientErr error
+		resp, clientErr = client.UpdateRates(ctx, req)
+		return clientErr
+	})
+	return resp, err
+}
+
 // withRandomShuffle gets all healthy frontends in the ring, randomly shuffles
 // them, and then calls f.
 func (c *ingestLimitsFrontendRingClient) withRandomShuffle(ctx context.Context, f func(ctx context.Context, client proto.IngestLimitsFrontendClient) error) error {
@@ -82,21 +94,21 @@ func (c *ingestLimitsFrontendRingClient) withRandomShuffle(ctx context.Context,
 
 type ingestLimits struct {
 	client         ingestLimitsFrontendClient
-	requests       prometheus.Counter
-	requestsFailed prometheus.Counter
+	requests       *prometheus.CounterVec
+	requestsFailed *prometheus.CounterVec
 }
 
 func newIngestLimits(client ingestLimitsFrontendClient, r prometheus.Registerer) *ingestLimits {
 	return &ingestLimits{
 		client: client,
-		requests: promauto.With(r).NewCounter(prometheus.CounterOpts{
+		requests: promauto.With(r).NewCounterVec(prometheus.CounterOpts{
 			Name: "loki_distributor_ingest_limits_requests_total",
 			Help: "The total number of requests.",
-		}),
-		requestsFailed: promauto.With(r).NewCounter(prometheus.CounterOpts{
+		}, []string{"operation"}),
+		requestsFailed: promauto.With(r).NewCounterVec(prometheus.CounterOpts{
 			Name: "loki_distributor_ingest_limits_requests_failed_total",
 			Help: "The total number of requests that failed.",
-		}),
+		}, []string{"operation"}),
 	}
 }
 
@@ -144,20 +156,16 @@ func (l *ingestLimits) EnforceLimits(ctx context.Context, tenant string, streams
 // 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()
+func (l *ingestLimits) ExceedsLimits(ctx context.Context, tenant string, streams []KeyedStream) ([]*proto.ExceedsLimitsResult, error) {
+	l.requests.WithLabelValues("ExceedsLimits").Inc()
 	req, err := newExceedsLimitsRequest(tenant, streams)
 	if err != nil {
 		l.requestsFailed.Inc()
 		return nil, err
 	}
 	resp, err := l.client.ExceedsLimits(ctx, req)
 	if err != nil {
-		l.requestsFailed.Inc()
+		l.requestsFailed.WithLabelValues("ExceedsLimits").Inc()
 		return nil, err
 	}
 	return resp.Results, nil
@@ -182,3 +190,41 @@ func newExceedsLimitsRequest(tenant string, streams []KeyedStream) (*proto.Excee
 		Streams: streamMetadata,
 	}, nil
 }
+
+// UpdateRates updates the rates for the streams and returns a slice of the
+// updated rates for all streams. Any streams that could not have rates updated
+// have a rate of zero.
+func (l *ingestLimits) UpdateRates(ctx context.Context, tenant string, streams []KeyedStream) ([]*proto.UpdateRatesResult, error) {
+	l.requests.WithLabelValues("UpdateRates").Inc()
+	req, err := newUpdateRatesRequest(tenant, streams)
+	if err != nil {
+		l.requestsFailed.Inc()
+		return nil, err
+	}
+	resp, err := l.client.UpdateRates(ctx, req)
+	if err != nil {
+		l.requestsFailed.WithLabelValues("UpdateRates").Inc()
+		return nil, err
+	}
+	return resp.Results, nil
+}
+
+func newUpdateRatesRequest(tenant string, streams []KeyedStream) (*proto.UpdateRatesRequest, error) {
+	// The distributor sends the hashes of all streams in the request to the
+	// limits-frontend. The limits-frontend is responsible for deciding if
+	// the request would exceed the tenants limits, and if so, which streams
+	// from the request caused it to exceed its limits.
+	streamMetadata := make([]*proto.StreamMetadata, 0, len(streams))
+	for _, stream := range streams {
+		entriesSize, structuredMetadataSize := calculateStreamSizes(stream.Stream)
+		streamMetadata = append(streamMetadata, &proto.StreamMetadata{
+			StreamHash:      stream.HashKeyNoShard,
+			TotalSize:       entriesSize + structuredMetadataSize,
+			IngestionPolicy: stream.Policy,
+		})
+	}
+	return &proto.UpdateRatesRequest{
+		Tenant:  tenant,
+		Streams: streamMetadata,
+	}, nil
+}