feat: refactor frontend cache, add metrics, filter out cached streams (#20860)

pkg/limits/frontend/client.go

@@ -9,6 +9,8 @@ import (
 	"time"
 
 	"github.com/bits-and-blooms/bloom/v3"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 
 	"github.com/grafana/loki/v3/pkg/limits/proto"
 )
@@ -22,61 +24,59 @@ type limitsClient interface {
 	UpdateRates(context.Context, *proto.UpdateRatesRequest) ([]*proto.UpdateRatesResponse, error)
 }
 
-// A cacheLimitsClient uses a cache to reduce the load on limits backends.
+// A cacheLimitsClient uses caches to reduce the load on limits backends.
 type cacheLimitsClient struct {
-	ttl    time.Duration
-	onMiss limitsClient
-	// The fields below MUST NOT be used without mtx.
-	mtx          sync.RWMutex
-	knownStreams *bloom.BloomFilter
-	lastExpired  time.Time
+	acceptedStreamsCache *acceptedStreamsCache
+	onMiss               limitsClient
 }
 
 // newCacheLimitsClient returns a new cache limits client.
-func newCacheLimitsClient(
-	ttl, maxJitter time.Duration,
-	knownStreams *bloom.BloomFilter,
-	onMiss limitsClient,
-) *cacheLimitsClient {
+func newCacheLimitsClient(acceptedStreamsCache *acceptedStreamsCache, onMiss limitsClient) *cacheLimitsClient {
 	return &cacheLimitsClient{
-		ttl:          ttl,
-		knownStreams: knownStreams,
-		lastExpired:  time.Now().Add(randDuration(maxJitter)),
-		onMiss:       onMiss,
+		acceptedStreamsCache: acceptedStreamsCache,
+		onMiss:               onMiss,
 	}
 }
 
 // ExceedsLimits implements the [limitsClient] interface.
 func (c *cacheLimitsClient) ExceedsLimits(ctx context.Context, req *proto.ExceedsLimitsRequest) ([]*proto.ExceedsLimitsResponse, error) {
-	c.expireTTL()
-	// If the exact same request has been seen before, and all streams were
-	// accepted, we can assume it will continue to be accepted.
-	if c.hasKnownStreams(req) {
+	c.acceptedStreamsCache.ExpireTTL()
+	// Remove streams that have been accepted from the request. This means
+	// we just check streams we haven't seen before, which reduces the
+	// number of requests we need to make to the limits backends.
+	c.acceptedStreamsCache.FilterInPlace(req)
+	if len(req.Streams) == 0 {
 		return []*proto.ExceedsLimitsResponse{}, nil
 	}
-	// Need to check with the limits service.
+	// Need to check remaining streams with the limits service.
 	resps, err := c.onMiss.ExceedsLimits(ctx, req)
 	if err != nil {
-		return resps, err
+		return nil, err
 	}
-	// We do not cache rejected streams at this time, so rejections must be
-	// filtered out before updating the cache.
+	numRejected := 0
+	for _, resp := range resps {
+		numRejected += len(resp.Results)
+	}
+	// Fast path, all streams rejected.
+	if numRejected == len(req.Streams) {
+		return resps, nil
+	}
+	// There are some accepted streams we haven't seen before, so add them
+	// to the cache. We do not cache rejected streams at this time, so
+	// rejections must be filtered out before updating the cache.
 	rejected := make(map[uint64]struct{})
 	for _, resp := range resps {
 		for _, res := range resp.Results {
 			rejected[res.StreamHash] = struct{}{}
 		}
 	}
-	c.mtx.Lock()
-	defer c.mtx.Unlock()
+	accepted := make([]*proto.StreamMetadata, 0, len(req.Streams))
 	for _, s := range req.Streams {
-		// If the stream was not rejected, add it to the cache.
 		if _, ok := rejected[s.StreamHash]; !ok {
-			b := bytes.Buffer{}
-			encodeStreamToBuf(&b, req.Tenant, s)
-			c.knownStreams.Add(b.Bytes())
+			accepted = append(accepted, s)
 		}
 	}
+	c.acceptedStreamsCache.Update(req.Tenant, accepted)
 	return resps, nil
 }
 
@@ -85,8 +85,45 @@ func (c *cacheLimitsClient) UpdateRates(ctx context.Context, req *proto.UpdateRa
 	return c.onMiss.UpdateRates(ctx, req)
 }
 
-// expireTTL expires the caches if the TTL has been exceeded.
-func (c *cacheLimitsClient) expireTTL() {
+type acceptedStreamsCache struct {
+	ttl time.Duration
+
+	// The fields below MUST NOT be used without mtx.
+	mtx         sync.RWMutex
+	bf          *bloom.BloomFilter
+	lastExpired time.Time
+
+	// Metrics.
+	cacheSize         prometheus.Gauge
+	cacheSizeEstimate prometheus.GaugeFunc
+}
+
+func newAcceptedStreamsCache(ttl, maxJitter time.Duration, cacheSize int, r prometheus.Registerer) *acceptedStreamsCache {
+	c := &acceptedStreamsCache{
+		ttl:         ttl,
+		bf:          bloom.NewWithEstimates(uint(cacheSize), 0.01),
+		lastExpired: time.Now().Add(randDuration(maxJitter)),
+	}
+	c.cacheSize = promauto.With(r).NewGauge(prometheus.GaugeOpts{
+		Name: "loki_ingest_limits_frontend_accepted_streams_cache_size",
+		Help: "Max size of the accepted streams cache.",
+	})
+	// This is static.
+	c.cacheSize.Set(float64(cacheSize))
+	// This is quite an expensive metric to compute so we moved it to a func.
+	c.cacheSizeEstimate = promauto.With(r).NewGaugeFunc(prometheus.GaugeOpts{
+		Name: "loki_ingest_limits_frontend_accepted_streams_cache_size_estimate",
+		Help: "Estimate size of the accepted streams cache.",
+	}, func() float64 {
+		c.mtx.RLock()
+		defer c.mtx.RUnlock()
+		return float64(c.bf.ApproximatedSize())
+	})
+	return c
+}
+
+// ExpireTTL expires the caches if the TTL has been exceeded.
+func (c *acceptedStreamsCache) ExpireTTL() {
 	// Fast path, first check the TTL with a read lock.
 	c.mtx.RLock()
 	lastExpired := c.lastExpired
@@ -100,25 +137,39 @@ func (c *cacheLimitsClient) expireTTL() {
 	c.mtx.Lock()
 	defer c.mtx.Unlock()
 	if time.Since(c.lastExpired) > c.ttl {
-		c.knownStreams.ClearAll()
+		c.bf.ClearAll()
 		c.lastExpired = time.Now()
 	}
 }
 
-// hasKnownStreams returns true if all streams in req are known streams.
-func (c *cacheLimitsClient) hasKnownStreams(req *proto.ExceedsLimitsRequest) bool {
+// FilterInPlace removes streams that are present in the cache.
+func (c *acceptedStreamsCache) FilterInPlace(req *proto.ExceedsLimitsRequest) {
 	// b is re-used. The data built from it MUST NOT escape this function.
 	b := bytes.Buffer{}
+	// See https://go.dev/wiki/SliceTricks.
+	filtered := req.Streams[:0]
 	c.mtx.RLock()
 	defer c.mtx.RUnlock()
 	for _, s := range req.Streams {
 		b.Reset()
 		encodeStreamToBuf(&b, req.Tenant, s)
-		if !c.knownStreams.Test(b.Bytes()) {
-			return false
+		if !c.bf.Test(b.Bytes()) {
+			filtered = append(filtered, s)
 		}
 	}
-	return true
+	req.Streams = filtered
+}
+
+func (c *acceptedStreamsCache) Update(tenant string, streams []*proto.StreamMetadata) {
+	// b is re-used. The data built from it MUST NOT escape this function.
+	b := bytes.Buffer{}
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	for _, s := range streams {
+		b.Reset()
+		encodeStreamToBuf(&b, tenant, s)
+		c.bf.Add(b.Bytes())
+	}
 }
 
 // randDuration returns a random duration between [0, d].
@@ -129,6 +180,6 @@ func randDuration(d time.Duration) time.Duration {
 // encodeStreamToBuf encodes the stream to the buffer.
 func encodeStreamToBuf(b *bytes.Buffer, tenant string, s *proto.StreamMetadata) {
 	b.Write([]byte(tenant))
-	// [bytes.Buffer] never return an error, it will panic instead.
+	// [bytes.Buffer] never returns an error, it will panic instead.
 	_ = binary.Write(b, binary.LittleEndian, s.StreamHash)
 }

pkg/limits/frontend/client_test.go

@@ -2,12 +2,11 @@ package frontend
 
 import (
 	"bytes"
-	"encoding/binary"
 	"testing"
 	"testing/synctest"
 	"time"
 
-	"github.com/bits-and-blooms/bloom/v3"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/loki/v3/pkg/limits"
@@ -15,47 +14,18 @@ import (
 )
 
 func TestCacheLimitsClient(t *testing.T) {
-	t.Run("streams accepted", func(t *testing.T) {
-		// When a stream is accepted, it should be inserted into known streams.
+	t.Run("accepted streams cached, rejected streams returned", func(t *testing.T) {
+		// When a stream is accepted, it should be inserted into the cache.
 		// We will assert this behavior later.
-		knownStreams := bloom.NewWithEstimates(10, 0.01)
+		cache := newAcceptedStreamsCache(time.Minute, 15*time.Second, 10, prometheus.NewRegistry())
 		onMiss := &mockLimitsClient{
 			t: t,
-			// Expect one stream 0x1 from the tenant "test".
+			// Expect two stream 0x1 and 0x2 from the tenant "test".
 			expectedExceedsLimitsRequest: &proto.ExceedsLimitsRequest{
 				Tenant:  "test",
-				Streams: []*proto.StreamMetadata{{StreamHash: 0x1}},
+				Streams: []*proto.StreamMetadata{{StreamHash: 0x1}, {StreamHash: 0x2}},
 			},
-			// All streams accepted.
-			exceedsLimitsResponses: []*proto.ExceedsLimitsResponse{},
-		}
-		client := newCacheLimitsClient(time.Minute, 15*time.Second, knownStreams, onMiss)
-		resps, err := client.ExceedsLimits(t.Context(), &proto.ExceedsLimitsRequest{
-			Tenant:  "test",
-			Streams: []*proto.StreamMetadata{{StreamHash: 0x1}},
-		})
-		// No streams should be rejected.
-		require.NoError(t, err)
-		require.Len(t, resps, 0)
-		// The cache should contain the stream 0x1 for the tenant "test".
-		// We don't use [encodeStreamToBuf] so we can test it.
-		b := bytes.Buffer{}
-		b.Write([]byte("test"))
-		_ = binary.Write(&b, binary.LittleEndian, uint64(1))
-		require.True(t, knownStreams.Test(b.Bytes()))
-	})
-
-	t.Run("streams rejected", func(t *testing.T) {
-		// When a stream is rejected, it should not be cached.
-		knownStreams := bloom.NewWithEstimates(10, 0.01)
-		onMiss := &mockLimitsClient{
-			t: t,
-			// Expect one stream 0x1 from the tenant "test".
-			expectedExceedsLimitsRequest: &proto.ExceedsLimitsRequest{
-				Tenant:  "test",
-				Streams: []*proto.StreamMetadata{{StreamHash: 0x1}},
-			},
-			// The stream should be rejected.
+			// Reject stream 0x1.
 			exceedsLimitsResponses: []*proto.ExceedsLimitsResponse{{
 				Results: []*proto.ExceedsLimitsResult{{
 					StreamHash: 0x1,
@@ -63,45 +33,118 @@ func TestCacheLimitsClient(t *testing.T) {
 				}},
 			}},
 		}
-		client := newCacheLimitsClient(time.Minute, 15*time.Second, knownStreams, onMiss)
+		client := newCacheLimitsClient(cache, onMiss)
 		resps, err := client.ExceedsLimits(t.Context(), &proto.ExceedsLimitsRequest{
 			Tenant:  "test",
-			Streams: []*proto.StreamMetadata{{StreamHash: 0x1}},
+			Streams: []*proto.StreamMetadata{{StreamHash: 0x1}, {StreamHash: 0x2}},
 		})
+		// The stream 0x1 should have been rejected, and should be absent
+		// from the cache.
 		require.NoError(t, err)
 		require.Len(t, resps, 1)
-		// No bits should have been set.
-		require.Equal(t, uint(0), knownStreams.BitSet().Count())
+		require.Len(t, resps[0].Results, 1)
+		require.Equal(t, uint64(0x1), resps[0].Results[0].StreamHash)
+		b := bytes.Buffer{}
+		encodeStreamToBuf(&b, "test", &proto.StreamMetadata{StreamHash: 0x1})
+		require.False(t, cache.bf.Test(b.Bytes()))
+		// The cache should contain the stream 0x2 for the tenant "test".
+		b.Reset()
+		encodeStreamToBuf(&b, "test", &proto.StreamMetadata{StreamHash: 0x2})
+		require.True(t, cache.bf.Test(b.Bytes()))
 	})
 
-	t.Run("cache is expired after TTL", func(t *testing.T) {
+	t.Run("cached streams are returned", func(t *testing.T) {
+		cache := newAcceptedStreamsCache(time.Minute, 15*time.Second, 10, prometheus.NewRegistry())
+		cache.Update("test", []*proto.StreamMetadata{{StreamHash: 0x1}})
+		onMiss := &mockLimitsClient{
+			t: t,
+			// Expect one stream 0x2 because 0x1 is cached.
+			expectedExceedsLimitsRequest: &proto.ExceedsLimitsRequest{
+				Tenant:  "test",
+				Streams: []*proto.StreamMetadata{{StreamHash: 0x2}},
+			},
+			exceedsLimitsResponses: []*proto.ExceedsLimitsResponse{},
+		}
+		client := newCacheLimitsClient(cache, onMiss)
+		resps, err := client.ExceedsLimits(t.Context(), &proto.ExceedsLimitsRequest{
+			Tenant:  "test",
+			Streams: []*proto.StreamMetadata{{StreamHash: 0x1}, {StreamHash: 0x2}},
+		})
+		require.NoError(t, err)
+		require.Len(t, resps, 0)
+	})
+}
+
+func TestAcceptedStreamsCache(t *testing.T) {
+	t.Run("cache is cleared after TTL elapsed", func(t *testing.T) {
 		synctest.Test(t, func(t *testing.T) {
-			knownStreams := bloom.NewWithEstimates(10, 0.01)
-			client := newCacheLimitsClient(time.Minute, 15*time.Second, knownStreams, &mockLimitsClient{})
+			c := newAcceptedStreamsCache(time.Minute, 15*time.Second, 10, prometheus.NewRegistry())
 			// Remove jitter for tests.
-			client.lastExpired = time.Now()
+			c.lastExpired = time.Now()
 
 			now := time.Now()
-			require.Equal(t, now, client.lastExpired)
+			require.Equal(t, now, c.lastExpired)
 			// No bits should have been set.
-			require.Equal(t, uint(0), knownStreams.BitSet().Count())
+			require.Equal(t, uint(0), c.bf.BitSet().Count())
 
 			// Advance the clock, no reset should happen.
 			time.Sleep(time.Second)
-			client.expireTTL()
-			require.Equal(t, now, client.lastExpired)
+			c.ExpireTTL()
+			require.Equal(t, now, c.lastExpired)
 
 			// Add some data to the cache.
-			knownStreams.Add([]byte("test"))
-			require.Greater(t, knownStreams.BitSet().Count(), uint(0))
+			c.bf.Add([]byte("test"))
+			require.Greater(t, c.bf.BitSet().Count(), uint(0))
 
 			// Advance the clock past the TTL (include the jitter).
 			time.Sleep(time.Minute + (5 * time.Second))
 			now = time.Now()
-			client.expireTTL()
-			require.Equal(t, now, client.lastExpired)
+			c.ExpireTTL()
+			require.Equal(t, now, c.lastExpired)
 			// The bits should have been reset.
-			require.Equal(t, uint(0), knownStreams.BitSet().Count())
+			require.Equal(t, uint(0), c.bf.BitSet().Count())
 		})
 	})
+
+	t.Run("request is filtered in place, accepted streams are removed", func(t *testing.T) {
+		c := newAcceptedStreamsCache(time.Minute, 15*time.Second, 10, prometheus.NewRegistry())
+		// Create a stream and add it to the cache.
+		s1 := &proto.StreamMetadata{StreamHash: 0x1}
+		c.Update("test", []*proto.StreamMetadata{s1})
+		// Create a second stream but do not add it to the cache.
+		s2 := &proto.StreamMetadata{StreamHash: 0x2}
+		req := &proto.ExceedsLimitsRequest{
+			Tenant:  "test",
+			Streams: []*proto.StreamMetadata{s1, s2},
+		}
+		c.FilterInPlace(req)
+		require.Len(t, req.Streams, 1)
+		require.Equal(t, uint64(0x2), req.Streams[0].StreamHash)
+	})
+
+	t.Run("cache contains streams", func(t *testing.T) {
+		c := newAcceptedStreamsCache(time.Minute, 15*time.Second, 10, prometheus.NewRegistry())
+		// Create a stream, add it to the cache, and then check that it is
+		// present.
+		s1 := &proto.StreamMetadata{StreamHash: 0x1}
+		c.Update("test", []*proto.StreamMetadata{s1})
+		b := bytes.Buffer{}
+		encodeStreamToBuf(&b, "test", s1)
+		require.True(t, c.bf.Test(b.Bytes()))
+		// Create a second stream without adding it to the cache, and then
+		// check that it is absent.
+		s2 := &proto.StreamMetadata{StreamHash: 0x2}
+		b.Reset()
+		encodeStreamToBuf(&b, "test", s2)
+		require.False(t, c.bf.Test(b.Bytes()))
+		// Add the second stream to the cache, and then check both streams
+		// are present.
+		c.Update("test", []*proto.StreamMetadata{s2})
+		b.Reset()
+		encodeStreamToBuf(&b, "test", s1)
+		require.True(t, c.bf.Test(b.Bytes()))
+		b.Reset()
+		encodeStreamToBuf(&b, "test", s2)
+		require.True(t, c.bf.Test(b.Bytes()))
+	})
 }

pkg/limits/frontend/frontend.go

@@ -9,7 +9,6 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/bits-and-blooms/bloom/v3"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/ring"
@@ -86,7 +85,7 @@ func New(cfg Config, ringName string, limitsRing ring.ReadRing, logger log.Logge
 	// Set up the limits client.
 	f.limitsClient = newRingLimitsClient(limitsRing, clientPool, cfg.NumPartitions, f.assignedPartitionsCache, logger, reg)
 	if cfg.CacheTTL > 0 {
-		f.limitsClient = newCacheLimitsClient(cfg.CacheTTL, cfg.CacheTTLJitter, bloom.NewWithEstimates(1000000, 0.01), f.limitsClient)
+		f.limitsClient = newCacheLimitsClient(newAcceptedStreamsCache(cfg.CacheTTL, cfg.CacheTTLJitter, 1000000, reg), f.limitsClient)
 	}
 	lifecycler, err := ring.NewLifecycler(cfg.LifecyclerConfig, f, RingName, RingKey, true, logger, reg)
 	if err != nil {