feat(ingest-limits): Implement global tenant rate limiting (#16727)

Author: periklis

docs/sources/shared/configuration.md

@@ -883,6 +883,16 @@ ingest_limits:
   # CLI flag: -ingest-limits.window-size
   [window_size: <duration> | default = 1h]
 
+  # The time window for rate calculation. This should match the window used in
+  # Prometheus rate() queries for consistency.
+  # CLI flag: -ingest-limits.rate-window
+  [rate_window: <duration> | default = 5m]
+
+  # The granularity of time buckets used for sliding window rate calculation.
+  # Smaller buckets provide more precise rate tracking but require more memory.
+  # CLI flag: -ingest-limits.bucket-duration
+  [bucket_duration: <duration> | default = 1m]
+
   lifecycler:
     ring:
       kvstore:
@@ -1186,6 +1196,10 @@ ingest_limits_frontend:
     # CLI flag: -ingest-limits-frontend.lifecycler.ID
     [id: <string> | default = "<hostname>"]
 
+  # The period to recheck per tenant ingestion rate limit configuration.
+  # CLI flag: -ingest-limits-frontend.recheck-period
+  [recheck_period: <duration> | default = 10s]
+
 ingest_limits_frontend_client:
   # Configures client gRPC connections to limits service.
   # The CLI flags prefix for this block configuration is:

pkg/distributor/distributor.go

@@ -1161,21 +1161,27 @@ func (d *Distributor) exceedsLimits(ctx context.Context, tenantID string, stream
 	// 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.
-	streamHashes := make([]*logproto.StreamMetadata, 0, len(streams))
+	streamMetadata := make([]*logproto.StreamMetadata, 0, len(streams))
 	for _, stream := range streams {
 		// Add the stream hash to FNV-1.
 		buf := make([]byte, binary.MaxVarintLen64)
 		binary.PutUvarint(buf, stream.HashKeyNoShard)
 		_, _ = h.Write(buf)
+
+		// Calculate the size of the stream.
+		entriesSize, structuredMetadataSize := calculateStreamSizes(stream.Stream)
+
 		// Add the stream hash to the request. This is sent to limits-frontend.
-		streamHashes = append(streamHashes, &logproto.StreamMetadata{
-			StreamHash: stream.HashKeyNoShard,
+		streamMetadata = append(streamMetadata, &logproto.StreamMetadata{
+			StreamHash:             stream.HashKeyNoShard,
+			EntriesSize:            entriesSize,
+			StructuredMetadataSize: structuredMetadataSize,
 		})
 	}
 
 	req := logproto.ExceedsLimitsRequest{
 		Tenant:  tenantID,
-		Streams: streamHashes,
+		Streams: streamMetadata,
 	}
 
 	// Get the limits-frontend instances from the ring.
@@ -1275,6 +1281,8 @@ func (d *Distributor) sendStreamToKafka(ctx context.Context, stream KeyedStream,
 		return fmt.Errorf("failed to marshal write request to records: %w", err)
 	}
 
+	entriesSize, structuredMetadataSize := calculateStreamSizes(stream.Stream)
+
 	// However, unlike stream records, the distributor writes stream metadata
 	// records to one of a fixed number of partitions, the size of which is
 	// determined ahead of time. It does not use a ring. The reason for this
@@ -1286,6 +1294,8 @@ func (d *Distributor) sendStreamToKafka(ctx context.Context, stream KeyedStream,
 		d.cfg.KafkaConfig.Topic,
 		tenant,
 		stream.HashKeyNoShard,
+		entriesSize,
+		structuredMetadataSize,
 	)
 	if err != nil {
 		return fmt.Errorf("failed to marshal metadata: %w", err)
@@ -1403,6 +1413,15 @@ func calculateShards(rate int64, pushSize, desiredRate int) int {
 	return int(math.Ceil(shards))
 }
 
+func calculateStreamSizes(stream logproto.Stream) (uint64, uint64) {
+	var entriesSize, structuredMetadataSize uint64
+	for _, entry := range stream.Entries {
+		entriesSize += uint64(len(entry.Line))
+		structuredMetadataSize += uint64(util.StructuredMetadataSize(entry.StructuredMetadata))
+	}
+	return entriesSize, structuredMetadataSize
+}
+
 // newRingAndLifecycler creates a new distributor ring and lifecycler with all required lifecycler delegates
 func newRingAndLifecycler(cfg RingConfig, instanceCount *atomic.Uint32, logger log.Logger, reg prometheus.Registerer, metricsNamespace string) (*ring.Ring, *ring.BasicLifecycler, error) {
 	kvStore, err := kv.NewClient(cfg.KVStore, ring.GetCodec(), kv.RegistererWithKVName(reg, "distributor-lifecycler"), logger)

pkg/kafka/encoding.go

@@ -197,14 +197,16 @@ func sovPush(x uint64) (n int) {
 
 // EncodeStreamMetadata encodes the stream metadata into a Kafka record
 // using the tenantID as the key and partition as the target partition
-func EncodeStreamMetadata(partition int32, topic string, tenantID string, streamHash uint64) (*kgo.Record, error) {
+func EncodeStreamMetadata(partition int32, topic, tenantID string, streamHash, entriesSize, structuredMetadataSize uint64) (*kgo.Record, error) {
 	// Validate stream hash
 	if streamHash == 0 {
 		return nil, fmt.Errorf("invalid stream hash '%d'", streamHash)
 	}
 
 	metadata := logproto.StreamMetadata{
-		StreamHash: streamHash,
+		StreamHash:             streamHash,
+		EntriesSize:            entriesSize,
+		StructuredMetadataSize: structuredMetadataSize,
 	}
 
 	// Encode the metadata into a byte slice

pkg/kafka/encoding_test.go

@@ -153,35 +153,51 @@ func generateRandomString(length int) string {
 
 func TestEncodeDecodeStreamMetadata(t *testing.T) {
 	tests := []struct {
-		name      string
-		hash      uint64
-		partition int32
-		topic     string
-		tenantID  string
-		expectErr bool
+		name                   string
+		hash                   uint64
+		partition              int32
+		topic                  string
+		tenantID               string
+		entriesSize            uint64
+		structuredMetadataSize uint64
+		expectErr              bool
 	}{
 		{
-			name:      "Valid metadata",
-			hash:      12345,
-			partition: 1,
-			topic:     "logs",
-			tenantID:  "tenant-1",
-			expectErr: false,
+			name:                   "Valid metadata",
+			hash:                   12345,
+			partition:              1,
+			topic:                  "logs",
+			tenantID:               "tenant-1",
+			entriesSize:            1024,
+			structuredMetadataSize: 512,
+			expectErr:              false,
 		},
 		{
-			name:      "Zero hash - should error",
-			hash:      0,
-			partition: 3,
-			topic:     "traces",
-			tenantID:  "tenant-3",
-			expectErr: true,
+			name:                   "Valid metadata with zero sizes",
+			hash:                   67890,
+			partition:              2,
+			topic:                  "metrics",
+			tenantID:               "tenant-2",
+			entriesSize:            0,
+			structuredMetadataSize: 0,
+			expectErr:              false,
+		},
+		{
+			name:                   "Zero hash - should error",
+			hash:                   0,
+			partition:              3,
+			topic:                  "traces",
+			tenantID:               "tenant-3",
+			entriesSize:            2048,
+			structuredMetadataSize: 1024,
+			expectErr:              true,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// Encode metadata
-			record, err := EncodeStreamMetadata(tt.partition, tt.topic, tt.tenantID, tt.hash)
+			record, err := EncodeStreamMetadata(tt.partition, tt.topic, tt.tenantID, tt.hash, tt.entriesSize, tt.structuredMetadataSize)
 			if tt.expectErr {
 				require.Error(t, err)
 				require.Nil(t, record)
@@ -201,6 +217,8 @@ func TestEncodeDecodeStreamMetadata(t *testing.T) {
 
 			// Verify decoded values
 			require.Equal(t, tt.hash, metadata.StreamHash)
+			require.Equal(t, tt.entriesSize, metadata.EntriesSize)
+			require.Equal(t, tt.structuredMetadataSize, metadata.StructuredMetadataSize)
 		})
 	}
 

pkg/limits/frontend/frontend.go

@@ -11,9 +11,11 @@ import (
 	"flag"
 	"fmt"
 	"net/http"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/grafana/dskit/limiter"
 	"github.com/grafana/dskit/ring"
 	"github.com/grafana/dskit/services"
 	"github.com/grafana/dskit/user"
@@ -34,11 +36,13 @@ const (
 type Config struct {
 	ClientConfig     limits_client.Config  `yaml:"client_config"`
 	LifecyclerConfig ring.LifecyclerConfig `yaml:"lifecycler,omitempty"`
+	RecheckPeriod    time.Duration         `yaml:"recheck_period"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	cfg.ClientConfig.RegisterFlagsWithPrefix("ingest-limits-frontend", f)
 	cfg.LifecyclerConfig.RegisterFlagsWithPrefix("ingest-limits-frontend.", f, util_log.Logger)
+	f.DurationVar(&cfg.RecheckPeriod, "ingest-limits-frontend.recheck-period", 10*time.Second, "The period to recheck per tenant ingestion rate limit configuration.")
 }
 
 func (cfg *Config) Validate() error {
@@ -71,7 +75,8 @@ func New(cfg Config, ringName string, limitsRing ring.ReadRing, limits Limits, l
 
 	factory := limits_client.NewPoolFactory(cfg.ClientConfig)
 	pool := limits_client.NewPool(ringName, cfg.ClientConfig.PoolConfig, limitsRing, factory, logger)
-	limitsSrv := NewRingIngestLimitsService(limitsRing, pool, limits, logger, reg)
+	rateLimiter := limiter.NewRateLimiter(newIngestionRateStrategy(limits), cfg.RecheckPeriod)
+	limitsSrv := NewRingIngestLimitsService(limitsRing, pool, limits, rateLimiter, logger, reg)
 
 	f := &Frontend{
 		cfg:    cfg,