github.com/hack0072008/kafka-go@v1.0.1/groupbalancer_test.go (about) 1 package kafka 2 3 import ( 4 "bytes" 5 "encoding/json" 6 "reflect" 7 "strconv" 8 "testing" 9 ) 10 11 func TestFindMembersByTopic(t *testing.T) { 12 a1 := GroupMember{ 13 ID: "a", 14 Topics: []string{"topic-1"}, 15 } 16 a12 := GroupMember{ 17 ID: "a", 18 Topics: []string{"topic-1", "topic-2"}, 19 } 20 b23 := GroupMember{ 21 ID: "b", 22 Topics: []string{"topic-2", "topic-3"}, 23 } 24 25 tests := map[string]struct { 26 Members []GroupMember 27 Expected map[string][]GroupMember 28 }{ 29 "empty": { 30 Expected: map[string][]GroupMember{}, 31 }, 32 "one member, one topic": { 33 Members: []GroupMember{a1}, 34 Expected: map[string][]GroupMember{ 35 "topic-1": { 36 a1, 37 }, 38 }, 39 }, 40 "one member, multiple topics": { 41 Members: []GroupMember{a12}, 42 Expected: map[string][]GroupMember{ 43 "topic-1": { 44 a12, 45 }, 46 "topic-2": { 47 a12, 48 }, 49 }, 50 }, 51 "multiple members, multiple topics": { 52 Members: []GroupMember{a12, b23}, 53 Expected: map[string][]GroupMember{ 54 "topic-1": { 55 a12, 56 }, 57 "topic-2": { 58 a12, 59 b23, 60 }, 61 "topic-3": { 62 b23, 63 }, 64 }, 65 }, 66 } 67 68 for label, test := range tests { 69 t.Run(label, func(t *testing.T) { 70 membersByTopic := findMembersByTopic(test.Members) 71 if !reflect.DeepEqual(test.Expected, membersByTopic) { 72 t.Errorf("expected %#v; got %#v", test.Expected, membersByTopic) 73 } 74 }) 75 } 76 } 77 78 func TestRangeAssignGroups(t *testing.T) { 79 newMeta := func(memberID string, topics ...string) GroupMember { 80 return GroupMember{ 81 ID: memberID, 82 Topics: topics, 83 } 84 } 85 86 newPartitions := func(partitionCount int, topics ...string) []Partition { 87 partitions := make([]Partition, 0, len(topics)*partitionCount) 88 for _, topic := range topics { 89 for partition := 0; partition < partitionCount; partition++ { 90 partitions = append(partitions, Partition{ 91 Topic: topic, 92 ID: partition, 93 }) 94 } 95 } 96 return partitions 97 } 98 99 tests := map[string]struct { 100 Members []GroupMember 101 Partitions []Partition 102 Expected GroupMemberAssignments 103 }{ 104 "empty": { 105 Expected: GroupMemberAssignments{}, 106 }, 107 "one member, one topic, one partition": { 108 Members: []GroupMember{ 109 newMeta("a", "topic-1"), 110 }, 111 Partitions: newPartitions(1, "topic-1"), 112 Expected: GroupMemberAssignments{ 113 "a": map[string][]int{ 114 "topic-1": {0}, 115 }, 116 }, 117 }, 118 "one member, one topic, multiple partitions": { 119 Members: []GroupMember{ 120 newMeta("a", "topic-1"), 121 }, 122 Partitions: newPartitions(3, "topic-1"), 123 Expected: GroupMemberAssignments{ 124 "a": map[string][]int{ 125 "topic-1": {0, 1, 2}, 126 }, 127 }, 128 }, 129 "multiple members, one topic, one partition": { 130 Members: []GroupMember{ 131 newMeta("a", "topic-1"), 132 newMeta("b", "topic-1"), 133 }, 134 Partitions: newPartitions(1, "topic-1"), 135 Expected: GroupMemberAssignments{ 136 "a": map[string][]int{}, 137 "b": map[string][]int{ 138 "topic-1": {0}, 139 }, 140 }, 141 }, 142 "multiple members, one topic, multiple partitions": { 143 Members: []GroupMember{ 144 newMeta("a", "topic-1"), 145 newMeta("b", "topic-1"), 146 }, 147 Partitions: newPartitions(3, "topic-1"), 148 Expected: GroupMemberAssignments{ 149 "a": map[string][]int{ 150 "topic-1": {0}, 151 }, 152 "b": map[string][]int{ 153 "topic-1": {1, 2}, 154 }, 155 }, 156 }, 157 "multiple members, multiple topics, multiple partitions": { 158 Members: []GroupMember{ 159 newMeta("a", "topic-1", "topic-2"), 160 newMeta("b", "topic-2", "topic-3"), 161 }, 162 Partitions: newPartitions(3, "topic-1", "topic-2", "topic-3"), 163 Expected: GroupMemberAssignments{ 164 "a": map[string][]int{ 165 "topic-1": {0, 1, 2}, 166 "topic-2": {0}, 167 }, 168 "b": map[string][]int{ 169 "topic-2": {1, 2}, 170 "topic-3": {0, 1, 2}, 171 }, 172 }, 173 }, 174 } 175 176 for label, test := range tests { 177 t.Run(label, func(t *testing.T) { 178 assignments := RangeGroupBalancer{}.AssignGroups(test.Members, test.Partitions) 179 if !reflect.DeepEqual(test.Expected, assignments) { 180 buf := bytes.NewBuffer(nil) 181 encoder := json.NewEncoder(buf) 182 encoder.SetIndent("", " ") 183 184 buf.WriteString("expected: ") 185 encoder.Encode(test.Expected) 186 buf.WriteString("got: ") 187 encoder.Encode(assignments) 188 189 t.Error(buf.String()) 190 } 191 }) 192 } 193 } 194 195 // For 66 members, 213 partitions, each member should get 213/66 = 3.22 partitions. 196 // This means that in practice, each member should get either 3 or 4 partitions 197 // assigned to it. Any other number is a failure. 198 func TestRangeAssignGroupsUnbalanced(t *testing.T) { 199 members := []GroupMember{} 200 for i := 0; i < 66; i++ { 201 members = append(members, GroupMember{ 202 ID: strconv.Itoa(i), 203 Topics: []string{"topic-1"}, 204 }) 205 } 206 partitions := []Partition{} 207 for i := 0; i < 213; i++ { 208 partitions = append(partitions, Partition{ 209 ID: i, 210 Topic: "topic-1", 211 }) 212 } 213 214 assignments := RangeGroupBalancer{}.AssignGroups(members, partitions) 215 if len(assignments) != len(members) { 216 t.Fatalf("Assignment count mismatch: %d != %d", len(assignments), len(members)) 217 } 218 219 for _, m := range assignments { 220 if len(m["topic-1"]) < 3 || len(m["topic-1"]) > 4 { 221 t.Fatalf("Expected assignment of 3 or 4 partitions, got %d", len(m["topic-1"])) 222 } 223 } 224 } 225 226 func TestRoundRobinAssignGroups(t *testing.T) { 227 newPartitions := func(partitionCount int, topics ...string) []Partition { 228 partitions := make([]Partition, 0, len(topics)*partitionCount) 229 for _, topic := range topics { 230 for partition := 0; partition < partitionCount; partition++ { 231 partitions = append(partitions, Partition{ 232 Topic: topic, 233 ID: partition, 234 }) 235 } 236 } 237 return partitions 238 } 239 240 tests := map[string]struct { 241 Members []GroupMember 242 Partitions []Partition 243 Expected GroupMemberAssignments 244 }{ 245 "empty": { 246 Expected: GroupMemberAssignments{}, 247 }, 248 "one member, one topic, one partition": { 249 Members: []GroupMember{ 250 { 251 ID: "a", 252 Topics: []string{"topic-1"}, 253 }, 254 }, 255 Partitions: newPartitions(1, "topic-1"), 256 Expected: GroupMemberAssignments{ 257 "a": map[string][]int{ 258 "topic-1": {0}, 259 }, 260 }, 261 }, 262 "one member, one topic, multiple partitions": { 263 Members: []GroupMember{ 264 { 265 ID: "a", 266 Topics: []string{"topic-1"}, 267 }, 268 }, 269 Partitions: newPartitions(3, "topic-1"), 270 Expected: GroupMemberAssignments{ 271 "a": map[string][]int{ 272 "topic-1": {0, 1, 2}, 273 }, 274 }, 275 }, 276 "multiple members, one topic, one partition": { 277 Members: []GroupMember{ 278 { 279 ID: "a", 280 Topics: []string{"topic-1"}, 281 }, 282 { 283 ID: "b", 284 Topics: []string{"topic-1"}, 285 }, 286 }, 287 Partitions: newPartitions(1, "topic-1"), 288 Expected: GroupMemberAssignments{ 289 "a": map[string][]int{ 290 "topic-1": {0}, 291 }, 292 "b": map[string][]int{}, 293 }, 294 }, 295 "multiple members, multiple topics, multiple partitions": { 296 Members: []GroupMember{ 297 { 298 ID: "a", 299 Topics: []string{"topic-1", "topic-2"}, 300 }, 301 { 302 ID: "b", 303 Topics: []string{"topic-2", "topic-3"}, 304 }, 305 }, 306 Partitions: newPartitions(3, "topic-1", "topic-2", "topic-3"), 307 Expected: GroupMemberAssignments{ 308 "a": map[string][]int{ 309 "topic-1": {0, 1, 2}, 310 "topic-2": {0, 2}, 311 }, 312 "b": map[string][]int{ 313 "topic-2": {1}, 314 "topic-3": {0, 1, 2}, 315 }, 316 }, 317 }, 318 } 319 320 for label, test := range tests { 321 t.Run(label, func(t *testing.T) { 322 assignments := RoundRobinGroupBalancer{}.AssignGroups(test.Members, test.Partitions) 323 if !reflect.DeepEqual(test.Expected, assignments) { 324 buf := bytes.NewBuffer(nil) 325 encoder := json.NewEncoder(buf) 326 encoder.SetIndent("", " ") 327 328 buf.WriteString("expected: ") 329 encoder.Encode(test.Expected) 330 buf.WriteString("got: ") 331 encoder.Encode(assignments) 332 333 t.Error(buf.String()) 334 } 335 }) 336 } 337 } 338 339 func TestFindMembersByTopicSortsByMemberID(t *testing.T) { 340 topic := "topic-1" 341 a := GroupMember{ 342 ID: "a", 343 Topics: []string{topic}, 344 } 345 b := GroupMember{ 346 ID: "b", 347 Topics: []string{topic}, 348 } 349 c := GroupMember{ 350 ID: "c", 351 Topics: []string{topic}, 352 } 353 354 testCases := map[string]struct { 355 Data []GroupMember 356 Expected []GroupMember 357 }{ 358 "in order": { 359 Data: []GroupMember{a, b}, 360 Expected: []GroupMember{a, b}, 361 }, 362 "out of order": { 363 Data: []GroupMember{a, c, b}, 364 Expected: []GroupMember{a, b, c}, 365 }, 366 } 367 368 for label, test := range testCases { 369 t.Run(label, func(t *testing.T) { 370 membersByTopic := findMembersByTopic(test.Data) 371 372 if actual := membersByTopic[topic]; !reflect.DeepEqual(test.Expected, actual) { 373 t.Errorf("expected %v; got %v", test.Expected, actual) 374 } 375 }) 376 } 377 } 378 379 func TestRackAffinityGroupBalancer(t *testing.T) { 380 t.Run("User Data", func(t *testing.T) { 381 t.Run("unknown zone", func(t *testing.T) { 382 b := RackAffinityGroupBalancer{} 383 zone, err := b.UserData() 384 if err != nil { 385 t.Fatal(err) 386 } 387 if string(zone) != "" { 388 t.Fatalf("expected empty zone but got %s", zone) 389 } 390 }) 391 392 t.Run("configure zone", func(t *testing.T) { 393 b := RackAffinityGroupBalancer{Rack: "zone1"} 394 zone, err := b.UserData() 395 if err != nil { 396 t.Fatal(err) 397 } 398 if string(zone) != "zone1" { 399 t.Fatalf("expected zone1 az but got %s", zone) 400 } 401 }) 402 }) 403 404 t.Run("Balance", func(t *testing.T) { 405 b := RackAffinityGroupBalancer{} 406 407 brokers := map[string]Broker{ 408 "z1": {ID: 1, Rack: "z1"}, 409 "z2": {ID: 2, Rack: "z2"}, 410 "z3": {ID: 2, Rack: "z3"}, 411 "": {}, 412 } 413 414 tests := []struct { 415 name string 416 memberCounts map[string]int 417 partitionCounts map[string]int 418 result map[string]map[string]int 419 }{ 420 { 421 name: "unknown and known zones", 422 memberCounts: map[string]int{ 423 "": 1, 424 "z1": 1, 425 "z2": 1, 426 }, 427 partitionCounts: map[string]int{ 428 "z1": 5, 429 "z2": 4, 430 "": 9, 431 }, 432 result: map[string]map[string]int{ 433 "z1": {"": 1, "z1": 5}, 434 "z2": {"": 2, "z2": 4}, 435 "": {"": 6}, 436 }, 437 }, 438 { 439 name: "all unknown", 440 memberCounts: map[string]int{ 441 "": 5, 442 }, 443 partitionCounts: map[string]int{ 444 "": 103, 445 }, 446 result: map[string]map[string]int{ 447 "": {"": 103}, 448 }, 449 }, 450 { 451 name: "remainder stays local", 452 memberCounts: map[string]int{ 453 "z1": 3, 454 "z2": 3, 455 "z3": 3, 456 }, 457 partitionCounts: map[string]int{ 458 "z1": 20, 459 "z2": 19, 460 "z3": 20, 461 }, 462 result: map[string]map[string]int{ 463 "z1": {"z1": 20}, 464 "z2": {"z2": 19}, 465 "z3": {"z3": 20}, 466 }, 467 }, 468 { 469 name: "imbalanced partitions", 470 memberCounts: map[string]int{ 471 "z1": 1, 472 "z2": 1, 473 "z3": 1, 474 }, 475 partitionCounts: map[string]int{ 476 "z1": 7, 477 "z2": 0, 478 "z3": 7, 479 }, 480 result: map[string]map[string]int{ 481 "z1": {"z1": 5}, 482 "z2": {"z1": 2, "z3": 2}, 483 "z3": {"z3": 5}, 484 }, 485 }, 486 { 487 name: "imbalanced members", 488 memberCounts: map[string]int{ 489 "z1": 5, 490 "z2": 3, 491 "z3": 1, 492 }, 493 partitionCounts: map[string]int{ 494 "z1": 9, 495 "z2": 9, 496 "z3": 9, 497 }, 498 result: map[string]map[string]int{ 499 "z1": {"z1": 9, "z3": 6}, 500 "z2": {"z2": 9}, 501 "z3": {"z3": 3}, 502 }, 503 }, 504 { 505 name: "no consumers in zone", 506 memberCounts: map[string]int{ 507 "z2": 10, 508 }, 509 partitionCounts: map[string]int{ 510 "z1": 20, 511 "z3": 19, 512 }, 513 result: map[string]map[string]int{ 514 "z2": {"z1": 20, "z3": 19}, 515 }, 516 }, 517 } 518 519 for _, tt := range tests { 520 t.Run(tt.name, func(t *testing.T) { 521 522 // create members per the distribution in the test case. 523 var members []GroupMember 524 for zone, count := range tt.memberCounts { 525 for i := 0; i < count; i++ { 526 members = append(members, GroupMember{ 527 ID: zone + ":" + strconv.Itoa(len(members)+1), 528 Topics: []string{"test"}, 529 UserData: []byte(zone), 530 }) 531 } 532 } 533 534 // create partitions per the distribution in the test case. 535 var partitions []Partition 536 for zone, count := range tt.partitionCounts { 537 for i := 0; i < count; i++ { 538 partitions = append(partitions, Partition{ 539 ID: len(partitions), 540 Topic: "test", 541 Leader: brokers[zone], 542 }) 543 } 544 } 545 546 res := b.AssignGroups(members, partitions) 547 548 // verification #1...all members must be assigned and with the 549 // correct load. 550 minLoad := len(partitions) / len(members) 551 maxLoad := minLoad 552 if len(partitions)%len(members) != 0 { 553 maxLoad++ 554 } 555 for _, member := range members { 556 assignments, _ := res[member.ID]["test"] 557 if len(assignments) < minLoad || len(assignments) > maxLoad { 558 t.Errorf("expected between %d and %d partitions for member %s", minLoad, maxLoad, member.ID) 559 } 560 } 561 562 // verification #2...all partitions are assigned, and the distribution 563 // per source zone matches. 564 partsPerZone := make(map[string]map[string]int) 565 uniqueParts := make(map[int]struct{}) 566 for id, topicToPartitions := range res { 567 568 for topic, assignments := range topicToPartitions { 569 if topic != "test" { 570 t.Fatalf("wrong topic...expected test but got %s", topic) 571 } 572 573 var member GroupMember 574 for _, m := range members { 575 if id == m.ID { 576 member = m 577 break 578 } 579 } 580 if member.ID == "" { 581 t.Fatal("empty member ID returned") 582 } 583 584 var partition Partition 585 for _, id := range assignments { 586 587 uniqueParts[id] = struct{}{} 588 589 for _, p := range partitions { 590 if p.ID == int(id) { 591 partition = p 592 break 593 } 594 } 595 if partition.Topic == "" { 596 t.Fatal("empty topic ID returned") 597 } 598 counts, ok := partsPerZone[string(member.UserData)] 599 if !ok { 600 counts = make(map[string]int) 601 partsPerZone[string(member.UserData)] = counts 602 } 603 counts[partition.Leader.Rack]++ 604 } 605 } 606 } 607 608 if len(partitions) != len(uniqueParts) { 609 t.Error("not all partitions were assigned") 610 } 611 if !reflect.DeepEqual(tt.result, partsPerZone) { 612 t.Errorf("wrong balanced zones. expected %v but got %v", tt.result, partsPerZone) 613 } 614 }) 615 } 616 }) 617 618 t.Run("Multi Topic", func(t *testing.T) { 619 b := RackAffinityGroupBalancer{} 620 621 brokers := map[string]Broker{ 622 "z1": {ID: 1, Rack: "z1"}, 623 "z2": {ID: 2, Rack: "z2"}, 624 "z3": {ID: 2, Rack: "z3"}, 625 "": {}, 626 } 627 628 members := []GroupMember{ 629 { 630 ID: "z1", 631 Topics: []string{"topic1", "topic2"}, 632 UserData: []byte("z1"), 633 }, 634 { 635 ID: "z2", 636 Topics: []string{"topic2", "topic3"}, 637 UserData: []byte("z2"), 638 }, 639 { 640 ID: "z3", 641 Topics: []string{"topic3", "topic1"}, 642 UserData: []byte("z3"), 643 }, 644 } 645 646 partitions := []Partition{ 647 { 648 ID: 1, 649 Topic: "topic1", 650 Leader: brokers["z1"], 651 }, 652 { 653 ID: 2, 654 Topic: "topic1", 655 Leader: brokers["z3"], 656 }, 657 { 658 ID: 1, 659 Topic: "topic2", 660 Leader: brokers["z1"], 661 }, 662 { 663 ID: 2, 664 Topic: "topic2", 665 Leader: brokers["z2"], 666 }, 667 { 668 ID: 1, 669 Topic: "topic3", 670 Leader: brokers["z3"], 671 }, 672 { 673 ID: 2, 674 Topic: "topic3", 675 Leader: brokers["z2"], 676 }, 677 } 678 679 expected := GroupMemberAssignments{ 680 "z1": {"topic1": []int{1}, "topic2": []int{1}}, 681 "z2": {"topic2": []int{2}, "topic3": []int{2}}, 682 "z3": {"topic3": []int{1}, "topic1": []int{2}}, 683 } 684 685 res := b.AssignGroups(members, partitions) 686 if !reflect.DeepEqual(expected, res) { 687 t.Fatalf("incorrect group assignment. expected %v but got %v", expected, res) 688 } 689 }) 690 }