Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: range balance strategy not like reference #2245

Merged
merged 1 commit into from
Jun 7, 2022
Merged

fix: range balance strategy not like reference #2245

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 16 additions & 7 deletions balance_strategy.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)))
Copy link
Contributor Author

@njhartwell njhartwell Jun 6, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The float math version is more concise, but gives different results than the reference implementation (e.g. when there are more members than available partitions, it assigns them to the middle of the list instead of the beginning).

extra := 0
if i < consumersWithExtraPartition {
extra = 1
}
max := min + partitionsPerConsumer + extra
plan.Add(memberID, topic, partitions[min:max]...)
}
},
Expand Down
51 changes: 37 additions & 14 deletions balance_strategy_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand All @@ -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)
}
})
}
}

Expand Down