From 3d317e1fb775c76cad88d11a926ae6e59cf20721 Mon Sep 17 00:00:00 2001
From: Nathan Hartwell <nathan.hartwell@shipt.com>
Date: Mon, 6 Jun 2022 16:48:31 -0500
Subject: [PATCH] fix: range balance strategy not like reference

https://kafka.apache.org/31/javadoc/org/apache/kafka/clients/consumer/RangeAssignor.html
---
 balance_strategy.go      | 23 ++++++++++++------
 balance_strategy_test.go | 51 +++++++++++++++++++++++++++++-----------
 2 files changed, 53 insertions(+), 21 deletions(-)

diff --git a/balance_strategy.go b/balance_strategy.go
index 9855bf443..49c1a40bd 100644
--- a/balance_strategy.go
+++ b/balance_strategy.go
@@ -58,18 +58,27 @@ type BalanceStrategy interface {
 // --------------------------------------------------------------------
 
 // BalanceStrategyRange is the default and assigns partitions as ranges to consumer group members.
-// Example with one topic T with six partitions (0..5) and two members (M1, M2):
-//   M1: {T: [0, 1, 2]}
-//   M2: {T: [3, 4, 5]}
+// This follows the same logic as
+// https://kafka.apache.org/31/javadoc/org/apache/kafka/clients/consumer/RangeAssignor.html
+//
+// Example with two topics T1 and T2 with six partitions each (0..5) and two members (M1, M2):
+//   M1: {T1: [0, 1, 2], T2: [0, 1, 2]}
+//   M2: {T2: [3, 4, 5], T2: [3, 4, 5]}
 var BalanceStrategyRange = &balanceStrategy{
 	name: RangeBalanceStrategyName,
 	coreFn: func(plan BalanceStrategyPlan, memberIDs []string, topic string, partitions []int32) {
-		step := float64(len(partitions)) / float64(len(memberIDs))
+		partitionsPerConsumer := len(partitions) / len(memberIDs)
+		consumersWithExtraPartition := len(partitions) % len(memberIDs)
+
+		sort.Strings(memberIDs)
 
 		for i, memberID := range memberIDs {
-			pos := float64(i)
-			min := int(math.Floor(pos*step + 0.5))
-			max := int(math.Floor((pos+1)*step + 0.5))
+			min := i*partitionsPerConsumer + int(math.Min(float64(consumersWithExtraPartition), float64(i)))
+			extra := 0
+			if i < consumersWithExtraPartition {
+				extra = 1
+			}
+			max := min + partitionsPerConsumer + extra
 			plan.Add(memberID, topic, partitions[min:max]...)
 		}
 	},
diff --git a/balance_strategy_test.go b/balance_strategy_test.go
index 44aee3ec0..2d2351d11 100644
--- a/balance_strategy_test.go
+++ b/balance_strategy_test.go
@@ -13,27 +13,48 @@ import (
 
 func TestBalanceStrategyRange(t *testing.T) {
 	tests := []struct {
+		name     string
 		members  map[string][]string
 		topics   map[string][]int32
 		expected BalanceStrategyPlan
 	}{
 		{
+			name:    "2 members, 2 topics, 4 partitions each",
 			members: map[string][]string{"M1": {"T1", "T2"}, "M2": {"T1", "T2"}},
 			topics:  map[string][]int32{"T1": {0, 1, 2, 3}, "T2": {0, 1, 2, 3}},
 			expected: BalanceStrategyPlan{
-				"M1": map[string][]int32{"T1": {0, 1}, "T2": {2, 3}},
-				"M2": map[string][]int32{"T1": {2, 3}, "T2": {0, 1}},
+				"M1": map[string][]int32{"T1": {0, 1}, "T2": {0, 1}},
+				"M2": map[string][]int32{"T1": {2, 3}, "T2": {2, 3}},
 			},
 		},
 		{
+			name:    "2 members, 2 topics, 4 partitions each (different member ids)",
+			members: map[string][]string{"M3": {"T1", "T2"}, "M4": {"T1", "T2"}},
+			topics:  map[string][]int32{"T1": {0, 1, 2, 3}, "T2": {0, 1, 2, 3}},
+			expected: BalanceStrategyPlan{
+				"M3": map[string][]int32{"T1": {0, 1}, "T2": {0, 1}},
+				"M4": map[string][]int32{"T1": {2, 3}, "T2": {2, 3}},
+			},
+		},
+		{
+			name:    "3 members, 1 topic, 1 partition each",
+			members: map[string][]string{"M1": {"T1"}, "M2": {"T1"}, "M3": {"T1"}},
+			topics:  map[string][]int32{"T1": {0}},
+			expected: BalanceStrategyPlan{
+				"M1": map[string][]int32{"T1": {0}},
+			},
+		},
+		{
+			name:    "2 members, 2 topics, 3 partitions each",
 			members: map[string][]string{"M1": {"T1", "T2"}, "M2": {"T1", "T2"}},
 			topics:  map[string][]int32{"T1": {0, 1, 2}, "T2": {0, 1, 2}},
 			expected: BalanceStrategyPlan{
-				"M1": map[string][]int32{"T1": {0, 1}, "T2": {2}},
-				"M2": map[string][]int32{"T1": {2}, "T2": {0, 1}},
+				"M1": map[string][]int32{"T1": {0, 1}, "T2": {0, 1}},
+				"M2": map[string][]int32{"T1": {2}, "T2": {2}},
 			},
 		},
 		{
+			name:    "2 members, 2 topics, different subscriptions",
 			members: map[string][]string{"M1": {"T1"}, "M2": {"T1", "T2"}},
 			topics:  map[string][]int32{"T1": {0, 1}, "T2": {0, 1}},
 			expected: BalanceStrategyPlan{
@@ -49,17 +70,19 @@ func TestBalanceStrategyRange(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		members := make(map[string]ConsumerGroupMemberMetadata)
-		for memberID, topics := range test.members {
-			members[memberID] = ConsumerGroupMemberMetadata{Topics: topics}
-		}
+		t.Run(test.name, func(t *testing.T) {
+			members := make(map[string]ConsumerGroupMemberMetadata)
+			for memberID, topics := range test.members {
+				members[memberID] = ConsumerGroupMemberMetadata{Topics: topics}
+			}
 
-		actual, err := strategy.Plan(members, test.topics)
-		if err != nil {
-			t.Errorf("Unexpected error %v", err)
-		} else if !reflect.DeepEqual(actual, test.expected) {
-			t.Errorf("Plan does not match expectation\nexpected: %#v\nactual: %#v", test.expected, actual)
-		}
+			actual, err := strategy.Plan(members, test.topics)
+			if err != nil {
+				t.Errorf("Unexpected error %v", err)
+			} else if !reflect.DeepEqual(actual, test.expected) {
+				t.Errorf("Plan does not match expectation\nexpected: %#v\nactual: %#v", test.expected, actual)
+			}
+		})
 	}
 }