feat(ingest-limits): Add partition id caching (#16964)

Author: periklis

docs/sources/shared/configuration.md

@@ -1209,6 +1209,10 @@ ingest_limits_frontend:
   # CLI flag: -ingest-limits-frontend.num-partitions
   [num_partitions: <int> | default = 64]
 
+  # The TTL for the stream usage cache.
+  # CLI flag: -ingest-limits-frontend.partition-id-cache-ttl
+  [partition_id_cache_ttl: <duration> | default = 1m]
+
 ingest_limits_frontend_client:
   # Configures client gRPC connections to limits service.
   # The CLI flags prefix for this block configuration is:

go.mod

@@ -132,6 +132,7 @@ require (
 	github.com/grafana/loki/pkg/push v0.0.0-20240924133635-758364c7775f
 	github.com/heroku/x v0.4.3
 	github.com/influxdata/tdigest v0.0.2-0.20210216194612-fc98d27c9e8b
+	github.com/jellydator/ttlcache/v3 v3.3.0
 	github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db
 	github.com/ncw/swift/v2 v2.0.3
 	github.com/parquet-go/parquet-go v0.25.0

go.sum

@@ -776,6 +776,8 @@ github.com/jcmturner/gokrb5/v8 v8.4.4 h1:x1Sv4HaTpepFkXbt2IkL29DXRf8sOfZXo8eRKh6
 github.com/jcmturner/gokrb5/v8 v8.4.4/go.mod h1:1btQEpgT6k+unzCwX1KdWMEwPPkkgBtP+F6aCACiMrs=
 github.com/jcmturner/rpc/v2 v2.0.3 h1:7FXXj8Ti1IaVFpSAziCZWNzbNuZmnvw/i6CqLNdWfZY=
 github.com/jcmturner/rpc/v2 v2.0.3/go.mod h1:VUJYCIDm3PVOEHw8sgt091/20OJjskO/YJki3ELg/Hc=
+github.com/jellydator/ttlcache/v3 v3.3.0 h1:BdoC9cE81qXfrxeb9eoJi9dWrdhSuwXMAnHTbnBm4Wc=
+github.com/jellydator/ttlcache/v3 v3.3.0/go.mod h1:bj2/e0l4jRnQdrnSTaGTsh4GSXvMjQcy41i7th0GVGw=
 github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af/go.mod h1:Nht3zPeWKUH0NzdCt2Blrr5ys8VGpn0CEB0cQHVjt7k=
 github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
 github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=

pkg/limits/frontend/cache.go

@@ -0,0 +1,18 @@
+package frontend
+
+import (
+	"time"
+
+	"github.com/jellydator/ttlcache/v3"
+
+	"github.com/grafana/loki/v3/pkg/logproto"
+)
+
+type PartitionConsumersCache = *ttlcache.Cache[string, logproto.GetAssignedPartitionsResponse]
+
+func NewPartitionConsumerCache(ttl time.Duration) PartitionConsumersCache {
+	return ttlcache.New(
+		ttlcache.WithTTL[string, logproto.GetAssignedPartitionsResponse](ttl),
+		ttlcache.WithDisableTouchOnHit[string, logproto.GetAssignedPartitionsResponse](),
+	)
+}

pkg/limits/frontend/config.go

@@ -13,17 +13,19 @@ import (
 
 // Config contains the config for an ingest-limits-frontend.
 type Config struct {
-	ClientConfig     limits_client.Config  `yaml:"client_config"`
-	LifecyclerConfig ring.LifecyclerConfig `yaml:"lifecycler,omitempty"`
-	RecheckPeriod    time.Duration         `yaml:"recheck_period"`
-	NumPartitions    int                   `yaml:"num_partitions"`
+	ClientConfig        limits_client.Config  `yaml:"client_config"`
+	LifecyclerConfig    ring.LifecyclerConfig `yaml:"lifecycler,omitempty"`
+	RecheckPeriod       time.Duration         `yaml:"recheck_period"`
+	NumPartitions       int                   `yaml:"num_partitions"`
+	PartitionIDCacheTTL time.Duration         `yaml:"partition_id_cache_ttl"`
 }
 
 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.")
 	f.IntVar(&cfg.NumPartitions, "ingest-limits-frontend.num-partitions", 64, "The number of partitions to use for the ring.")
+	f.DurationVar(&cfg.PartitionIDCacheTTL, "ingest-limits-frontend.partition-id-cache-ttl", 1*time.Minute, "The TTL for the stream usage cache.")
 }
 
 func (cfg *Config) Validate() error {

pkg/limits/frontend/frontend.go

@@ -58,10 +58,11 @@ type Frontend struct {
 	cfg    Config
 	logger log.Logger
 
-	limits      Limits
-	rateLimiter *limiter.RateLimiter
-	streamUsage StreamUsageGatherer
-	metrics     *metrics
+	limits           Limits
+	rateLimiter      *limiter.RateLimiter
+	streamUsage      StreamUsageGatherer
+	partitionIDCache PartitionConsumersCache
+	metrics          *metrics
 
 	subservices        *services.Manager
 	subservicesWatcher *services.FailureWatcher
@@ -84,15 +85,17 @@ func New(cfg Config, ringName string, limitsRing ring.ReadRing, limits Limits, l
 	)
 
 	rateLimiter := limiter.NewRateLimiter(newRateLimitsAdapter(limits), cfg.RecheckPeriod)
-	streamUsage := NewRingStreamUsageGatherer(limitsRing, clientPool, logger, cfg.NumPartitions)
+	partitionIDCache := NewPartitionConsumerCache(cfg.PartitionIDCacheTTL)
+	streamUsage := NewRingStreamUsageGatherer(limitsRing, clientPool, logger, partitionIDCache, cfg.NumPartitions)
 
 	f := &Frontend{
-		cfg:         cfg,
-		logger:      logger,
-		limits:      limits,
-		rateLimiter: rateLimiter,
-		streamUsage: streamUsage,
-		metrics:     newMetrics(reg),
+		cfg:              cfg,
+		logger:           logger,
+		limits:           limits,
+		rateLimiter:      rateLimiter,
+		streamUsage:      streamUsage,
+		partitionIDCache: partitionIDCache,
+		metrics:          newMetrics(reg),
 	}
 
 	lifecycler, err := ring.NewLifecycler(cfg.LifecyclerConfig, f, RingName, RingKey, true, logger, reg)
@@ -135,6 +138,8 @@ func (f *Frontend) starting(ctx context.Context) (err error) {
 		return fmt.Errorf("failed to start subservices: %w", err)
 	}
 
+	go f.partitionIDCache.Start()
+
 	return nil
 }
 
@@ -150,6 +155,7 @@ func (f *Frontend) running(ctx context.Context) error {
 
 // stopping implements services.Service.
 func (f *Frontend) stopping(_ error) error {
+	f.partitionIDCache.Stop()
 	return services.StopManagerAndAwaitStopped(context.Background(), f.subservices)
 }
 

pkg/limits/frontend/frontend_test.go

@@ -305,11 +305,12 @@ func TestFrontend_ExceedsLimits(t *testing.T) {
 				ingestionRate:    test.ingestionRate,
 			}
 			rl := limiter.NewRateLimiter(newRateLimitsAdapter(l), 10*time.Second)
+			cache := NewPartitionConsumerCache(1 * time.Millisecond)
 
 			f := Frontend{
 				limits:      l,
 				rateLimiter: rl,
-				streamUsage: NewRingStreamUsageGatherer(readRing, clientPool, log.NewNopLogger(), 2),
+				streamUsage: NewRingStreamUsageGatherer(readRing, clientPool, log.NewNopLogger(), cache, 2),
 				metrics:     newMetrics(prometheus.NewRegistry()),
 			}
 

pkg/limits/frontend/http.go

@@ -3,6 +3,7 @@ package frontend
 import (
 	"encoding/json"
 	"net/http"
+	"text/template"
 
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/user"
@@ -11,6 +12,46 @@ import (
 	"github.com/grafana/loki/v3/pkg/util"
 )
 
+const ringStreamUsageTemplate = `
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Ring Stream Usage Cache</title>
+    <style>
+        table { border-collapse: collapse; width: 100%; }
+        th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
+        th { background-color: #f2f2f2; }
+        form { display: inline; }
+        button { padding: 5px 10px; margin: 2px; }
+    </style>
+</head>
+<body>
+    <h1>Ring Stream Usage Cache</h1>
+    <form method="POST">
+        <button type="submit">Clear All Cache</button>
+    </form>
+    <table>
+        <tr>
+            <th>Instance Address</th>
+            <th>Partitions</th>
+            <th>Actions</th>
+        </tr>
+        {{range $addr, $partitions := .Entries}}
+        <tr>
+            <td>{{$addr}}</td>
+            <td>{{range $i, $p := $partitions}}{{if $i}}, {{end}}{{$p}}{{end}}</td>
+            <td>
+                <form method="POST">
+                    <input type="hidden" name="instance" value="{{$addr}}">
+                    <button type="submit">Clear Cache</button>
+                </form>
+            </td>
+        </tr>
+        {{end}}
+    </table>
+</body>
+</html>`
+
 type httpExceedsLimitsRequest struct {
 	TenantID     string   `json:"tenantID"`
 	StreamHashes []uint64 `json:"streamHashes"`
@@ -61,3 +102,46 @@ func (f *Frontend) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		Results: resp.Results,
 	})
 }
+
+// PartitionConsumersCacheHandler handles the GET request to display the cache.
+func (f *Frontend) PartitionConsumersCacheHandler(w http.ResponseWriter, _ *http.Request) {
+	data := struct {
+		Entries map[string][]int32
+	}{
+		Entries: make(map[string][]int32),
+	}
+
+	for addr, entry := range f.partitionIDCache.Items() {
+		assignedPartitions := entry.Value().AssignedPartitions
+		for partition := range assignedPartitions {
+			data.Entries[addr] = append(data.Entries[addr], partition)
+		}
+	}
+
+	w.Header().Set("Content-Type", "text/html; charset=utf-8")
+	tmpl := template.Must(template.New("cache").Parse(ringStreamUsageTemplate))
+	if err := tmpl.Execute(w, data); err != nil {
+		http.Error(w, "Failed to render template", http.StatusInternalServerError)
+		return
+	}
+}
+
+// PartitionConsumersCacheEvictHandler handles the POST request to clear the cache.
+func (f *Frontend) PartitionConsumersCacheEvictHandler(w http.ResponseWriter, r *http.Request) {
+	if err := r.ParseForm(); err != nil {
+		http.Error(w, "Failed to parse form", http.StatusBadRequest)
+		return
+	}
+
+	instance := r.FormValue("instance")
+	if instance == "" {
+		// Clear all cache
+		f.partitionIDCache.DeleteAll()
+	} else {
+		// Clear specific instance
+		f.partitionIDCache.Delete(instance)
+	}
+
+	// Redirect back to the GET page
+	http.Redirect(w, r, r.URL.Path, http.StatusSeeOther)
+}