hanzo/logs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Introduce a rate batcher to batch rate updates (#20784)

Browse Source
This commit is contained in:
Dylan Guedes
2026-02-12 17:31:50 -03:00
committed by GitHub
parent d85fb0317f
commit b9a0a15f2a
6 changed files with 679 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/sources/shared/configuration.md
+5
View File
@@ -3318,6 +3318,11 @@ dataobj_tee:
# Enables optional debug metrics.
# CLI flag: -distributor.dataobj-tee.debug-metrics-enabled
[debug_metrics_enabled: <boolean> | default = false]
# Duration to accumulate rate updates before sending to limits frontend. Set
# to 0 to disable batching.
# CLI flag: -distributor.dataobj-tee.rate-batch-window
[rate_batch_window: <duration> | default = 0s]
```
### etcd
pkg/distributor/dataobj_tee.go
+55 -16
View File
@@ -6,9 +6,11 @@ import (
"flag"
"fmt"
"strconv"
"time"
"github.com/go-kit/log"
"github.com/go-kit/log/level"
"github.com/grafana/dskit/services"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promauto"
"github.com/twmb/franz-go/pkg/kgo"
@@ -26,11 +28,12 @@ const (
)
type DataObjTeeConfig struct {
Enabled bool `yaml:"enabled"`
Topic string `yaml:"topic"`
MaxBufferedBytes int `yaml:"max_buffered_bytes"`
PerPartitionRateBytes int `yaml:"per_partition_rate_bytes"`
DebugMetricsEnabled bool `yaml:"debug_metrics_enabled"`
Enabled bool `yaml:"enabled"`
Topic string `yaml:"topic"`
MaxBufferedBytes int `yaml:"max_buffered_bytes"`
PerPartitionRateBytes int `yaml:"per_partition_rate_bytes"`
DebugMetricsEnabled bool `yaml:"debug_metrics_enabled"`
RateBatchWindow time.Duration `yaml:"rate_batch_window"`
}
func (c *DataObjTeeConfig) RegisterFlags(f *flag.FlagSet) {
@@ -39,6 +42,7 @@ func (c *DataObjTeeConfig) RegisterFlags(f *flag.FlagSet) {
f.IntVar(&c.MaxBufferedBytes, "distributor.dataobj-tee.max-buffered-bytes", 100<<20, "Maximum number of bytes to buffer.")
f.IntVar(&c.PerPartitionRateBytes, "distributor.dataobj-tee.per-partition-rate-bytes", 1024*1024, "The per-tenant partition rate (bytes/sec).")
f.BoolVar(&c.DebugMetricsEnabled, "distributor.dataobj-tee.debug-metrics-enabled", false, "Enables optional debug metrics.")
f.DurationVar(&c.RateBatchWindow, "distributor.dataobj-tee.rate-batch-window", 0, "Duration to accumulate rate updates before sending to limits frontend. Set to 0 to disable batching.")
}
func (c *DataObjTeeConfig) Validate() error {
@@ -62,6 +66,7 @@ func (c *DataObjTeeConfig) Validate() error {
type DataObjTee struct {
cfg *DataObjTeeConfig
limitsClient *ingestLimits
rateBatcher *rateBatcher // nil if batching is disabled
limits Limits
kafkaClient *kgo.Client
resolver *SegmentationPartitionResolver
@@ -84,7 +89,7 @@ func NewDataObjTee(
logger log.Logger,
r prometheus.Registerer,
) (*DataObjTee, error) {
return &DataObjTee{
t := &DataObjTee{
cfg: cfg,
resolver: resolver,
kafkaClient: kafkaClient,
@@ -109,7 +114,21 @@ func NewDataObjTee(
Name: "loki_distributor_dataobj_tee_produced_records_total",
Help: "Total number of records produced to each partition.",
}, []string{"partition", "tenant", "segmentation_key"}),
}, nil
}
// Create rate batcher if batching is enabled.
if cfg.RateBatchWindow > 0 {
t.rateBatcher = newRateBatcher(
RateBatcherConfig{
BatchWindow: cfg.RateBatchWindow,
},
limitsClient,
logger,
r,
)
}
return t, nil
}
// A SegmentedStream is a KeyedStream with a segmentation key.
@@ -140,15 +159,26 @@ func (t *DataObjTee) Duplicate(ctx context.Context, tenant string, streams []Key
SegmentationKeyHash: segmentationKey.Sum64(),
})
}
rates, err := t.limitsClient.UpdateRates(ctx, tenant, segmentationKeyStreams)
if err != nil {
level.Error(t.logger).Log("msg", "failed to update rates", "err", err)
}
// fastRates is a temporary lookup table that lets us find the rate
// for a segmentation key in constant time.
fastRates := make(map[uint64]uint64, len(rates))
for _, rate := range rates {
fastRates[rate.StreamHash] = rate.Rate
// fastRates is a lookup table that lets us find the rate for a segmentation
// key in constant time.
var fastRates map[uint64]uint64
if t.rateBatcher != nil {
// Batching enabled: add to batch and get last known rates
// (stats from this push request aren't included as we don't know if the push would be accepted).
// New streams will have rate=0 until first flush includes them.
fastRates = t.rateBatcher.Add(tenant, segmentationKeyStreams)
} else {
// Batching disabled: call UpdateRates synchronously.
fastRates = make(map[uint64]uint64, len(segmentationKeyStreams))
rates, err := t.limitsClient.UpdateRates(ctx, tenant, segmentationKeyStreams)
if err != nil {
level.Error(t.logger).Log("msg", "failed to update rates", "err", err)
}
for _, rate := range rates {
fastRates[rate.StreamHash] = rate.Rate
}
}
// We use max to prevent negative values becoming large positive values
@@ -216,3 +246,12 @@ func (t *DataObjTee) observeDuplicate(partition int32, tenant, segmentationKey s
segmentationKeyLabelValue,
).Inc()
}
// RateBatcher returns the rate batcher service if batching is enabled, nil otherwise.
// This is used to add the batcher to the distributor's subservices for lifecycle management.
func (t *DataObjTee) RateBatcher() services.Service {
if t.rateBatcher == nil {
return nil
}
return t.rateBatcher
}
pkg/distributor/distributor.go
+4
View File
@@ -327,6 +327,10 @@ func New(
return nil, fmt.Errorf("failed to create data object tee: %w", err)
}
tee = WrapTee(tee, dataObjTee)
if rateBatcher := dataObjTee.RateBatcher(); rateBatcher != nil {
servs = append(servs, rateBatcher)
}
}
}
pkg/distributor/ingest_limits.go
+9 -1
View File
@@ -202,12 +202,20 @@ func newExceedsLimitsRequest(tenant string, streams []KeyedStream) (*proto.Excee
// 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 []SegmentedStream) ([]*proto.UpdateRatesResult, error) {
l.requests.WithLabelValues("UpdateRates").Inc()
req, err := newUpdateRatesRequest(tenant, streams)
if err != nil {
// We update `UpdateRates` here because we have clients directly calling `UpdateRatesRaw`.
l.requests.WithLabelValues("UpdateRates").Inc()
l.requestsFailed.WithLabelValues("UpdateRates").Inc()
return nil, err
}
return l.UpdateRatesRaw(ctx, req)
}
// UpdateRatesRaw sends a pre-built UpdateRatesRequest to the frontend.
// This is used by the rate batcher which accumulates stream data over time.
func (l *ingestLimits) UpdateRatesRaw(ctx context.Context, req *proto.UpdateRatesRequest) ([]*proto.UpdateRatesResult, error) {
l.requests.WithLabelValues("UpdateRates").Inc()
resp, err := l.client.UpdateRates(ctx, req)
if err != nil {
l.requestsFailed.WithLabelValues("UpdateRates").Inc()
pkg/distributor/rate_batcher.go
+250
View File
@@ -0,0 +1,250 @@
package distributor
import (
"context"
"sync"
"time"
"github.com/go-kit/log"
"github.com/go-kit/log/level"
"github.com/grafana/dskit/services"
"github.com/grafana/dskit/user"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promauto"
"github.com/grafana/loki/v3/pkg/limits/proto"
)
// RateBatcherConfig contains the configuration for the RateBatcher.
type RateBatcherConfig struct {
// BatchWindow is the duration to accumulate rate updates before flushing.
BatchWindow time.Duration
}
// rateBatcherClient is the interface for sending rate updates.
// This allows the batcher to use the ingestLimits wrapper which tracks metrics.
type rateBatcherClient interface {
UpdateRatesRaw(ctx context.Context, req *proto.UpdateRatesRequest) ([]*proto.UpdateRatesResult, error)
}
// rateBatcher accumulates UpdateRates requests and dispatches them in batches.
// This is a fire-and-forget mechanism - callers add streams to the batch and
// don't wait for results. The batch is periodically flushed to the frontend.
// Results from UpdateRates are stored and can be looked up for partition resolution.
type rateBatcher struct {
services.Service
cfg RateBatcherConfig
client rateBatcherClient
logger log.Logger
// pending accumulates streams to be sent in the next batch.
// Map: tenant -> segmentationKeyHash -> *proto.StreamMetadata
pendingMu sync.Mutex
pending map[string]map[uint64]*proto.StreamMetadata
// rates stores the last known rate for each stream, updated on each flush.
// Map: tenant -> segmentationKeyHash -> rate (bytes/sec)
ratesMu sync.RWMutex
rates map[string]map[uint64]uint64
// Metrics
batchesSent prometheus.Counter
batchesFailed prometheus.Counter
streamsPerBatch prometheus.Histogram
pendingStreams prometheus.Gauge
streamsFlushed prometheus.Counter
}
// newRateBatcher creates a new rate batcher.
func newRateBatcher(cfg RateBatcherConfig, client rateBatcherClient, logger log.Logger, reg prometheus.Registerer) *rateBatcher {
b := &rateBatcher{
cfg: cfg,
client: client,
logger: logger,
pending: make(map[string]map[uint64]*proto.StreamMetadata),
rates: make(map[string]map[uint64]uint64),
batchesSent: promauto.With(reg).NewCounter(prometheus.CounterOpts{
Name: "loki_distributor_rate_batcher_batches_sent_total",
Help: "Total number of batches sent to the limits frontend.",
}),
batchesFailed: promauto.With(reg).NewCounter(prometheus.CounterOpts{
Name: "loki_distributor_rate_batcher_batches_failed_total",
Help: "Total number of batches that failed to send.",
}),
streamsPerBatch: promauto.With(reg).NewHistogram(prometheus.HistogramOpts{
Name: "loki_distributor_rate_batcher_streams_per_batch",
Help: "Number of unique streams per batch.",
Buckets: prometheus.ExponentialBuckets(10, 2, 10), // 10, 20, 40, ... 5120
}),
pendingStreams: promauto.With(reg).NewGauge(prometheus.GaugeOpts{
Name: "loki_distributor_rate_batcher_pending_streams",
Help: "Current number of streams pending in the batch.",
}),
streamsFlushed: promauto.With(reg).NewCounter(prometheus.CounterOpts{
Name: "loki_distributor_rate_batcher_streams_flushed_total",
Help: "Total number of streams flushed to the limits frontend.",
}),
}
b.Service = services.NewBasicService(nil, b.running, nil)
return b
}
// running is the main loop that periodically flushes batches.
func (b *rateBatcher) running(ctx context.Context) error {
ticker := time.NewTicker(b.cfg.BatchWindow)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return nil
case <-ticker.C:
b.flush(ctx)
}
}
}
// Add adds streams to the pending batch and returns the last known rates for them.
// This is a non-blocking operation. Streams with unknown rates will have rate=0.
func (b *rateBatcher) Add(tenant string, streams []SegmentedStream) map[uint64]uint64 {
if len(streams) == 0 {
return nil
}
// Get known rates first (read lock).
rates := make(map[uint64]uint64, len(streams))
b.ratesMu.RLock()
tenantRates := b.rates[tenant]
for _, stream := range streams {
if tenantRates != nil {
rates[stream.SegmentationKeyHash] = tenantRates[stream.SegmentationKeyHash]
} else {
rates[stream.SegmentationKeyHash] = 0
}
}
b.ratesMu.RUnlock()
// Add to pending batch (separate lock).
b.pendingMu.Lock()
tenantPending, ok := b.pending[tenant]
if !ok {
tenantPending = make(map[uint64]*proto.StreamMetadata)
b.pending[tenant] = tenantPending
}
for _, stream := range streams {
hash := stream.SegmentationKeyHash
totalSize := uint64(stream.Stream.Size())
// If we already have this stream in the pending batch, accumulate the size.
if existing, ok := tenantPending[hash]; ok {
existing.TotalSize += totalSize
} else {
tenantPending[hash] = &proto.StreamMetadata{
StreamHash: hash,
TotalSize: totalSize,
IngestionPolicy: stream.Policy,
}
}
}
// Update pending streams gauge.
var total int
for _, t := range b.pending {
total += len(t)
}
b.pendingStreams.Set(float64(total))
b.pendingMu.Unlock()
return rates
}
// flush sends all pending streams to the frontend.
func (b *rateBatcher) flush(ctx context.Context) {
// Swap out the pending map so we don't hold the lock during the RPC.
b.pendingMu.Lock()
toFlush := b.pending
b.pending = make(map[string]map[uint64]*proto.StreamMetadata)
b.pendingMu.Unlock()
b.pendingStreams.Set(0)
if len(toFlush) == 0 {
return
}
// Count total streams for metrics.
var totalStreams int
for _, streams := range toFlush {
totalStreams += len(streams)
}
b.streamsPerBatch.Observe(float64(totalStreams))
// Send each tenant's streams to the frontend.
for tenant, streams := range toFlush {
if len(streams) == 0 {
continue
}
// Convert map to slice for the request.
metadata := make([]*proto.StreamMetadata, 0, len(streams))
for _, m := range streams {
metadata = append(metadata, m)
}
req := &proto.UpdateRatesRequest{
Tenant: tenant,
Streams: metadata,
}
// Inject tenant ID into context for the RPC.
tenantCtx := user.InjectOrgID(ctx, tenant)
results, err := b.client.UpdateRatesRaw(tenantCtx, req)
if err != nil {
level.Error(b.logger).Log(
"msg", "failed to flush rate batch",
"tenant", tenant,
"streams", len(metadata),
"err", err,
)
b.batchesFailed.Inc()
continue
}
b.batchesSent.Inc()
b.streamsFlushed.Add(float64(len(metadata)))
// Store the rates for future lookups.
b.storeRates(tenant, results)
}
}
// storeRates replaces the rates for a tenant with the results from the latest flush.
// This ensures we don't accumulate stale rates for streams that stopped sending.
func (b *rateBatcher) storeRates(tenant string, results []*proto.UpdateRatesResult) {
b.ratesMu.Lock()
defer b.ratesMu.Unlock()
if len(results) == 0 {
delete(b.rates, tenant)
return
}
tenantRates := make(map[uint64]uint64, len(results))
for _, result := range results {
tenantRates[result.StreamHash] = result.Rate
}
b.rates[tenant] = tenantRates
}
// GetRate returns the last known rate for a stream, or 0 if unknown.
func (b *rateBatcher) GetRate(tenant string, segmentationKeyHash uint64) uint64 {
b.ratesMu.RLock()
defer b.ratesMu.RUnlock()
if tenantRates, ok := b.rates[tenant]; ok {
return tenantRates[segmentationKeyHash]
}
return 0
}
pkg/distributor/rate_batcher_test.go
+356
View File
@@ -0,0 +1,356 @@
package distributor
import (
"context"
"sync"
"testing"
"time"
"github.com/go-kit/log"
"github.com/grafana/dskit/services"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/require"
"github.com/grafana/loki/v3/pkg/limits/proto"
"github.com/grafana/loki/v3/pkg/logproto"
)
type mockUpdateRatesClient struct {
mu sync.Mutex
calls int
requests []*proto.UpdateRatesRequest
customRate uint64 // If non-zero, return this rate instead of 1000.
}
func (m *mockUpdateRatesClient) UpdateRatesRaw(_ context.Context, req *proto.UpdateRatesRequest) ([]*proto.UpdateRatesResult, error) {
m.mu.Lock()
defer m.mu.Unlock()
m.calls++
m.requests = append(m.requests, req)
// Return customRate if set, otherwise 1000.
rate := m.customRate
if rate == 0 {
rate = 1000
}
results := make([]*proto.UpdateRatesResult, len(req.Streams))
for i, stream := range req.Streams {
results[i] = &proto.UpdateRatesResult{
StreamHash: stream.StreamHash,
Rate: rate,
}
}
return results, nil
}
func TestRateBatcher_Add_AccumulatesStreams(t *testing.T) {
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: time.Hour, // Long window so we control when flush happens
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// Add some streams.
streams := []SegmentedStream{
{
KeyedStream: KeyedStream{
Stream: logproto.Stream{
Labels: `{app="test"}`,
Entries: []logproto.Entry{{Timestamp: time.Now(), Line: "test"}},
},
Policy: "default",
},
SegmentationKeyHash: 123,
},
{
KeyedStream: KeyedStream{
Stream: logproto.Stream{
Labels: `{app="test2"}`,
Entries: []logproto.Entry{{Timestamp: time.Now(), Line: "test2"}},
},
Policy: "default",
},
SegmentationKeyHash: 456,
},
}
batcher.Add("tenant1", streams)
batcher.Add("tenant1", streams) // Add same streams again - should accumulate size
// No calls yet - batching.
require.Equal(t, 0, client.calls)
// Manually flush.
batcher.flush(context.Background())
// Should have made one call.
require.Equal(t, 1, client.calls)
require.Len(t, client.requests, 1)
require.Equal(t, "tenant1", client.requests[0].Tenant)
require.Len(t, client.requests[0].Streams, 2) // 2 unique streams
}
func TestRateBatcher_AccumulatesSize(t *testing.T) {
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: time.Hour,
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// Add stream with some entries.
stream1 := []SegmentedStream{
{
KeyedStream: KeyedStream{
Stream: logproto.Stream{
Labels: `{app="test"}`,
Entries: []logproto.Entry{{Timestamp: time.Now(), Line: "hello"}},
},
},
SegmentationKeyHash: 123,
},
}
// Add stream again with more entries.
stream2 := []SegmentedStream{
{
KeyedStream: KeyedStream{
Stream: logproto.Stream{
Labels: `{app="test"}`,
Entries: []logproto.Entry{{Timestamp: time.Now(), Line: "world!"}},
},
},
SegmentationKeyHash: 123, // Same hash
},
}
batcher.Add("tenant1", stream1)
batcher.Add("tenant1", stream2)
batcher.flush(context.Background())
// Should have one stream with accumulated size.
require.Equal(t, 1, client.calls)
require.Len(t, client.requests[0].Streams, 1)
// Size should be accumulated: first Add size + second Add size.
// The batcher uses stream.Stream.Size() (protobuf size).
expectedTotal := uint64(stream1[0].Stream.Size()) + uint64(stream2[0].Stream.Size())
require.Equal(t, expectedTotal, client.requests[0].Streams[0].TotalSize)
}
func TestRateBatcher_MultipleTenants(t *testing.T) {
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: time.Hour,
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// Add streams for multiple tenants.
batcher.Add("tenant1", []SegmentedStream{{SegmentationKeyHash: 123}})
batcher.Add("tenant2", []SegmentedStream{{SegmentationKeyHash: 456}})
batcher.flush(context.Background())
// Should have made 2 calls (one per tenant).
require.Equal(t, 2, client.calls)
// Find requests by tenant.
tenantRequests := make(map[string]*proto.UpdateRatesRequest)
for _, req := range client.requests {
tenantRequests[req.Tenant] = req
}
require.Contains(t, tenantRequests, "tenant1")
require.Contains(t, tenantRequests, "tenant2")
require.Len(t, tenantRequests["tenant1"].Streams, 1)
require.Len(t, tenantRequests["tenant2"].Streams, 1)
}
func TestRateBatcher_EmptyFlush(t *testing.T) {
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: time.Hour,
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// Flush with nothing pending.
batcher.flush(context.Background())
// Should not have made any calls.
require.Equal(t, 0, client.calls)
}
func TestRateBatcher_ServiceLifecycle(t *testing.T) {
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: 50 * time.Millisecond,
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// Start the service.
ctx := context.Background()
require.NoError(t, services.StartAndAwaitRunning(ctx, batcher))
// Add some streams.
batcher.Add("tenant1", []SegmentedStream{{SegmentationKeyHash: 123}})
// Wait for automatic flush.
time.Sleep(100 * time.Millisecond)
// Should have flushed automatically.
client.mu.Lock()
calls := client.calls
client.mu.Unlock()
require.GreaterOrEqual(t, calls, 1)
// Stop the service.
require.NoError(t, services.StopAndAwaitTerminated(ctx, batcher))
}
func TestRateBatcher_StoresRatesFromFlush(t *testing.T) {
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: time.Hour,
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// Initially, rates are unknown (0).
rate := batcher.GetRate("tenant1", 123)
require.Equal(t, uint64(0), rate)
// Add streams and flush.
batcher.Add("tenant1", []SegmentedStream{{SegmentationKeyHash: 123}})
batcher.Add("tenant1", []SegmentedStream{{SegmentationKeyHash: 456}})
batcher.flush(context.Background())
// After flush, rates should be stored (mock returns 1000 for all).
rate = batcher.GetRate("tenant1", 123)
require.Equal(t, uint64(1000), rate)
rate = batcher.GetRate("tenant1", 456)
require.Equal(t, uint64(1000), rate)
// Unknown stream still returns 0.
rate = batcher.GetRate("tenant1", 789)
require.Equal(t, uint64(0), rate)
// Different tenant returns 0.
rate = batcher.GetRate("tenant2", 123)
require.Equal(t, uint64(0), rate)
}
func TestRateBatcher_AddReturnsRates(t *testing.T) {
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: time.Hour,
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// First Add returns 0 for all (no prior flush).
rates := batcher.Add("tenant1", []SegmentedStream{
{SegmentationKeyHash: 123},
{SegmentationKeyHash: 456},
})
require.Equal(t, uint64(0), rates[123])
require.Equal(t, uint64(0), rates[456])
// Flush to populate rates.
batcher.flush(context.Background())
// Second Add returns last known rates.
rates = batcher.Add("tenant1", []SegmentedStream{
{SegmentationKeyHash: 123},
{SegmentationKeyHash: 456},
{SegmentationKeyHash: 789}, // New stream, unknown rate.
})
require.Equal(t, uint64(1000), rates[123])
require.Equal(t, uint64(1000), rates[456])
require.Equal(t, uint64(0), rates[789]) // Unknown stream.
}
func TestRateBatcher_RatesUpdatedOnSubsequentFlush(t *testing.T) {
// Custom client that returns different rates based on call count.
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: time.Hour,
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// First flush.
batcher.Add("tenant1", []SegmentedStream{{SegmentationKeyHash: 123}})
batcher.flush(context.Background())
require.Equal(t, uint64(1000), batcher.GetRate("tenant1", 123))
// Modify mock to return different rate.
client.mu.Lock()
client.customRate = 5000
client.mu.Unlock()
// Second flush updates the rate.
batcher.Add("tenant1", []SegmentedStream{{SegmentationKeyHash: 123}})
batcher.flush(context.Background())
require.Equal(t, uint64(5000), batcher.GetRate("tenant1", 123))
}
func TestRateBatcher_RatesClearedForInactiveStreams(t *testing.T) {
client := &mockUpdateRatesClient{}
batcher := newRateBatcher(
RateBatcherConfig{
BatchWindow: time.Hour,
},
client,
log.NewNopLogger(),
prometheus.NewRegistry(),
)
// First flush with streams 123 and 456.
batcher.Add("tenant1", []SegmentedStream{
{SegmentationKeyHash: 123},
{SegmentationKeyHash: 456},
})
batcher.flush(context.Background())
require.Equal(t, uint64(1000), batcher.GetRate("tenant1", 123))
require.Equal(t, uint64(1000), batcher.GetRate("tenant1", 456))
// Second flush with only stream 123 (456 became inactive).
batcher.Add("tenant1", []SegmentedStream{{SegmentationKeyHash: 123}})
batcher.flush(context.Background())
// Stream 123 still has a rate.
require.Equal(t, uint64(1000), batcher.GetRate("tenant1", 123))
// Stream 456 was cleared (not in last batch).
require.Equal(t, uint64(0), batcher.GetRate("tenant1", 456))
}