Handle PR comments and improved tests

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

Author: roystchiang

pkg/ring/ring.go

@@ -342,29 +342,26 @@ 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
 			}
+		}
 
-			// We should only add instance zone if we are not going to extend,
+		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++
+		} else if r.cfg.ZoneAwarenessEnabled && info.Zone != "" {
+			// We should only add the 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)
-			}
+			distinctZones = append(distinctZones, info.Zone)
 		}
 
-		distinctHosts = append(distinctHosts, info.InstanceID)
 		instances = append(instances, instance)
 	}
 

pkg/ring/ring_test.go

@@ -135,56 +135,77 @@ func TestAddIngesterReplacesExistingTokens(t *testing.T) {
 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,
+	tests := map[string]struct {
+		replicationFactor int
+		expectedInstances int
+		expectedZones     int
+	}{
+		"should succeed if there are enough instances per zone on RF = 3": {
+			replicationFactor: 3,
+			expectedInstances: 3,
+			expectedZones:     3,
+		},
+		"should succeed if there are enough instances per zone on RF = 2": {
+			replicationFactor: 2,
+			expectedInstances: 2,
+			expectedZones:     2,
 		},
-		ringDesc:            r,
-		ringTokens:          r.GetTokens(),
-		ringTokensByZone:    r.getTokensByZone(),
-		ringInstanceByToken: r.getTokensInfo(),
-		ringZones:           getZones(r.getTokensByZone()),
-		strategy:            NewDefaultReplicationStrategy(),
 	}
 
-	_, bufHosts, bufZones := MakeBuffersForGet()
+	for testName, testData := range tests {
+		t.Run(testName, func(t *testing.T) {
+			r := NewDesc()
+			instances := map[string]InstanceDesc{
+				"instance-1": {Addr: "127.0.0.1", Zone: "zone-a", State: ACTIVE},
+				"instance-2": {Addr: "127.0.0.2", Zone: "zone-a", State: ACTIVE},
+				"instance-3": {Addr: "127.0.0.3", Zone: "zone-b", State: ACTIVE},
+				"instance-4": {Addr: "127.0.0.4", Zone: "zone-b", State: ACTIVE},
+				"instance-5": {Addr: "127.0.0.5", Zone: "zone-c", State: LEAVING},
+				"instance-6": {Addr: "127.0.0.6", Zone: "zone-c", 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:    testData.replicationFactor,
+					ZoneAwarenessEnabled: true,
+				},
+				ringDesc:            r,
+				ringTokens:          r.GetTokens(),
+				ringTokensByZone:    r.getTokensByZone(),
+				ringInstanceByToken: r.getTokensInfo(),
+				ringZones:           getZones(r.getTokensByZone()),
+				strategy:            NewDefaultReplicationStrategy(),
+			}
 
-	// Use the GenerateTokens to get an array of random uint32 values.
-	testValues := GenerateTokens(testCount, nil)
+			_, bufHosts, bufZones := MakeBuffersForGet()
 
-	for i := 0; i < testCount; i++ {
-		set, err := ring.Get(testValues[i], Write, instancesList, bufHosts, bufZones)
-		require.NoError(t, err)
+			// Use the GenerateTokens to get an array of random uint32 values.
+			testValues := GenerateTokens(testCount, nil)
 
-		distinctZones := map[string]int{}
-		for _, instance := range set.Instances {
-			distinctZones[instance.Zone]++
-		}
+			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)
+				assert.Len(t, set.Instances, testData.expectedInstances)
+				assert.Len(t, distinctZones, testData.expectedZones)
+			}
+		})
 	}
 }