fix: support multi-zone ingesters when converting global to local limits for streams in limiter.go

go.mod

@@ -50,7 +50,7 @@ require (
 	github.com/gorilla/mux v1.8.0
 	github.com/gorilla/websocket v1.5.0
 	github.com/grafana/cloudflare-go v0.0.0-20230110200409-c627cf6792f2
-	github.com/grafana/dskit v0.0.0-20240528015923-27d7d41066d3
+	github.com/grafana/dskit v0.0.0-20240626184720-35810fdf1c6d
 	github.com/grafana/go-gelf/v2 v2.0.1
 	github.com/grafana/gomemcache v0.0.0-20240229205252-cd6a66d6fb56
 	github.com/grafana/regexp v0.0.0-20221122212121-6b5c0a4cb7fd

go.sum

@@ -1017,8 +1017,8 @@ github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWm
 github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grafana/cloudflare-go v0.0.0-20230110200409-c627cf6792f2 h1:qhugDMdQ4Vp68H0tp/0iN17DM2ehRo1rLEdOFe/gB8I=
 github.com/grafana/cloudflare-go v0.0.0-20230110200409-c627cf6792f2/go.mod h1:w/aiO1POVIeXUQyl0VQSZjl5OAGDTL5aX+4v0RA1tcw=
-github.com/grafana/dskit v0.0.0-20240528015923-27d7d41066d3 h1:k8vINlI4w+RYc37NRwQlRe/IHYoEbu6KAe2XdGDeV1U=
-github.com/grafana/dskit v0.0.0-20240528015923-27d7d41066d3/go.mod h1:HvSf3uf8Ps2vPpzHeAFyZTdUcbVr+Rxpq1xcx7J/muc=
+github.com/grafana/dskit v0.0.0-20240626184720-35810fdf1c6d h1:CD8PWWX+9lYdgeMquSofmLErvCtk7jb+3/W/SH6oo/k=
+github.com/grafana/dskit v0.0.0-20240626184720-35810fdf1c6d/go.mod h1:HvSf3uf8Ps2vPpzHeAFyZTdUcbVr+Rxpq1xcx7J/muc=
 github.com/grafana/go-gelf/v2 v2.0.1 h1:BOChP0h/jLeD+7F9mL7tq10xVkDG15he3T1zHuQaWak=
 github.com/grafana/go-gelf/v2 v2.0.1/go.mod h1:lexHie0xzYGwCgiRGcvZ723bSNyNI8ZRD4s0CLobh90=
 github.com/grafana/gocql v0.0.0-20200605141915-ba5dc39ece85 h1:xLuzPoOzdfNb/RF/IENCw+oLVdZB4G21VPhkHBgwSHY=

pkg/ingester/ingester_test.go

@@ -1568,6 +1568,10 @@ func (r *readRingMock) ZonesCount() int {
 	return 1
 }
 
+func (r *readRingMock) HealthyInstancesInZoneCount() int {
+	return len(r.replicationSet.Instances)
+}
+
 func (r *readRingMock) Subring(_ uint32, _ int) ring.ReadRing {
 	return r
 }

pkg/ingester/limiter.go

@@ -20,6 +20,8 @@ const (
 // to count members
 type RingCount interface {
 	HealthyInstancesCount() int
+	HealthyInstancesInZoneCount() int
+	ZonesCount() int
 }
 
 type Limits interface {
@@ -106,22 +108,31 @@ func (l *Limiter) minNonZero(first, second int) int {
 }
 
 func (l *Limiter) convertGlobalToLocalLimit(globalLimit int) int {
-	if globalLimit == 0 {
+	if globalLimit == 0 || l.replicationFactor == 0 {
 		return 0
 	}
-	// todo: change to healthyInstancesInZoneCount() once
-	// Given we don't need a super accurate count (ie. when the ingesters
-	// topology changes) and we prefer to always be in favor of the tenant,
-	// we can use a per-ingester limit equal to:
-	// (global limit / number of ingesters) * replication factor
-	numIngesters := l.ring.HealthyInstancesCount()
 
-	// May happen because the number of ingesters is asynchronously updated.
-	// If happens, we just temporarily ignore the global limit.
+	zonesCount := l.ring.ZonesCount()
+	if zonesCount <= 1 {
+		return calculateLimitForSingleZone(globalLimit, l)
+	}
+
+	return calculateLimitForMultipleZones(globalLimit, zonesCount, l)
+}
+
+func calculateLimitForSingleZone(globalLimit int, l *Limiter) int {
+	numIngesters := l.ring.HealthyInstancesCount()
 	if numIngesters > 0 {
 		return int((float64(globalLimit) / float64(numIngesters)) * float64(l.replicationFactor))
 	}
+	return 0
+}
 
+func calculateLimitForMultipleZones(globalLimit, zonesCount int, l *Limiter) int {
+	ingestersInZone := l.ring.HealthyInstancesInZoneCount()
+	if ingestersInZone > 0 {
+		return int((float64(globalLimit) * float64(l.replicationFactor)) / float64(zonesCount) / float64(ingestersInZone))
+	}
 	return 0
 }
 

pkg/ingester/limiter_test.go

@@ -214,6 +214,14 @@ func (m *ringCountMock) HealthyInstancesCount() int {
 	return m.count
 }
 
+func (m *ringCountMock) ZonesCount() int {
+	return 1
+}
+
+func (m *ringCountMock) HealthyInstancesInZoneCount() int {
+	return m.count
+}
+
 // Assert some of the weirder (bug?) behavior of golang.org/x/time/rate
 func TestGoLimiter(t *testing.T) {
 	for _, tc := range []struct {
@@ -254,3 +262,59 @@ func TestGoLimiter(t *testing.T) {
 		})
 	}
 }
+
+type MockRing struct {
+	zonesCount                  int
+	healthyInstancesCount       int
+	healthyInstancesInZoneCount int
+}
+
+func (m *MockRing) ZonesCount() int {
+	return m.zonesCount
+}
+
+func (m *MockRing) HealthyInstancesCount() int {
+	return m.healthyInstancesCount
+}
+
+func (m *MockRing) HealthyInstancesInZoneCount() int {
+	return m.healthyInstancesInZoneCount
+}
+
+func TestConvertGlobalToLocalLimit(t *testing.T) {
+	tests := []struct {
+		name                        string
+		globalLimit                 int
+		zonesCount                  int
+		healthyInstancesCount       int
+		healthyInstancesInZoneCount int
+		replicationFactor           int
+		expectedLocalLimit          int
+	}{
+		{"GlobalLimitZero", 0, 1, 1, 1, 3, 0},
+		{"SingleZoneMultipleIngesters", 100, 1, 10, 10, 3, 30},
+		{"MultipleZones", 200, 3, 30, 10, 3, 20},
+		{"MultipleZonesNoHealthyIngesters", 200, 2, 0, 0, 3, 0},
+		{"MultipleZonesNoHealthyIngestersInZone", 200, 3, 10, 0, 3, 0},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			mockRing := &MockRing{
+				zonesCount:                  tc.zonesCount,
+				healthyInstancesCount:       tc.healthyInstancesCount,
+				healthyInstancesInZoneCount: tc.healthyInstancesInZoneCount,
+			}
+
+			limiter := &Limiter{
+				ring:              mockRing,
+				replicationFactor: tc.replicationFactor,
+			}
+
+			localLimit := limiter.convertGlobalToLocalLimit(tc.globalLimit)
+			if localLimit != tc.expectedLocalLimit {
+				t.Errorf("expected %d, got %d", tc.expectedLocalLimit, localLimit)
+			}
+		})
+	}
+}