Fix issue where we are not returning all the available ingesters in the ring (#3977)

Signed-off-by: Roy Chiang <[email protected]>

Author: roystchiang

CHANGELOG.md

@@ -20,6 +20,7 @@
 * [BUGFIX] Distributor: reverted changes done to rate limiting in #3825. #3948
 * [BUGFIX] Ingester: Fix race condition when opening and closing tsdb concurrently. #3959
 * [BUGFIX] Querier: streamline tracing spans. #3924
+* [BUGFIX] Distributor: fix issue where we are not returning all the available ingesters in the ring. #3977
 
 ## 1.8.0 in progress
 

pkg/ring/ring.go

@@ -342,23 +342,29 @@ func (r *Ring) Get(key uint32, op Operation, bufDescs []InstanceDesc, bufHosts,
 			continue
 		}
 
+		instance := r.ringDesc.Ingesters[info.InstanceID]
+
+		// Check whether the replica set should be extended given we're including
+		// this instance.
+		shouldExtendReplicaSet := op.ShouldExtendReplicaSetOnState(instance.State)
+		if shouldExtendReplicaSet {
+			n++
+		}
+
 		// Ignore if the instances don't have a zone set.
 		if r.cfg.ZoneAwarenessEnabled && info.Zone != "" {
 			if util.StringsContain(distinctZones, info.Zone) {
 				continue
 			}
-			distinctZones = append(distinctZones, info.Zone)
-		}
 
-		distinctHosts = append(distinctHosts, info.InstanceID)
-		instance := r.ringDesc.Ingesters[info.InstanceID]
-
-		// Check whether the replica set should be extended given we're including
-		// this instance.
-		if op.ShouldExtendReplicaSetOnState(instance.State) {
-			n++
+			// We should only add instance zone if we are not going to extend,
+			// as we want to extend the instance in the same AZ.
+			if !shouldExtendReplicaSet {
+				distinctZones = append(distinctZones, info.Zone)
+			}
 		}
 
+		distinctHosts = append(distinctHosts, info.InstanceID)
 		instances = append(instances, instance)
 	}
 

pkg/ring/ring_test.go

@@ -132,6 +132,62 @@ func TestAddIngesterReplacesExistingTokens(t *testing.T) {
 	require.Equal(t, newTokens, r.Ingesters[ing1Name].Tokens)
 }
 
+func TestRing_Get_ZoneAwarenessWithIngesterLeaving(t *testing.T) {
+	const testCount = 10000
+
+	r := NewDesc()
+	instances := map[string]InstanceDesc{
+		"instance-1": {Addr: "127.0.0.1", Zone: "zone-a", Tokens: GenerateTokens(128, nil), State: ACTIVE},
+		"instance-2": {Addr: "127.0.0.2", Zone: "zone-a", Tokens: GenerateTokens(128, nil), State: ACTIVE},
+		"instance-3": {Addr: "127.0.0.3", Zone: "zone-b", Tokens: GenerateTokens(128, nil), State: ACTIVE},
+		"instance-4": {Addr: "127.0.0.4", Zone: "zone-b", Tokens: GenerateTokens(128, nil), State: ACTIVE},
+		"instance-5": {Addr: "127.0.0.5", Zone: "zone-c", Tokens: GenerateTokens(128, nil), State: LEAVING},
+		"instance-6": {Addr: "127.0.0.6", Zone: "zone-c", Tokens: GenerateTokens(128, nil), State: ACTIVE},
+	}
+	var prevTokens []uint32
+	for id, instance := range instances {
+		ingTokens := GenerateTokens(128, prevTokens)
+		r.AddIngester(id, instance.Addr, instance.Zone, ingTokens, instance.State, time.Now())
+		prevTokens = append(prevTokens, ingTokens...)
+	}
+	instancesList := make([]InstanceDesc, 0, len(r.GetIngesters()))
+	for _, v := range r.GetIngesters() {
+		instancesList = append(instancesList, v)
+	}
+
+	ring := Ring{
+		cfg: Config{
+			HeartbeatTimeout:     time.Hour,
+			ReplicationFactor:    3,
+			ZoneAwarenessEnabled: true,
+		},
+		ringDesc:            r,
+		ringTokens:          r.GetTokens(),
+		ringTokensByZone:    r.getTokensByZone(),
+		ringInstanceByToken: r.getTokensInfo(),
+		ringZones:           getZones(r.getTokensByZone()),
+		strategy:            NewDefaultReplicationStrategy(),
+	}
+
+	_, bufHosts, bufZones := MakeBuffersForGet()
+
+	// Use the GenerateTokens to get an array of random uint32 values.
+	testValues := GenerateTokens(testCount, nil)
+
+	for i := 0; i < testCount; i++ {
+		set, err := ring.Get(testValues[i], Write, instancesList, bufHosts, bufZones)
+		require.NoError(t, err)
+
+		distinctZones := map[string]int{}
+		for _, instance := range set.Instances {
+			distinctZones[instance.Zone]++
+		}
+
+		assert.Equal(t, len(set.Instances), 3)
+		assert.Equal(t, len(distinctZones), 3)
+	}
+}
+
 func TestRing_Get_ZoneAwareness(t *testing.T) {
 	// Number of tests to run.
 	const testCount = 10000