(chore): Simplify implementation of instance sort iterator (#11959)

Signed-off-by: Christian Haudum <[email protected]>
grafana · Feb 15, 2024 · 2177037 · 2177037
1 parent d0fae5c
commit 2177037
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 60 deletions.
diff --git a/pkg/bloomutils/iter.go b/pkg/bloomutils/iter.go
diff --git a/pkg/bloomutils/ring.go b/pkg/bloomutils/ring.go
@@ -107,31 +107,38 @@ func GetInstancesWithTokenRanges(id string, instances []ring.InstanceDesc) Insta
 // NewInstanceSortMergeIterator creates an iterator that yields instanceWithToken elements
 // where the token of the elements are sorted in ascending order.
 func NewInstanceSortMergeIterator(instances []ring.InstanceDesc) v1.Iterator[InstanceWithTokenRange] {
-	it := &sortMergeIterator[ring.InstanceDesc, uint32, InstanceWithTokenRange]{
-		items: instances,
-		transform: func(item ring.InstanceDesc, val uint32, prev *InstanceWithTokenRange) *InstanceWithTokenRange {
-			var prevToken uint32
-			if prev != nil {
-				prevToken = prev.MaxToken + 1
-			}
-			return &InstanceWithTokenRange{Instance: item, MinToken: prevToken, MaxToken: val}
-		},
-	}
-	sequences := make([]v1.PeekingIterator[v1.IndexedValue[uint32]], 0, len(instances))
-	for i := range instances {
-		sort.Slice(instances[i].Tokens, func(a, b int) bool {
-			return instances[i].Tokens[a] < instances[i].Tokens[b]
-		})
-		iter := v1.NewIterWithIndex[uint32](v1.NewSliceIter(instances[i].Tokens), i)
-		sequences = append(sequences, v1.NewPeekingIter[v1.IndexedValue[uint32]](iter))
+
+	tokenIters := make([]v1.PeekingIterator[v1.IndexedValue[uint32]], 0, len(instances))
+	for i, inst := range instances {
+		sort.Slice(inst.Tokens, func(a, b int) bool { return inst.Tokens[a] < inst.Tokens[b] })
+		itr := v1.NewIterWithIndex(v1.NewSliceIter[uint32](inst.Tokens), i)
+		tokenIters = append(tokenIters, v1.NewPeekingIter[v1.IndexedValue[uint32]](itr))
 	}
-	it.heap = v1.NewHeapIterator(
-		func(i, j v1.IndexedValue[uint32]) bool {
-			return i.Value() < j.Value()
+
+	heapIter := v1.NewHeapIterator[v1.IndexedValue[uint32]](
+		func(iv1, iv2 v1.IndexedValue[uint32]) bool {
+			return iv1.Value() < iv2.Value()
 		},
-		sequences...,
+		tokenIters...,
 	)
-	it.err = nil
 
-	return it
+	prevToken := -1
+	return v1.NewDedupingIter[v1.IndexedValue[uint32], InstanceWithTokenRange](
+		func(iv v1.IndexedValue[uint32], iwtr InstanceWithTokenRange) bool {
+			return false
+		},
+		func(iv v1.IndexedValue[uint32]) InstanceWithTokenRange {
+			minToken, maxToken := uint32(prevToken+1), iv.Value()
+			prevToken = int(maxToken)
+			return InstanceWithTokenRange{
+				Instance: instances[iv.Index()],
+				MinToken: minToken,
+				MaxToken: maxToken,
+			}
+		},
+		func(iv v1.IndexedValue[uint32], iwtr InstanceWithTokenRange) InstanceWithTokenRange {
+			panic("must not be called, because Eq() is always false")
+		},
+		v1.NewPeekingIter(heapIter),
+	)
 }
diff --git a/pkg/bloomutils/ring_test.go b/pkg/bloomutils/ring_test.go
@@ -11,6 +11,11 @@ import (
 )
 
 func TestBloomGatewayClient_SortInstancesByToken(t *testing.T) {
+	//          | 1  2  3  4  5  6  7  8  9  |
+	// ---------+----------------------------+
+	// ID 1     |             *           *  |
+	// ID 2     |       *           *        |
+	// ID 3     | *                          |
 	input := []ring.InstanceDesc{
 		{Id: "1", Tokens: []uint32{5, 9}},
 		{Id: "2", Tokens: []uint32{3, 7}},