Solved all TODOs

Signed-off-by: Marco Pracucci <[email protected]>

Author: pracucci

pkg/ring/model.go

@@ -397,7 +397,11 @@ type TokenDesc struct {
 
 // getTokens returns sorted list of tokens with ingester IDs, owned by each ingester in the ring.
 func (d *Desc) getTokens() []TokenDesc {
-	tokens := make([]TokenDesc, 0, d.getTokensCountEstimation())
+	numTokens := 0
+	for _, ing := range d.Ingesters {
+		numTokens += len(ing.Tokens)
+	}
+	tokens := make([]TokenDesc, 0, numTokens)
 	for key, ing := range d.Ingesters {
 		for _, token := range ing.Tokens {
 			tokens = append(tokens, TokenDesc{Token: token, Ingester: key, Zone: ing.GetZone()})
@@ -408,13 +412,12 @@ func (d *Desc) getTokens() []TokenDesc {
 	return tokens
 }
 
-// TODO comment
-// TODO test me
+// getTokensByZone returns instances tokens grouped by zone. Tokens within each zone
+// are guaranteed to be sorted.
 func (d *Desc) getTokensByZone() map[string][]TokenDesc {
 	zones := map[string][]TokenDesc{}
 
 	for key, ing := range d.Ingesters {
-		// TODO initialize zones[ing.Zone] with a slice allocated for getTokensCountEstimation() ?
 		for _, token := range ing.Tokens {
 			zones[ing.Zone] = append(zones[ing.Zone], TokenDesc{Token: token, Ingester: key, Zone: ing.GetZone()})
 		}
@@ -428,17 +431,6 @@ func (d *Desc) getTokensByZone() map[string][]TokenDesc {
 	return zones
 }
 
-// getTokensCountEstimation returns a quick estimation of the number of tokens
-// across all instances in the ring, assuming each instance has the same number
-// of tokens.
-func (d *Desc) getTokensCountEstimation() int {
-	for _, ing := range d.Ingesters {
-		return len(ing.Tokens) * len(d.Ingesters)
-	}
-
-	return 0
-}
-
 func GetOrCreateRingDesc(d interface{}) *Desc {
 	if d == nil {
 		return NewDesc()

pkg/ring/model_test.go

@@ -186,3 +186,57 @@ func TestDesc_Ready(t *testing.T) {
 		t.Fatal("expected ready, got", err)
 	}
 }
+
+func TestDesc_getTokensByZone(t *testing.T) {
+	tests := map[string]struct {
+		desc     *Desc
+		expected map[string][]TokenDesc
+	}{
+		"empty ring": {
+			desc:     &Desc{Ingesters: map[string]IngesterDesc{}},
+			expected: map[string][]TokenDesc{},
+		},
+		"single zone": {
+			desc: &Desc{Ingesters: map[string]IngesterDesc{
+				"instance-1": {Addr: "127.0.0.1", Tokens: []uint32{1, 5}, Zone: ""},
+				"instance-2": {Addr: "127.0.0.1", Tokens: []uint32{2, 4}, Zone: ""},
+				"instance-3": {Addr: "127.0.0.1", Tokens: []uint32{3, 6}, Zone: ""},
+			}},
+			expected: map[string][]TokenDesc{
+				"": {
+					{Token: 1, Ingester: "instance-1", Zone: ""},
+					{Token: 2, Ingester: "instance-2", Zone: ""},
+					{Token: 3, Ingester: "instance-3", Zone: ""},
+					{Token: 4, Ingester: "instance-2", Zone: ""},
+					{Token: 5, Ingester: "instance-1", Zone: ""},
+					{Token: 6, Ingester: "instance-3", Zone: ""},
+				},
+			},
+		},
+		"multiple zones": {
+			desc: &Desc{Ingesters: map[string]IngesterDesc{
+				"instance-1": {Addr: "127.0.0.1", Tokens: []uint32{1, 5}, Zone: "zone-1"},
+				"instance-2": {Addr: "127.0.0.1", Tokens: []uint32{2, 4}, Zone: "zone-1"},
+				"instance-3": {Addr: "127.0.0.1", Tokens: []uint32{3, 6}, Zone: "zone-2"},
+			}},
+			expected: map[string][]TokenDesc{
+				"zone-1": {
+					{Token: 1, Ingester: "instance-1", Zone: "zone-1"},
+					{Token: 2, Ingester: "instance-2", Zone: "zone-1"},
+					{Token: 4, Ingester: "instance-2", Zone: "zone-1"},
+					{Token: 5, Ingester: "instance-1", Zone: "zone-1"},
+				},
+				"zone-2": {
+					{Token: 3, Ingester: "instance-3", Zone: "zone-2"},
+					{Token: 6, Ingester: "instance-3", Zone: "zone-2"},
+				},
+			},
+		},
+	}
+
+	for testName, testData := range tests {
+		t.Run(testName, func(t *testing.T) {
+			assert.Equal(t, testData.expected, testData.desc.getTokensByZone())
+		})
+	}
+}

pkg/ring/ring_test.go

@@ -248,21 +248,8 @@ func TestRing_ShuffleShard_Stability(t *testing.T) {
 		shardSizes     = []int{3, 6, 9, 12, 15}
 	)
 
-	// Initialise the ring instances.
-	instances := make(map[string]IngesterDesc, numInstances)
-	for i := 1; i <= numInstances; i++ {
-		id := fmt.Sprintf("instance-%d", i)
-		instances[id] = IngesterDesc{
-			Addr:      fmt.Sprintf("127.0.0.%d", i),
-			Timestamp: time.Now().Unix(),
-			State:     ACTIVE,
-			Tokens:    GenerateTokens(128, nil),
-			Zone:      fmt.Sprintf("zone-%d", i%numZones),
-		}
-	}
-
 	// Initialise the ring.
-	ringDesc := &Desc{Ingesters: instances}
+	ringDesc := &Desc{Ingesters: generateRingInstances(numInstances, numZones)}
 	ring := Ring{
 		cfg:              Config{HeartbeatTimeout: time.Hour},
 		ringDesc:         ringDesc,
@@ -706,7 +693,6 @@ func generateTokensLinear(instanceID, numInstances, numTokens int) []uint32 {
 	return tokens
 }
 
-// TODO use across the tests
 func generateRingInstances(numInstances, numZones int) map[string]IngesterDesc {
 	instances := make(map[string]IngesterDesc, numInstances)
 
@@ -729,7 +715,6 @@ func generateRingInstance(id, zone int) (string, IngesterDesc) {
 }
 
 // compareReplicationSets returns the list of instance addresses which differ between the two sets.
-// TODO check if can be used to simplify other tests
 func compareReplicationSets(first, second ReplicationSet) (added, removed []string) {
 	for _, instance := range first.Ingesters {
 		if !second.Includes(instance.Addr) {

pkg/ring/util.go

@@ -82,7 +82,8 @@ func WaitInstanceState(ctx context.Context, r *Ring, instanceID string, state In
 	return backoff.Err()
 }
 
-// TODO comment (guarantee sorting)
+// getZones return the list zones from the provided tokens. The returned list
+// is guaranteed to be sorted.
 func getZones(tokens map[string][]TokenDesc) []string {
 	var zones []string
 
@@ -94,7 +95,7 @@ func getZones(tokens map[string][]TokenDesc) []string {
 	return zones
 }
 
-// TODO comment
+// searchToken returns the offset of the tokens entry holding the range for the provided key.
 func searchToken(tokens []TokenDesc, key uint32) int {
 	i := sort.Search(len(tokens), func(x int) bool {
 		return tokens[x].Token > key