fix(ingest-limits): Use stripe locking for metadata (#17150)

Author: periklis

pkg/limits/http.go

@@ -19,11 +19,6 @@ type httpTenantLimitsResponse struct {
 // ServeHTTP implements the http.Handler interface.
 // It returns the current stream counts and status per tenant as a JSON response.
 func (s *IngestLimits) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	// TODO(grobinson): Avoid acquiring the mutex for the entire duration
-	// of the request.
-	s.mtx.RLock()
-	defer s.mtx.RUnlock()
-
 	tenant := mux.Vars(r)["tenant"]
 	if tenant == "" {
 		http.Error(w, "invalid tenant", http.StatusBadRequest)
@@ -43,20 +38,22 @@ func (s *IngestLimits) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		response      httpTenantLimitsResponse
 	)
 
-	for _, partitions := range s.metadata[tenant] {
-		for _, stream := range partitions {
-			if stream.lastSeenAt >= cutoff {
-				activeStreams++
+	s.metadata.All(func(tenantID string, _ int32, stream Stream) {
+		if tenantID != tenant {
+			return
+		}
+
+		if stream.LastSeenAt >= cutoff {
+			activeStreams++
 
-				// Calculate size only within the rate window
-				for _, bucket := range stream.rateBuckets {
-					if bucket.timestamp >= rateWindowCutoff {
-						totalSize += bucket.size
-					}
+			// Calculate size only within the rate window
+			for _, bucket := range stream.RateBuckets {
+				if bucket.Timestamp >= rateWindowCutoff {
+					totalSize += bucket.Size
 				}
 			}
 		}
-	}
+	})
 
 	// Calculate rate using only data from within the rate window
 	calculatedRate := float64(totalSize) / s.cfg.WindowSize.Seconds()

pkg/limits/http_test.go

@@ -20,20 +20,30 @@ func TestIngestLimits_ServeHTTP(t *testing.T) {
 			RateWindow:     time.Minute,
 			BucketDuration: 30 * time.Second,
 		},
-		metadata: map[string]map[int32][]streamMetadata{
-			"tenant": {
-				0: {{
-					hash:      0x1,
-					totalSize: 100,
-					rateBuckets: []rateBucket{{
-						timestamp: time.Now().UnixNano(),
-						size:      1,
-					}},
-					lastSeenAt: time.Now().UnixNano(),
-				}},
+		metadata: &streamMetadata{
+			stripes: []map[string]map[int32][]Stream{
+				{
+					"tenant": {
+						0: {{
+							Hash:      0x1,
+							TotalSize: 100,
+							RateBuckets: []RateBucket{{
+								Timestamp: time.Now().UnixNano(),
+								Size:      1,
+							}},
+							LastSeenAt: time.Now().UnixNano(),
+						}},
+					},
+				},
 			},
+			locks: make([]stripeLock, 1),
 		},
 		logger: log.NewNopLogger(),
+		partitionManager: &PartitionManager{
+			partitions: map[int32]int64{
+				0: time.Now().UnixNano(),
+			},
+		},
 	}
 
 	// Set up a mux router for the test server otherwise mux.Vars() won't work.