feat(kafka): Remove rate limits for kafka ingestion (#14460)

Author: benclive

pkg/ingester/checkpoint_test.go

@@ -459,7 +459,7 @@ func Test_SeriesIterator(t *testing.T) {
 	limits, err := validation.NewOverrides(l, nil)
 	require.NoError(t, err)
 
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	for i := 0; i < 3; i++ {
 		inst, err := newInstance(defaultConfig(), defaultPeriodConfigs, fmt.Sprintf("%d", i), limiter, runtime.DefaultTenantConfigs(), noopWAL{}, NilMetrics, nil, nil, nil, nil, NewStreamRateCalculator(), nil, nil)
@@ -506,7 +506,8 @@ func Benchmark_SeriesIterator(b *testing.B) {
 
 	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
 	require.NoError(b, err)
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	for i := range instances {
 		inst, _ := newInstance(defaultConfig(), defaultPeriodConfigs, fmt.Sprintf("instance %d", i), limiter, runtime.DefaultTenantConfigs(), noopWAL{}, NilMetrics, nil, nil, nil, nil, NewStreamRateCalculator(), nil, nil)

pkg/ingester/ingester.go

@@ -404,18 +404,21 @@ func New(cfg Config, clientConfig client.Config, store Store, limits Limits, con
 		i.SetExtractorWrapper(i.cfg.SampleExtractorWrapper)
 	}
 
-	var limiterStrategy limiterRingStrategy
+	var streamCountLimiter limiterRingStrategy
 	var ownedStreamsStrategy ownershipStrategy
+	var streamRateLimiter RateLimiterStrategy
 	if i.cfg.KafkaIngestion.Enabled {
-		limiterStrategy = newPartitionRingLimiterStrategy(partitionRingWatcher, limits.IngestionPartitionsTenantShardSize)
+		streamCountLimiter = newPartitionRingLimiterStrategy(partitionRingWatcher, limits.IngestionPartitionsTenantShardSize)
 		ownedStreamsStrategy = newOwnedStreamsPartitionStrategy(i.ingestPartitionID, partitionRingWatcher, limits.IngestionPartitionsTenantShardSize, util_log.Logger)
+		streamRateLimiter = &NoLimitsStrategy{} // Kafka ingestion does not have per-stream rate limits, because we control the consumption speed.
 	} else {
-		limiterStrategy = newIngesterRingLimiterStrategy(i.lifecycler, cfg.LifecyclerConfig.RingConfig.ReplicationFactor)
+		streamCountLimiter = newIngesterRingLimiterStrategy(i.lifecycler, cfg.LifecyclerConfig.RingConfig.ReplicationFactor)
 		ownedStreamsStrategy = newOwnedStreamsIngesterStrategy(i.lifecycler.ID, i.readRing, util_log.Logger)
+		streamRateLimiter = &TenantBasedStrategy{limits: limits}
 	}
 	// Now that the lifecycler has been created, we can create the limiter
 	// which depends on it.
-	i.limiter = NewLimiter(limits, metrics, limiterStrategy)
+	i.limiter = NewLimiter(limits, metrics, streamCountLimiter, streamRateLimiter)
 	i.recalculateOwnedStreams = newRecalculateOwnedStreamsSvc(i.getInstances, ownedStreamsStrategy, cfg.OwnedStreamsCheckInterval, util_log.Logger)
 
 	return i, nil

pkg/ingester/instance.go

@@ -307,7 +307,7 @@ func (i *instance) createStream(ctx context.Context, pushReqStream logproto.Stre
 		return nil, fmt.Errorf("failed to create stream: %w", err)
 	}
 
-	s := newStream(chunkfmt, headfmt, i.cfg, i.limiter, i.instanceID, fp, sortedLabels, i.limiter.UnorderedWrites(i.instanceID), i.streamRateCalculator, i.metrics, i.writeFailures, i.configs)
+	s := newStream(chunkfmt, headfmt, i.cfg, i.limiter.rateLimitStrategy, i.instanceID, fp, sortedLabels, i.limiter.UnorderedWrites(i.instanceID), i.streamRateCalculator, i.metrics, i.writeFailures, i.configs)
 
 	// record will be nil when replaying the wal (we don't want to rewrite wal entries as we replay them).
 	if record != nil {
@@ -372,7 +372,7 @@ func (i *instance) createStreamByFP(ls labels.Labels, fp model.Fingerprint) (*st
 		return nil, fmt.Errorf("failed to create stream for fingerprint: %w", err)
 	}
 
-	s := newStream(chunkfmt, headfmt, i.cfg, i.limiter, i.instanceID, fp, sortedLabels, i.limiter.UnorderedWrites(i.instanceID), i.streamRateCalculator, i.metrics, i.writeFailures, i.configs)
+	s := newStream(chunkfmt, headfmt, i.cfg, i.limiter.rateLimitStrategy, i.instanceID, fp, sortedLabels, i.limiter.UnorderedWrites(i.instanceID), i.streamRateCalculator, i.metrics, i.writeFailures, i.configs)
 
 	i.onStreamCreated(s)
 

pkg/ingester/instance_test.go

@@ -78,7 +78,7 @@ var NilMetrics = newIngesterMetrics(nil, constants.Loki)
 func TestLabelsCollisions(t *testing.T) {
 	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
 	require.NoError(t, err)
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	i, err := newInstance(defaultConfig(), defaultPeriodConfigs, "test", limiter, loki_runtime.DefaultTenantConfigs(), noopWAL{}, NilMetrics, &OnceSwitch{}, nil, nil, nil, NewStreamRateCalculator(), nil, nil)
 	require.Nil(t, err)
@@ -106,7 +106,7 @@ func TestLabelsCollisions(t *testing.T) {
 func TestConcurrentPushes(t *testing.T) {
 	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
 	require.NoError(t, err)
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	inst, err := newInstance(defaultConfig(), defaultPeriodConfigs, "test", limiter, loki_runtime.DefaultTenantConfigs(), noopWAL{}, NilMetrics, &OnceSwitch{}, nil, nil, nil, NewStreamRateCalculator(), nil, nil)
 	require.Nil(t, err)
@@ -158,7 +158,7 @@ func TestConcurrentPushes(t *testing.T) {
 func TestGetStreamRates(t *testing.T) {
 	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
 	require.NoError(t, err)
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	inst, err := newInstance(defaultConfig(), defaultPeriodConfigs, "test", limiter, loki_runtime.DefaultTenantConfigs(), noopWAL{}, NilMetrics, &OnceSwitch{}, nil, nil, nil, NewStreamRateCalculator(), nil, nil)
 	require.NoError(t, err)
@@ -245,7 +245,7 @@ func labelHashNoShard(l labels.Labels) uint64 {
 func TestSyncPeriod(t *testing.T) {
 	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
 	require.NoError(t, err)
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	const (
 		syncPeriod = 1 * time.Minute
@@ -290,7 +290,7 @@ func setupTestStreams(t *testing.T) (*instance, time.Time, int) {
 	t.Helper()
 	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
 	require.NoError(t, err)
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 	indexShards := 2
 
 	// just some random values
@@ -315,7 +315,7 @@ func setupTestStreams(t *testing.T) (*instance, time.Time, int) {
 		require.NoError(t, err)
 		chunkfmt, headfmt, err := instance.chunkFormatAt(minTs(&testStream))
 		require.NoError(t, err)
-		chunk := newStream(chunkfmt, headfmt, cfg, limiter, "fake", 0, nil, true, NewStreamRateCalculator(), NilMetrics, nil, nil).NewChunk()
+		chunk := newStream(chunkfmt, headfmt, cfg, limiter.rateLimitStrategy, "fake", 0, nil, true, NewStreamRateCalculator(), NilMetrics, nil, nil).NewChunk()
 		for _, entry := range testStream.Entries {
 			dup, err := chunk.Append(&entry)
 			require.False(t, dup)
@@ -507,7 +507,7 @@ func makeRandomLabels() labels.Labels {
 func Benchmark_PushInstance(b *testing.B) {
 	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
 	require.NoError(b, err)
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	i, _ := newInstance(&Config{IndexShards: 1}, defaultPeriodConfigs, "test", limiter, loki_runtime.DefaultTenantConfigs(), noopWAL{}, NilMetrics, &OnceSwitch{}, nil, nil, nil, NewStreamRateCalculator(), nil, nil)
 	ctx := context.Background()
@@ -549,7 +549,7 @@ func Benchmark_instance_addNewTailer(b *testing.B) {
 	l.MaxLocalStreamsPerUser = 100000
 	limits, err := validation.NewOverrides(l, nil)
 	require.NoError(b, err)
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	ctx := context.Background()
 
@@ -575,7 +575,7 @@ func Benchmark_instance_addNewTailer(b *testing.B) {
 
 	b.Run("addTailersToNewStream", func(b *testing.B) {
 		for n := 0; n < b.N; n++ {
-			inst.addTailersToNewStream(newStream(chunkfmt, headfmt, nil, limiter, "fake", 0, lbs, true, NewStreamRateCalculator(), NilMetrics, nil, nil))
+			inst.addTailersToNewStream(newStream(chunkfmt, headfmt, nil, limiter.rateLimitStrategy, "fake", 0, lbs, true, NewStreamRateCalculator(), NilMetrics, nil, nil))
 		}
 	})
 }
@@ -1089,7 +1089,7 @@ func TestStreamShardingUsage(t *testing.T) {
 	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), limitsDefinition)
 	require.NoError(t, err)
 
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: limits})
 
 	defaultShardStreamsCfg := limiter.limits.ShardStreams("fake")
 	tenantShardStreamsCfg := limiter.limits.ShardStreams(customTenant1)
@@ -1454,7 +1454,7 @@ func defaultInstance(t *testing.T) *instance {
 		&ingesterConfig,
 		defaultPeriodConfigs,
 		"fake",
-		NewLimiter(overrides, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1)),
+		NewLimiter(overrides, NilMetrics, newIngesterRingLimiterStrategy(&ringCountMock{count: 1}, 1), &TenantBasedStrategy{limits: overrides}),
 		loki_runtime.DefaultTenantConfigs(),
 		noopWAL{},
 		NilMetrics,

pkg/ingester/limiter.go

@@ -38,9 +38,10 @@ type Limits interface {
 // Limiter implements primitives to get the maximum number of streams
 // an ingester can handle for a specific tenant
 type Limiter struct {
-	limits       Limits
-	ringStrategy limiterRingStrategy
-	metrics      *ingesterMetrics
+	limits            Limits
+	ringStrategy      limiterRingStrategy
+	metrics           *ingesterMetrics
+	rateLimitStrategy RateLimiterStrategy
 
 	mtx      sync.RWMutex
 	disabled bool
@@ -51,25 +52,28 @@ func (l *Limiter) DisableForWALReplay() {
 	defer l.mtx.Unlock()
 	l.disabled = true
 	l.metrics.limiterEnabled.Set(0)
+	l.rateLimitStrategy.SetDisabled(true)
 }
 
 func (l *Limiter) Enable() {
 	l.mtx.Lock()
 	defer l.mtx.Unlock()
 	l.disabled = false
 	l.metrics.limiterEnabled.Set(1)
+	l.rateLimitStrategy.SetDisabled(false)
 }
 
 type limiterRingStrategy interface {
 	convertGlobalToLocalLimit(int, string) int
 }
 
 // NewLimiter makes a new limiter
-func NewLimiter(limits Limits, metrics *ingesterMetrics, ingesterRingLimiterStrategy limiterRingStrategy) *Limiter {
+func NewLimiter(limits Limits, metrics *ingesterMetrics, ingesterRingLimiterStrategy limiterRingStrategy, rateLimitStrategy RateLimiterStrategy) *Limiter {
 	return &Limiter{
-		limits:       limits,
-		ringStrategy: ingesterRingLimiterStrategy,
-		metrics:      metrics,
+		limits:            limits,
+		ringStrategy:      ingesterRingLimiterStrategy,
+		metrics:           metrics,
+		rateLimitStrategy: rateLimitStrategy,
 	}
 }
 
@@ -231,16 +235,36 @@ func (l *streamCountLimiter) getSuppliers(tenant string) (streamCountSupplier, f
 
 type RateLimiterStrategy interface {
 	RateLimit(tenant string) validation.RateLimit
+	SetDisabled(bool)
 }
 
-func (l *Limiter) RateLimit(tenant string) validation.RateLimit {
+type TenantBasedStrategy struct {
+	disabled bool
+	limits   Limits
+}
+
+func (l *TenantBasedStrategy) RateLimit(tenant string) validation.RateLimit {
 	if l.disabled {
 		return validation.Unlimited
 	}
 
 	return l.limits.PerStreamRateLimit(tenant)
 }
 
+func (l *TenantBasedStrategy) SetDisabled(disabled bool) {
+	l.disabled = disabled
+}
+
+type NoLimitsStrategy struct{}
+
+func (l *NoLimitsStrategy) RateLimit(_ string) validation.RateLimit {
+	return validation.Unlimited
+}
+
+func (l *NoLimitsStrategy) SetDisabled(_ bool) {
+	// no-op
+}
+
 type StreamRateLimiter struct {
 	recheckPeriod time.Duration
 	recheckAt     time.Time

pkg/ingester/limiter_test.go

@@ -133,7 +133,7 @@ func TestStreamCountLimiter_AssertNewStreamAllowed(t *testing.T) {
 				ownedStreamCount: testData.ownedStreamCount,
 			}
 			strategy := &fixedStrategy{localLimit: testData.calculatedLocalLimit}
-			limiter := NewLimiter(limits, NilMetrics, strategy)
+			limiter := NewLimiter(limits, NilMetrics, strategy, &TenantBasedStrategy{limits: limits})
 			defaultCountSupplier := func() int {
 				return testData.streams
 			}
@@ -182,7 +182,7 @@ func TestLimiter_minNonZero(t *testing.T) {
 
 	for testName, testData := range tests {
 		t.Run(testName, func(t *testing.T) {
-			limiter := NewLimiter(nil, NilMetrics, nil)
+			limiter := NewLimiter(nil, NilMetrics, nil, nil)
 			assert.Equal(t, testData.expected, limiter.minNonZero(testData.first, testData.second))
 		})
 	}

pkg/ingester/owned_streams_test.go

@@ -17,7 +17,7 @@ func Test_OwnedStreamService(t *testing.T) {
 	require.NoError(t, err)
 	// Mock the ring
 	ring := &ringCountMock{count: 30}
-	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(ring, 3))
+	limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(ring, 3), &TenantBasedStrategy{limits: limits})
 
 	service := newOwnedStreamService("test", limiter)
 	require.Equal(t, 0, service.getOwnedStreamCount())

pkg/ingester/recalculate_owned_streams_test.go

@@ -70,7 +70,7 @@ func Test_recalculateOwnedStreams_recalculateWithIngesterStrategy(t *testing.T)
 				UseOwnedStreamCount:     testData.featureEnabled,
 			}, nil)
 			require.NoError(t, err)
-			limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(mockRing, 1))
+			limiter := NewLimiter(limits, NilMetrics, newIngesterRingLimiterStrategy(mockRing, 1), &TenantBasedStrategy{limits: limits})
 
 			tenant, err := newInstance(
 				defaultConfig(),