From 01915cf82ee752e91c90afe358c4707b72678247 Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Fri, 21 Jun 2019 15:22:16 -0700 Subject: [PATCH 1/9] Add sticky partition assignor --- balance_strategy.go | 898 ++++++++++++ balance_strategy_test.go | 2176 +++++++++++++++++++++++++++++ consumer_group.go | 28 +- examples/consumergroup/main.go | 25 +- sticky_assignor_user_data.go | 124 ++ sticky_assignor_user_data_test.go | 51 + 6 files changed, 3291 insertions(+), 11 deletions(-) create mode 100644 sticky_assignor_user_data.go create mode 100644 sticky_assignor_user_data_test.go diff --git a/balance_strategy.go b/balance_strategy.go index 2fce17fb6..7392cc738 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -5,6 +5,10 @@ import ( "sort" ) +const ( + defaultGeneration = -1 +) + // BalanceStrategyPlan is the results of any BalanceStrategy.Plan attempt. // It contains an allocation of topic/partitions by memberID in the form of // a `memberID -> topic -> partitions` map. @@ -68,6 +72,19 @@ var BalanceStrategyRoundRobin = &balanceStrategy{ }, } +// BalanceStrategySticky assigns partitions to members with an attempt to preserve earlier assignments +// while maintain a balanced partition distribution. +// Example with topic T with six partitions (0..5) and two members (M1, M2): +// M1: {T: [0, 2, 4]} +// M2: {T: [1, 3, 5]} +// +// On reassignment with an additional consumer, you might get an assignment plan like: +// M1: {T: [0, 2]} +// M2: {T: [1, 3]} +// M3: {T: [4, 5]} +// +var BalanceStrategySticky = &stickyBalanceStrategy{} + // -------------------------------------------------------------------- type balanceStrategy struct { @@ -127,3 +144,884 @@ func balanceStrategyHashValue(vv ...string) uint32 { } return h } + +type stickyBalanceStrategy struct { + movements partitionMovements +} + +// Name implements BalanceStrategy. +func (s *stickyBalanceStrategy) Name() string { return "sticky" } + +// Plan implements BalanceStrategy. +func (s *stickyBalanceStrategy) Plan(members map[string]ConsumerGroupMemberMetadata, topics map[string][]int32) (BalanceStrategyPlan, error) { + // track partition movements during generation of the partition assignment plan + s.movements = partitionMovements{ + Movements: make(map[topicPartitionAssignment]consumerPair), + PartitionMovementsByTopic: make(map[string]map[consumerPair]map[topicPartitionAssignment]bool), + } + + // prepopulate the current assignment state from userdata on the consumer group members + currentAssignment, prevAssignment, err := prepopulateCurrentAssignments(members) + if err != nil { + return nil, err + } + + // determine if we're dealing with a completely fresh assignment, or if there's existing assignment state + isFreshAssignment := false + if len(currentAssignment) == 0 { + isFreshAssignment = true + } + + // create a mapping of all current topic partitions and the consumers that can be assigned to them + partition2AllPotentialConsumers := make(map[topicPartitionAssignment][]string) + for topic, partitions := range topics { + for _, partition := range partitions { + partition2AllPotentialConsumers[topicPartitionAssignment{Topic: topic, Partition: partition}] = make([]string, 0) + } + } + + // create a mapping of all consumers to all potential topic partitions that can be assigned to them + // also, populate the mapping of partitions to potential consumers + consumer2AllPotentialPartitions := make(map[string][]topicPartitionAssignment, len(members)) + for memberID, meta := range members { + consumer2AllPotentialPartitions[memberID] = make([]topicPartitionAssignment, 0) + for _, topicSubscription := range meta.Topics { + // only evaluate topic subscriptions that are present in the supplied topics map + if _, found := topics[topicSubscription]; found { + for _, partition := range topics[topicSubscription] { + topicPartition := topicPartitionAssignment{Topic: topicSubscription, Partition: partition} + consumer2AllPotentialPartitions[memberID] = append(consumer2AllPotentialPartitions[memberID], topicPartition) + partition2AllPotentialConsumers[topicPartition] = append(partition2AllPotentialConsumers[topicPartition], memberID) + } + } + } + + // add this consumer to currentAssignment (with an empty topic partition assignment) if it does not already exist + if _, exists := currentAssignment[memberID]; !exists { + currentAssignment[memberID] = make([]topicPartitionAssignment, 0) + } + } + + // create a mapping of each partition to its current consumer, where possible + currentPartitionConsumer := make(map[topicPartitionAssignment]string, len(currentAssignment)) + for memberID, partitions := range currentAssignment { + for _, partition := range partitions { + currentPartitionConsumer[partition] = memberID + } + } + + // sort the topic partitions in order of priority for reassignment + sortedPartitions := sortPartitions(currentAssignment, prevAssignment, isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions) + unassignedPartitions := deepCopyPartitions(sortedPartitions) + for memberID, partitions := range currentAssignment { + // if a consumer that existed before (and had some partition assignments) is now removed, remove it from currentAssignment + if _, exists := members[memberID]; !exists { + for _, partition := range partitions { + delete(currentPartitionConsumer, partition) + } + delete(currentAssignment, memberID) + continue + } + + // otherwise (the consumer still exists) + updatedPartitions := deepCopyPartitions(partitions) + for _, partition := range partitions { + if _, exists := partition2AllPotentialConsumers[partition]; !exists { + // if this topic partition of this consumer no longer exists remove it from currentAssignment of the consumer + updatedPartitions = removeTopicPartitionFromMemberAssignments(updatedPartitions, partition) + delete(currentPartitionConsumer, partition) + } else if _, exists := topics[partition.Topic]; !exists { + // if this partition cannot remain assigned to its current consumer because the consumer + // is no longer subscribed to its topic remove it from currentAssignment of the consumer + updatedPartitions = removeTopicPartitionFromMemberAssignments(updatedPartitions, partition) + } else { + // otherwise, remove the topic partition from those that need to be assigned only if + // its current consumer is still subscribed to its topic (because it is already assigned + // and we would want to preserve that assignment as much as possible) + unassignedPartitions = removeTopicPartitionFromMemberAssignments(unassignedPartitions, partition) + } + } + currentAssignment[memberID] = updatedPartitions + } + + // at this point we have preserved all valid topic partition to consumer assignments and removed + // all invalid topic partitions and invalid consumers. Now we need to assign unassignedPartitions + // to consumers so that the topic partition assignments are as balanced as possible. + + // an ascending sorted set of consumers based on how many topic partitions are already assigned to them + sortedCurrentSubscriptions := sortMemberIDsByPartitionAssignments(currentAssignment) + s.balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer) + + // Assemble plan + plan := make(BalanceStrategyPlan, len(currentAssignment)) + for memberID, assignments := range currentAssignment { + if len(assignments) == 0 { + plan[memberID] = make(map[string][]int32, 0) + } else { + for _, assignment := range assignments { + plan.Add(memberID, assignment.Topic, assignment.Partition) + } + } + } + return plan, nil +} + +// Balance assignments across consumers for maximum fairness and stickiness. +func (s *stickyBalanceStrategy) balance(currentAssignment map[string][]topicPartitionAssignment, prevAssignment map[topicPartitionAssignment]consumerGenerationPair, sortedPartitions []topicPartitionAssignment, unassignedPartitions []topicPartitionAssignment, sortedCurrentSubscriptions []string, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment, partition2AllPotentialConsumers map[topicPartitionAssignment][]string, currentPartitionConsumer map[topicPartitionAssignment]string) { + initializing := false + if len(sortedCurrentSubscriptions) == 0 || len(currentAssignment[sortedCurrentSubscriptions[0]]) == 0 { + initializing = true + } + + // assign all unassigned partitions + for _, partition := range unassignedPartitions { + // skip if there is no potential consumer for the partition + if len(partition2AllPotentialConsumers[partition]) == 0 { + continue + } + sortedCurrentSubscriptions = assignPartition(partition, sortedCurrentSubscriptions, currentAssignment, consumer2AllPotentialPartitions, currentPartitionConsumer) + } + + // narrow down the reassignment scope to only those partitions that can actually be reassigned + for partition := range partition2AllPotentialConsumers { + if !canTopicPartitionParticipateInReassignment(partition, partition2AllPotentialConsumers) { + sortedPartitions = removeTopicPartitionFromMemberAssignments(sortedPartitions, partition) + } + } + + // narrow down the reassignment scope to only those consumers that are subject to reassignment + fixedAssignments := make(map[string][]topicPartitionAssignment) + for memberID := range consumer2AllPotentialPartitions { + if !canConsumerParticipateInReassignment(memberID, currentAssignment, consumer2AllPotentialPartitions, partition2AllPotentialConsumers) { + fixedAssignments[memberID] = currentAssignment[memberID] + delete(currentAssignment, memberID) + sortedCurrentSubscriptions = sortMemberIDsByPartitionAssignments(currentAssignment) + } + } + + // create a deep copy of the current assignment so we can revert to it if we do not get a more balanced assignment later + preBalanceAssignment := deepCopyAssignment(currentAssignment) + preBalancePartitionConsumers := make(map[topicPartitionAssignment]string, len(currentPartitionConsumer)) + for k, v := range currentPartitionConsumer { + preBalancePartitionConsumers[k] = v + } + + reassignmentPerformed := s.performReassignments(sortedPartitions, currentAssignment, prevAssignment, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer) + + // if we are not preserving existing assignments and we have made changes to the current assignment + // make sure we are getting a more balanced assignment; otherwise, revert to previous assignment + if !initializing && reassignmentPerformed && getBalanceScore(currentAssignment) >= getBalanceScore(preBalanceAssignment) { + currentAssignment = deepCopyAssignment(preBalanceAssignment) + currentPartitionConsumer = make(map[topicPartitionAssignment]string, len(preBalancePartitionConsumers)) + for k, v := range preBalancePartitionConsumers { + currentPartitionConsumer[k] = v + } + } + + // add the fixed assignments (those that could not change) back + for consumer, assignments := range fixedAssignments { + currentAssignment[consumer] = assignments + } +} + +// Calculate the balance score of the given assignment, as the sum of assigned partitions size difference of all consumer pairs. +// A perfectly balanced assignment (with all consumers getting the same number of partitions) has a balance score of 0. +// Lower balance score indicates a more balanced assignment. +func getBalanceScore(assignment map[string][]topicPartitionAssignment) int { + consumer2AssignmentSize := make(map[string]int, len(assignment)) + for memberID, partitions := range assignment { + consumer2AssignmentSize[memberID] = len(partitions) + } + + var score float64 + for memberID, consumerAssignmentSize := range consumer2AssignmentSize { + delete(consumer2AssignmentSize, memberID) + for _, otherConsumerAssignmentSize := range consumer2AssignmentSize { + score += math.Abs(float64(consumerAssignmentSize - otherConsumerAssignmentSize)) + } + } + return int(score) +} + +// Determine whether the current assignment plan is balanced. +func isBalanced(currentAssignment map[string][]topicPartitionAssignment, sortedCurrentSubscriptions []string, allSubscriptions map[string][]topicPartitionAssignment) bool { + sortedCurrentSubscriptions = sortMemberIDsByPartitionAssignments(currentAssignment) + min := len(currentAssignment[sortedCurrentSubscriptions[0]]) + max := len(currentAssignment[sortedCurrentSubscriptions[len(sortedCurrentSubscriptions)-1]]) + if min >= max-1 { + // if minimum and maximum numbers of partitions assigned to consumers differ by at most one return true + return true + } + + // create a mapping from partitions to the consumer assigned to them + allPartitions := make(map[topicPartitionAssignment]string) + for memberID, partitions := range currentAssignment { + for _, partition := range partitions { + if _, exists := allPartitions[partition]; exists { + Logger.Printf("Topic %s Partition %d is assigned more than one consumer", partition.Topic, partition.Partition) + } + allPartitions[partition] = memberID + } + } + + // for each consumer that does not have all the topic partitions it can get make sure none of the topic partitions it + // could but did not get cannot be moved to it (because that would break the balance) + for _, memberID := range sortedCurrentSubscriptions { + consumerPartitions := currentAssignment[memberID] + consumerPartitionCount := len(consumerPartitions) + + // skip if this consumer already has all the topic partitions it can get + if consumerPartitionCount == len(allSubscriptions[memberID]) { + continue + } + + // otherwise make sure it cannot get any more + potentialTopicPartitions := allSubscriptions[memberID] + for _, partition := range potentialTopicPartitions { + if !memberAssignmentsIncludeTopicPartition(currentAssignment[memberID], partition) { + otherConsumer := allPartitions[partition] + otherConsumerPartitionCount := len(currentAssignment[otherConsumer]) + if consumerPartitionCount < otherConsumerPartitionCount { + return false + } + } + } + } + return true +} + +// Reassign all topic partitions that need reassignment until balanced. +func (s *stickyBalanceStrategy) performReassignments(reassignablePartitions []topicPartitionAssignment, currentAssignment map[string][]topicPartitionAssignment, prevAssignment map[topicPartitionAssignment]consumerGenerationPair, sortedCurrentSubscriptions []string, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment, partition2AllPotentialConsumers map[topicPartitionAssignment][]string, currentPartitionConsumer map[topicPartitionAssignment]string) bool { + reassignmentPerformed := false + modified := false + + // repeat reassignment until no partition can be moved to improve the balance + for { + modified = false + // reassign all reassignable partitions (starting from the partition with least potential consumers and if needed) + // until the full list is processed or a balance is achieved + for _, partition := range reassignablePartitions { + if isBalanced(currentAssignment, sortedCurrentSubscriptions, consumer2AllPotentialPartitions) { + break + } + + // the partition must have at least two consumers + if len(partition2AllPotentialConsumers[partition]) <= 1 { + Logger.Printf("Expected more than one potential consumer for partition %s topic %d", partition.Topic, partition.Partition) + } + + // the partition must have a consumer + consumer := currentPartitionConsumer[partition] + if consumer == "" { + Logger.Printf("Expected topic %s partition %d to be assigned to a consumer", partition.Topic, partition.Partition) + } + + if _, exists := prevAssignment[partition]; exists { + if len(currentAssignment[consumer]) > (len(currentAssignment[prevAssignment[partition].MemberID]) + 1) { + sortedCurrentSubscriptions = s.reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, prevAssignment[partition].MemberID) + reassignmentPerformed = true + modified = true + continue + } + } + + // check if a better-suited consumer exists for the partition; if so, reassign it + for _, otherConsumer := range partition2AllPotentialConsumers[partition] { + if len(currentAssignment[consumer]) > (len(currentAssignment[otherConsumer]) + 1) { + sortedCurrentSubscriptions = s.reassignPartitionToNewConsumer(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions) + reassignmentPerformed = true + modified = true + break + } + + } + } + if !modified { + return reassignmentPerformed + } + } +} + +// Identify a new consumer for a topic partition and reassign it. +func (s *stickyBalanceStrategy) reassignPartitionToNewConsumer(partition topicPartitionAssignment, currentAssignment map[string][]topicPartitionAssignment, sortedCurrentSubscriptions []string, currentPartitionConsumer map[topicPartitionAssignment]string, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment) []string { + for _, anotherConsumer := range sortedCurrentSubscriptions { + if memberAssignmentsIncludeTopicPartition(consumer2AllPotentialPartitions[anotherConsumer], partition) { + return s.reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, anotherConsumer) + } + } + return sortedCurrentSubscriptions +} + +// Reassign a specific partition to a new consumer +func (s *stickyBalanceStrategy) reassignPartition(partition topicPartitionAssignment, currentAssignment map[string][]topicPartitionAssignment, sortedCurrentSubscriptions []string, currentPartitionConsumer map[topicPartitionAssignment]string, newConsumer string) []string { + consumer := currentPartitionConsumer[partition] + // find the correct partition movement considering the stickiness requirement + partitionToBeMoved := s.movements.getTheActualPartitionToBeMoved(partition, consumer, newConsumer) + return s.processPartitionMovement(partitionToBeMoved, newConsumer, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer) +} + +// Track the movement of a topic partition after assignment +func (s *stickyBalanceStrategy) processPartitionMovement(partition topicPartitionAssignment, newConsumer string, currentAssignment map[string][]topicPartitionAssignment, sortedCurrentSubscriptions []string, currentPartitionConsumer map[topicPartitionAssignment]string) []string { + oldConsumer := currentPartitionConsumer[partition] + s.movements.movePartition(partition, oldConsumer, newConsumer) + + currentAssignment[oldConsumer] = removeTopicPartitionFromMemberAssignments(currentAssignment[oldConsumer], partition) + currentAssignment[newConsumer] = append(currentAssignment[newConsumer], partition) + currentPartitionConsumer[partition] = newConsumer + return sortMemberIDsByPartitionAssignments(currentAssignment) +} + +// Determine whether a specific consumer should be considered for topic partition assignment. +func canConsumerParticipateInReassignment(memberID string, currentAssignment map[string][]topicPartitionAssignment, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment, partition2AllPotentialConsumers map[topicPartitionAssignment][]string) bool { + currentPartitions := currentAssignment[memberID] + currentAssignmentSize := len(currentPartitions) + maxAssignmentSize := len(consumer2AllPotentialPartitions[memberID]) + if currentAssignmentSize > maxAssignmentSize { + Logger.Printf("The consumer %s is assigned more partitions than the maximum possible", memberID) + } + if currentAssignmentSize < maxAssignmentSize { + // if a consumer is not assigned all its potential partitions it is subject to reassignment + return true + } + for _, partition := range currentPartitions { + if canTopicPartitionParticipateInReassignment(partition, partition2AllPotentialConsumers) { + return true + } + } + return false +} + +// Only consider reassigning those topic partitions that have two or more potential consumers. +func canTopicPartitionParticipateInReassignment(partition topicPartitionAssignment, partition2AllPotentialConsumers map[topicPartitionAssignment][]string) bool { + return len(partition2AllPotentialConsumers[partition]) >= 2 +} + +// The assignment should improve the overall balance of the partition assignments to consumers. +func assignPartition(partition topicPartitionAssignment, sortedCurrentSubscriptions []string, currentAssignment map[string][]topicPartitionAssignment, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment, currentPartitionConsumer map[topicPartitionAssignment]string) []string { + updatedSubscriptions := make([]string, len(sortedCurrentSubscriptions)) + for i, s := range sortedCurrentSubscriptions { + updatedSubscriptions[i] = s + } + i := 0 + for _, memberID := range sortedCurrentSubscriptions { + if memberAssignmentsIncludeTopicPartition(consumer2AllPotentialPartitions[memberID], partition) { + updatedSubscriptions = removeIndexFromStringSlice(updatedSubscriptions, i) + currentAssignment[memberID] = append(currentAssignment[memberID], partition) + currentPartitionConsumer[partition] = memberID + updatedSubscriptions = append(updatedSubscriptions, memberID) + break + } + i++ + } + return updatedSubscriptions +} + +// Deserialize topic partition assignment data to aid with creation of a sticky assignment. +func deserializeTopicPartitionAssignment(userDataBytes []byte) (StickyAssignorUserData, error) { + userDataV1 := &StickyAssignorUserDataV1{} + if err := decode(userDataBytes, userDataV1); err != nil { + userDataV0 := &StickyAssignorUserDataV0{} + if err := decode(userDataBytes, userDataV0); err != nil { + return nil, err + } + return userDataV0, nil + } + return userDataV1, nil +} + +// filterAssignedPartitions returns a map of consumer group members to their list of previously-assigned topic partitions, limited +// to those topic partitions currently reported by the Kafka cluster. +func filterAssignedPartitions(currentAssignment map[string][]topicPartitionAssignment, partition2AllPotentialConsumers map[topicPartitionAssignment][]string) map[string][]topicPartitionAssignment { + assignments := deepCopyAssignment(currentAssignment) + for memberID, partitions := range assignments { + // perform in-place filtering + i := 0 + for _, partition := range partitions { + if _, exists := partition2AllPotentialConsumers[partition]; exists { + partitions[i] = partition + i++ + } + } + assignments[memberID] = partitions[:i] + } + return assignments +} + +func removeValueFromStringSlice(s []string, e string) []string { + for i, v := range s { + if v == e { + s = append(s[:i], s[i+1:]...) + break + } + } + return s +} + +func removeIndexFromStringSlice(s []string, i int) []string { + if len(s) == 0 { + return s + } + return append(s[:i], s[i+1:]...) +} + +func removeTopicPartitionFromMemberAssignments(assignments []topicPartitionAssignment, topic topicPartitionAssignment) []topicPartitionAssignment { + for i, assignment := range assignments { + if assignment == topic { + return append(assignments[:i], assignments[i+1:]...) + } + } + return assignments +} + +func memberAssignmentsIncludeTopicPartition(assignments []topicPartitionAssignment, topic topicPartitionAssignment) bool { + for _, assignment := range assignments { + if assignment == topic { + return true + } + } + return false +} + +func sortPartitions(currentAssignment map[string][]topicPartitionAssignment, partitionsWithADifferentPreviousAssignment map[topicPartitionAssignment]consumerGenerationPair, isFreshAssignment bool, partition2AllPotentialConsumers map[topicPartitionAssignment][]string, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment) []topicPartitionAssignment { + sortedPartitions := make([]topicPartitionAssignment, 0) + if !isFreshAssignment && areSubscriptionsIdentical(partition2AllPotentialConsumers, consumer2AllPotentialPartitions) { + // if this is a reassignment and the subscriptions are identical (all consumers can consumer from all topics) + // then we just need to simply list partitions in a round robin fashion (from consumers with + // most assigned partitions to those with least) + assignments := filterAssignedPartitions(currentAssignment, partition2AllPotentialConsumers) + + // sortedMemberIDs contains a descending-sorted list of consumers based on how many valid partitions are currently assigned to them + sortedMemberIDs := sortMemberIDsByPartitionAssignments(assignments) + for i := len(sortedMemberIDs)/2 - 1; i >= 0; i-- { + opp := len(sortedMemberIDs) - 1 - i + sortedMemberIDs[i], sortedMemberIDs[opp] = sortedMemberIDs[opp], sortedMemberIDs[i] + } + for { + // loop until no consumer-group members remain + if len(sortedMemberIDs) == 0 { + break + } + updatedMemberIDs := make([]string, 0) + for _, memberID := range sortedMemberIDs { + // partitions that were assigned to a different consumer last time + prevPartitions := make([]topicPartitionAssignment, 0) + for partition := range partitionsWithADifferentPreviousAssignment { + // from partitions that had a different consumer before, keep only those that are assigned to this consumer now + if memberAssignmentsIncludeTopicPartition(assignments[memberID], partition) { + prevPartitions = append(prevPartitions, partition) + } + } + + if len(prevPartitions) > 0 { + // if there is a partition on this consumer that was assigned to another consumer before mark it as good options for reassignment + partition := prevPartitions[0] + prevPartitions = append(prevPartitions[:0], prevPartitions[1:]...) + assignments[memberID] = removeTopicPartitionFromMemberAssignments(assignments[memberID], partition) + sortedPartitions = append(sortedPartitions, partition) + updatedMemberIDs = append(updatedMemberIDs, memberID) + } else if len(assignments[memberID]) > 0 { + // otherwise, mark any other one of the current partitions as a reassignment candidate + partition := assignments[memberID][0] + assignments[memberID] = append(assignments[memberID][:0], assignments[memberID][1:]...) + sortedPartitions = append(sortedPartitions, partition) + updatedMemberIDs = append(updatedMemberIDs, memberID) + } + } + sortedMemberIDs = updatedMemberIDs + } + + for partition := range partition2AllPotentialConsumers { + found := false + for _, p := range sortedPartitions { + if partition == p { + found = true + break + } + } + if !found { + sortedPartitions = append(sortedPartitions, partition) + } + } + } else { + // an ascending sorted set of topic partitions based on how many consumers can potentially use them + sortedPartitions = sortPartitionsByPotentialConsumerAssignments(partition2AllPotentialConsumers) + } + return sortedPartitions +} + +func sortMemberIDsByPartitionAssignments(assignments map[string][]topicPartitionAssignment) []string { + // sort the members by the number of partition assignments in ascending order + sortedMemberIDs := make([]string, 0, len(assignments)) + for memberID := range assignments { + sortedMemberIDs = append(sortedMemberIDs, memberID) + } + sort.SliceStable(sortedMemberIDs, func(i, j int) bool { + ret := len(assignments[sortedMemberIDs[i]]) - len(assignments[sortedMemberIDs[j]]) + if ret == 0 { + return sortedMemberIDs[i] < sortedMemberIDs[j] + } + return len(assignments[sortedMemberIDs[i]]) < len(assignments[sortedMemberIDs[j]]) + }) + return sortedMemberIDs +} + +func sortPartitionsByPotentialConsumerAssignments(partition2AllPotentialConsumers map[topicPartitionAssignment][]string) []topicPartitionAssignment { + // sort the members by the number of partition assignments in descending order + sortedPartionIDs := make([]topicPartitionAssignment, 0, len(partition2AllPotentialConsumers)) + for partition := range partition2AllPotentialConsumers { + sortedPartionIDs = append(sortedPartionIDs, partition) + } + sort.SliceStable(sortedPartionIDs, func(i, j int) bool { + return len(partition2AllPotentialConsumers[sortedPartionIDs[i]]) > len(partition2AllPotentialConsumers[sortedPartionIDs[j]]) + }) + return sortedPartionIDs +} + +func deepCopyPartitions(src []topicPartitionAssignment) []topicPartitionAssignment { + dst := make([]topicPartitionAssignment, len(src)) + for i, partition := range src { + dst[i] = partition + } + return dst +} + +func deepCopyAssignment(assignment map[string][]topicPartitionAssignment) map[string][]topicPartitionAssignment { + copy := make(map[string][]topicPartitionAssignment, len(assignment)) + for memberID, subscriptions := range assignment { + copy[memberID] = append(subscriptions[:0:0], subscriptions...) + } + return copy +} + +func areSubscriptionsIdentical(partition2AllPotentialConsumers map[topicPartitionAssignment][]string, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment) bool { + curMembers := make(map[string]int) + for _, cur := range partition2AllPotentialConsumers { + if len(curMembers) == 0 { + for _, curMembersElem := range cur { + curMembers[curMembersElem]++ + } + continue + } + + if len(curMembers) != len(cur) { + return false + } + + yMap := make(map[string]int) + for _, yElem := range cur { + yMap[yElem]++ + } + + for curMembersMapKey, curMembersMapVal := range curMembers { + if yMap[curMembersMapKey] != curMembersMapVal { + return false + } + } + } + + curPartitions := make(map[topicPartitionAssignment]int) + for _, cur := range consumer2AllPotentialPartitions { + if len(curPartitions) == 0 { + for _, curPartitionElem := range cur { + curPartitions[curPartitionElem]++ + } + continue + } + + if len(curPartitions) != len(cur) { + return false + } + + yMap := make(map[topicPartitionAssignment]int) + for _, yElem := range cur { + yMap[yElem]++ + } + + for curMembersMapKey, curMembersMapVal := range curPartitions { + if yMap[curMembersMapKey] != curMembersMapVal { + return false + } + } + } + return true +} + +// We need to process subscriptions' user data with each consumer's reported generation in mind +// higher generations overwrite lower generations in case of a conflict +// note that a conflict could exist only if user data is for different generations +func prepopulateCurrentAssignments(members map[string]ConsumerGroupMemberMetadata) (map[string][]topicPartitionAssignment, map[topicPartitionAssignment]consumerGenerationPair, error) { + currentAssignment := make(map[string][]topicPartitionAssignment) + prevAssignment := make(map[topicPartitionAssignment]consumerGenerationPair) + + // for each partition we create a sorted map of its consumers by generation + sortedPartitionConsumersByGeneration := make(map[topicPartitionAssignment]map[int]string) + for memberID, meta := range members { + consumerUserData, err := deserializeTopicPartitionAssignment(meta.UserData) + if err != nil { + return nil, nil, err + } + for _, partition := range consumerUserData.partitions() { + if consumers, exists := sortedPartitionConsumersByGeneration[partition]; exists { + if consumerUserData.hasGeneration() { + if _, generationExists := consumers[consumerUserData.generation()]; generationExists { + // same partition is assigned to two consumers during the same rebalance. + // log a warning and skip this record + Logger.Printf("Topic %s Partition %d is assigned to multiple consumers following sticky assignment generation %d", partition.Topic, partition.Partition, consumerUserData.generation()) + continue + } else { + consumers[consumerUserData.generation()] = memberID + } + } else { + consumers[defaultGeneration] = memberID + } + } else { + generation := defaultGeneration + if consumerUserData.hasGeneration() { + generation = consumerUserData.generation() + } + sortedPartitionConsumersByGeneration[partition] = map[int]string{generation: memberID} + } + } + } + + // prevAssignment holds the prior ConsumerGenerationPair (before current) of each partition + // current and previous consumers are the last two consumers of each partition in the above sorted map + for partition, consumers := range sortedPartitionConsumersByGeneration { + // sort consumers by generation in decreasing order + var generations []int + for generation := range consumers { + generations = append(generations, generation) + } + sort.Sort(sort.Reverse(sort.IntSlice(generations))) + + if len(generations) > 0 { + consumer := consumers[generations[0]] + currentConsumerAssignments, exists := currentAssignment[consumer] + if !exists { + currentConsumerAssignments = []topicPartitionAssignment{partition} + } else { + currentConsumerAssignments = append(currentConsumerAssignments, partition) + } + currentAssignment[consumer] = currentConsumerAssignments + + // check for previous assignment, if any + if len(generations) > 1 { + prevAssignment[partition] = consumerGenerationPair{ + MemberID: consumers[generations[1]], + Generation: generations[1], + } + } + } + } + return currentAssignment, prevAssignment, nil +} + +type consumerGenerationPair struct { + MemberID string + Generation int +} + +// consumerPair represents a pair of Kafka consumer ids involved in a partition reassignment. +type consumerPair struct { + SrcMemberID string + DstMemberID string +} + +// partitionMovements maintains some data structures to simplify lookup of partition movements among consumers. +type partitionMovements struct { + PartitionMovementsByTopic map[string]map[consumerPair]map[topicPartitionAssignment]bool + Movements map[topicPartitionAssignment]consumerPair +} + +func (p *partitionMovements) removeMovementRecordOfPartition(partition topicPartitionAssignment) consumerPair { + pair := p.Movements[partition] + delete(p.Movements, partition) + + partitionMovementsForThisTopic := p.PartitionMovementsByTopic[partition.Topic] + delete(partitionMovementsForThisTopic[pair], partition) + if len(partitionMovementsForThisTopic[pair]) == 0 { + delete(partitionMovementsForThisTopic, pair) + } + if len(p.PartitionMovementsByTopic[partition.Topic]) == 0 { + delete(p.PartitionMovementsByTopic, partition.Topic) + } + return pair +} + +func (p *partitionMovements) addPartitionMovementRecord(partition topicPartitionAssignment, pair consumerPair) { + p.Movements[partition] = pair + if _, exists := p.PartitionMovementsByTopic[partition.Topic]; !exists { + p.PartitionMovementsByTopic[partition.Topic] = make(map[consumerPair]map[topicPartitionAssignment]bool) + } + partitionMovementsForThisTopic := p.PartitionMovementsByTopic[partition.Topic] + if _, exists := partitionMovementsForThisTopic[pair]; !exists { + partitionMovementsForThisTopic[pair] = make(map[topicPartitionAssignment]bool) + } + partitionMovementsForThisTopic[pair][partition] = true +} + +func (p *partitionMovements) movePartition(partition topicPartitionAssignment, oldConsumer, newConsumer string) { + pair := consumerPair{ + SrcMemberID: oldConsumer, + DstMemberID: newConsumer, + } + if _, exists := p.Movements[partition]; exists { + // this partition has previously moved + existingPair := p.removeMovementRecordOfPartition(partition) + if existingPair.DstMemberID != oldConsumer { + Logger.Printf("Existing pair DstMemberID %s was not equal to the oldConsumer ID %s", existingPair.DstMemberID, oldConsumer) + } + if existingPair.SrcMemberID != newConsumer { + // the partition is not moving back to its previous consumer + p.addPartitionMovementRecord(partition, consumerPair{ + SrcMemberID: existingPair.SrcMemberID, + DstMemberID: newConsumer, + }) + } + } else { + p.addPartitionMovementRecord(partition, pair) + } +} + +func (p *partitionMovements) getTheActualPartitionToBeMoved(partition topicPartitionAssignment, oldConsumer, newConsumer string) topicPartitionAssignment { + if _, exists := p.PartitionMovementsByTopic[partition.Topic]; !exists { + return partition + } + if _, exists := p.Movements[partition]; exists { + // this partition has previously moved + if oldConsumer != p.Movements[partition].DstMemberID { + Logger.Printf("Partition movement DstMemberID %s was not equal to the oldConsumer ID %s", p.Movements[partition].DstMemberID, oldConsumer) + } + oldConsumer = p.Movements[partition].SrcMemberID + } + + partitionMovementsForThisTopic := p.PartitionMovementsByTopic[partition.Topic] + reversePair := consumerPair{ + SrcMemberID: newConsumer, + DstMemberID: oldConsumer, + } + if _, exists := partitionMovementsForThisTopic[reversePair]; !exists { + return partition + } + var reversePairPartition topicPartitionAssignment + for otherPartition := range partitionMovementsForThisTopic[reversePair] { + reversePairPartition = otherPartition + } + return reversePairPartition +} + +func (p *partitionMovements) isLinked(src, dst string, pairs []consumerPair, currentPath []string) ([]string, bool) { + if src == dst { + return currentPath, false + } + if len(pairs) == 0 { + return currentPath, false + } + for _, pair := range pairs { + if src == pair.SrcMemberID && dst == pair.DstMemberID { + currentPath = append(currentPath, src, dst) + return currentPath, true + } + } + + for _, pair := range pairs { + if pair.SrcMemberID == src { + // create a deep copy of the pairs, excluding the current pair + reducedSet := make([]consumerPair, len(pairs)-1) + i := 0 + for _, p := range pairs { + if p != pair { + reducedSet[i] = pair + i++ + } + } + + currentPath = append(currentPath, pair.SrcMemberID) + return p.isLinked(pair.DstMemberID, dst, reducedSet, currentPath) + } + } + return currentPath, false +} + +func (p *partitionMovements) in(cycle []string, cycles [][]string) bool { + superCycle := make([]string, len(cycle)-1) + for i := 0; i < len(cycle)-1; i++ { + superCycle[i] = cycle[i] + } + for _, c := range cycle { + superCycle = append(superCycle, c) + } + for _, foundCycle := range cycles { + if len(foundCycle) == len(cycle) && indexOfSubList(superCycle, foundCycle) != -1 { + return true + } + } + return false +} + +func (p *partitionMovements) hasCycles(pairs []consumerPair) bool { + cycles := make([][]string, 0) + for _, pair := range pairs { + // create a deep copy of the pairs, excluding the current pair + reducedPairs := make([]consumerPair, len(pairs)-1) + i := 0 + for _, p := range pairs { + if p != pair { + reducedPairs[i] = pair + i++ + } + } + if path, linked := p.isLinked(pair.DstMemberID, pair.SrcMemberID, reducedPairs, []string{pair.SrcMemberID}); linked { + if !p.in(path, cycles) { + cycles = append(cycles, path) + Logger.Printf("A cycle of length %d was found: %v", len(path)-1, path) + } + } + } + + // for now we want to make sure there is no partition movements of the same topic between a pair of consumers. + // the odds of finding a cycle among more than two consumers seem to be very low (according to various randomized + // tests with the given sticky algorithm) that it should not worth the added complexity of handling those cases. + for _, cycle := range cycles { + if len(cycle) == 3 { + return true + } + } + return false +} + +func (p *partitionMovements) isSticky() bool { + for topic, movements := range p.PartitionMovementsByTopic { + movementPairs := make([]consumerPair, len(movements)) + i := 0 + for pair := range movements { + movementPairs[i] = pair + i++ + } + if p.hasCycles(movementPairs) { + Logger.Printf("Stickiness is violated for topic %s", topic) + Logger.Printf("Partition movements for this topic occurred among the following consumer pairs: %v", movements) + return false + } + } + return true +} + +func indexOfSubList(source []string, target []string) int { + targetSize := len(target) + maxCandidate := len(source) - targetSize +nextCand: + for candidate := 0; candidate <= maxCandidate; candidate++ { + j := candidate + for i := 0; i < targetSize; i++ { + if target[i] != source[j] { + // Element mismatch, try next cand + continue nextCand + } + j++ + } + // All elements of candidate matched target + return candidate + } + return -1 +} diff --git a/balance_strategy_test.go b/balance_strategy_test.go index 047157f37..ef23d0ebf 100644 --- a/balance_strategy_test.go +++ b/balance_strategy_test.go @@ -1,8 +1,13 @@ package sarama import ( + "fmt" + "math" + "math/rand" "reflect" + "sort" "testing" + "time" ) func TestBalanceStrategyRange(t *testing.T) { @@ -100,3 +105,2174 @@ func TestBalanceStrategyRoundRobin(t *testing.T) { } } } + +func Test_deserializeTopicPartitionAssignment(t *testing.T) { + type args struct { + userDataBytes []byte + } + tests := []struct { + name string + args args + want StickyAssignorUserData + wantErr bool + }{ + { + name: "Nil userdata bytes", + args: args{}, + want: &StickyAssignorUserDataV1{}, + }, + { + name: "Non-empty invalid userdata bytes", + args: args{ + userDataBytes: []byte{ + 0x00, 0x00, + 0x00, 0x00, 0x00, 0x01, + 0x00, 0x03, 'f', 'o', 'o', + }, + }, + wantErr: true, + }, + { + name: "Valid v0 userdata bytes", + args: args{ + userDataBytes: []byte{ + 0x00, 0x00, 0x00, 0x01, 0x00, 0x03, 0x74, 0x30, + 0x33, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, + 0x05, + }, + }, + want: &StickyAssignorUserDataV0{ + Topics: map[string][]int32{"t03": {5}}, + topicPartitions: []topicPartitionAssignment{ + topicPartitionAssignment{ + Topic: "t03", + Partition: 5, + }, + }, + }, + }, + { + name: "Valid v1 userdata bytes", + args: args{ + userDataBytes: []byte{ + 0x00, 0x00, 0x00, 0x01, 0x00, 0x03, 0x74, 0x30, + 0x36, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x04, 0xff, 0xff, 0xff, + 0xff, + }, + }, + want: &StickyAssignorUserDataV1{ + Topics: map[string][]int32{"t06": {0, 4}}, + Generation: -1, + topicPartitions: []topicPartitionAssignment{ + topicPartitionAssignment{ + Topic: "t06", + Partition: 0, + }, + topicPartitionAssignment{ + Topic: "t06", + Partition: 4, + }, + }, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := deserializeTopicPartitionAssignment(tt.args.userDataBytes) + if (err != nil) != tt.wantErr { + t.Errorf("deserializeTopicPartitionAssignment() error = %v, wantErr %v", err, tt.wantErr) + return + } + if !reflect.DeepEqual(got, tt.want) { + t.Errorf("deserializeTopicPartitionAssignment() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_prepopulateCurrentAssignments(t *testing.T) { + type args struct { + members map[string]ConsumerGroupMemberMetadata + } + tests := []struct { + name string + args args + wantCurrentAssignments map[string][]topicPartitionAssignment + wantPrevAssignments map[topicPartitionAssignment]consumerGenerationPair + wantErr bool + }{ + { + name: "Empty map", + wantCurrentAssignments: map[string][]topicPartitionAssignment{}, + wantPrevAssignments: map[topicPartitionAssignment]consumerGenerationPair{}, + }, + { + name: "Single consumer", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "c01": ConsumerGroupMemberMetadata{ + Version: 2, + UserData: []byte{ + 0x00, 0x00, 0x00, 0x01, 0x00, 0x03, 0x74, 0x30, + 0x36, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x04, 0xff, 0xff, 0xff, + 0xff, + }, + }, + }, + }, + wantCurrentAssignments: map[string][]topicPartitionAssignment{ + "c01": []topicPartitionAssignment{ + topicPartitionAssignment{ + Topic: "t06", + Partition: 0, + }, + topicPartitionAssignment{ + Topic: "t06", + Partition: 4, + }, + }, + }, + wantPrevAssignments: map[topicPartitionAssignment]consumerGenerationPair{}, + }, + { + name: "Duplicate consumer assignments in metadata", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "c01": ConsumerGroupMemberMetadata{ + Version: 2, + UserData: []byte{ + 0x00, 0x00, 0x00, 0x01, 0x00, 0x03, 0x74, 0x30, + 0x36, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x04, 0xff, 0xff, 0xff, + 0xff, + }, + }, + "c02": ConsumerGroupMemberMetadata{ + Version: 2, + UserData: []byte{ + 0x00, 0x00, 0x00, 0x01, 0x00, 0x03, 0x74, 0x30, + 0x36, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x04, 0xff, 0xff, 0xff, + 0xff, + }, + }, + }, + }, + wantCurrentAssignments: map[string][]topicPartitionAssignment{ + "c01": []topicPartitionAssignment{ + topicPartitionAssignment{ + Topic: "t06", + Partition: 0, + }, + topicPartitionAssignment{ + Topic: "t06", + Partition: 4, + }, + }, + }, + wantPrevAssignments: map[topicPartitionAssignment]consumerGenerationPair{}, + }, + { + name: "Different generations (5, 6) of consumer assignments in metadata", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "c01": ConsumerGroupMemberMetadata{ + Version: 2, + UserData: []byte{ + 0x00, 0x00, 0x00, 0x01, 0x00, 0x03, 0x74, 0x30, + 0x36, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, + 0x05, + }, + }, + "c02": ConsumerGroupMemberMetadata{ + Version: 2, + UserData: []byte{ + 0x00, 0x00, 0x00, 0x01, 0x00, 0x03, 0x74, 0x30, + 0x36, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, + 0x06, + }, + }, + }, + }, + wantCurrentAssignments: map[string][]topicPartitionAssignment{ + "c01": []topicPartitionAssignment{ + topicPartitionAssignment{ + Topic: "t06", + Partition: 0, + }, + topicPartitionAssignment{ + Topic: "t06", + Partition: 4, + }, + }, + }, + wantPrevAssignments: map[topicPartitionAssignment]consumerGenerationPair{ + topicPartitionAssignment{ + Topic: "t06", + Partition: 0, + }: consumerGenerationPair{ + Generation: 5, + MemberID: "c01", + }, + topicPartitionAssignment{ + Topic: "t06", + Partition: 4, + }: consumerGenerationPair{ + Generation: 5, + MemberID: "c01", + }, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, gotPrevAssignments, err := prepopulateCurrentAssignments(tt.args.members) + + if (err != nil) != tt.wantErr { + t.Errorf("prepopulateCurrentAssignments() error = %v, wantErr %v", err, tt.wantErr) + } + + if !reflect.DeepEqual(gotPrevAssignments, tt.wantPrevAssignments) { + t.Errorf("deserializeTopicPartitionAssignment() prevAssignments = %v, want %v", gotPrevAssignments, tt.wantPrevAssignments) + } + }) + } +} + +func Test_areSubscriptionsIdentical(t *testing.T) { + type args struct { + partition2AllPotentialConsumers map[topicPartitionAssignment][]string + consumer2AllPotentialPartitions map[string][]topicPartitionAssignment + } + tests := []struct { + name string + args args + want bool + }{ + { + name: "Empty consumers and partitions", + args: args{ + partition2AllPotentialConsumers: make(map[topicPartitionAssignment][]string), + consumer2AllPotentialPartitions: make(map[string][]topicPartitionAssignment), + }, + want: true, + }, + { + name: "Topic partitions with identical consumer entries", + args: args{ + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c1", "c2", "c3"}, + }, + consumer2AllPotentialPartitions: make(map[string][]topicPartitionAssignment), + }, + want: true, + }, + { + name: "Topic partitions with mixed up consumer entries", + args: args{ + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c2", "c3", "c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c3", "c1", "c2"}, + }, + consumer2AllPotentialPartitions: make(map[string][]topicPartitionAssignment), + }, + want: true, + }, + { + name: "Topic partitions with different consumer entries", + args: args{ + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c2", "c3", "c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"cX", "c1", "c2"}, + }, + consumer2AllPotentialPartitions: make(map[string][]topicPartitionAssignment), + }, + want: false, + }, + { + name: "Topic partitions with different number of consumer entries", + args: args{ + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c2", "c3", "c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c1", "c2"}, + }, + consumer2AllPotentialPartitions: make(map[string][]topicPartitionAssignment), + }, + want: false, + }, + { + name: "Consumers with identical topic partitions", + args: args{ + partition2AllPotentialConsumers: make(map[topicPartitionAssignment][]string), + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + }, + }, + want: true, + }, + { + name: "Consumer2 with mixed up consumer entries", + args: args{ + partition2AllPotentialConsumers: make(map[topicPartitionAssignment][]string), + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}, topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 2}, topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}}, + }, + }, + want: true, + }, + { + name: "Consumer2 with different consumer entries", + args: args{ + partition2AllPotentialConsumers: make(map[topicPartitionAssignment][]string), + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}, topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "tX", Partition: 2}, topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}}, + }, + }, + want: false, + }, + { + name: "Consumer2 with different number of consumer entries", + args: args{ + partition2AllPotentialConsumers: make(map[topicPartitionAssignment][]string), + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}, topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}}, + }, + }, + want: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := areSubscriptionsIdentical(tt.args.partition2AllPotentialConsumers, tt.args.consumer2AllPotentialPartitions); got != tt.want { + t.Errorf("areSubscriptionsIdentical() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_sortMemberIDsByPartitionAssignments(t *testing.T) { + type args struct { + assignments map[string][]topicPartitionAssignment + } + tests := []struct { + name string + args args + want []string + }{ + { + name: "Null assignments", + want: make([]string, 0), + }, + { + name: "Single assignment", + args: args{ + assignments: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + }, + }, + want: []string{"c1"}, + }, + { + name: "Multiple assignments with different partition counts", + args: args{ + assignments: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + "c2": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + "c3": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 3}, + topicPartitionAssignment{Topic: "t1", Partition: 4}, + topicPartitionAssignment{Topic: "t1", Partition: 5}, + }, + }, + }, + want: []string{"c1", "c2", "c3"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := sortMemberIDsByPartitionAssignments(tt.args.assignments); !reflect.DeepEqual(got, tt.want) { + t.Errorf("sortMemberIDsByPartitionAssignments() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_sortPartitions(t *testing.T) { + type args struct { + currentAssignment map[string][]topicPartitionAssignment + partitionsWithADifferentPreviousAssignment map[topicPartitionAssignment]consumerGenerationPair + isFreshAssignment bool + partition2AllPotentialConsumers map[topicPartitionAssignment][]string + consumer2AllPotentialPartitions map[string][]topicPartitionAssignment + } + tests := []struct { + name string + args args + want []topicPartitionAssignment + }{ + { + name: "Empty everything", + want: make([]topicPartitionAssignment, 0), + }, + { + name: "Base case", + args: args{ + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 2}}, + }, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c2", "c3", "c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c3", "c1", "c2"}, + }, + }, + }, + { + name: "Partitions assigned to a different consumer last time", + args: args{ + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + }, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c2", "c3", "c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c3", "c1", "c2"}, + }, + partitionsWithADifferentPreviousAssignment: map[topicPartitionAssignment]consumerGenerationPair{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: consumerGenerationPair{Generation: 1, MemberID: "c2"}, + }, + }, + }, + { + name: "Partitions assigned to a different consumer last time", + args: args{ + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}}, + }, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c2", "c3", "c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c3", "c1", "c2"}, + }, + partitionsWithADifferentPreviousAssignment: map[topicPartitionAssignment]consumerGenerationPair{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: consumerGenerationPair{Generation: 1, MemberID: "c2"}, + }, + }, + }, + { + name: "Fresh assignment", + args: args{ + isFreshAssignment: true, + currentAssignment: map[string][]topicPartitionAssignment{}, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}, topicPartitionAssignment{Topic: "t1", Partition: 1}, topicPartitionAssignment{Topic: "t1", Partition: 2}}, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2", "c3"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c2", "c3", "c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c3", "c1", "c2"}, + }, + partitionsWithADifferentPreviousAssignment: map[topicPartitionAssignment]consumerGenerationPair{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: consumerGenerationPair{Generation: 1, MemberID: "c2"}, + }, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := sortPartitions(tt.args.currentAssignment, tt.args.partitionsWithADifferentPreviousAssignment, tt.args.isFreshAssignment, tt.args.partition2AllPotentialConsumers, tt.args.consumer2AllPotentialPartitions) + if tt.want != nil && !reflect.DeepEqual(got, tt.want) { + t.Errorf("sortPartitions() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_filterAssignedPartitions(t *testing.T) { + type args struct { + currentAssignment map[string][]topicPartitionAssignment + partition2AllPotentialConsumers map[topicPartitionAssignment][]string + } + tests := []struct { + name string + args args + want map[string][]topicPartitionAssignment + }{ + { + name: "All partitions accounted for", + args: args{ + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}}, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c2"}, + }, + }, + want: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}}, + }, + }, + { + name: "One consumer using an unrecognized partition", + args: args{ + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}}, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1"}, + }, + }, + want: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c2": []topicPartitionAssignment{}, + }, + }, + { + name: "Interleaved consumer removal", + args: args{ + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c2": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 1}}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 2}}, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c3"}, + }, + }, + want: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 0}}, + "c2": []topicPartitionAssignment{}, + "c3": []topicPartitionAssignment{topicPartitionAssignment{Topic: "t1", Partition: 2}}, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := filterAssignedPartitions(tt.args.currentAssignment, tt.args.partition2AllPotentialConsumers); !reflect.DeepEqual(got, tt.want) { + t.Errorf("filterAssignedPartitions() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_canConsumerParticipateInReassignment(t *testing.T) { + type args struct { + memberID string + currentAssignment map[string][]topicPartitionAssignment + consumer2AllPotentialPartitions map[string][]topicPartitionAssignment + partition2AllPotentialConsumers map[topicPartitionAssignment][]string + } + tests := []struct { + name string + args args + want bool + }{ + { + name: "Consumer has been assigned partitions not available to it", + args: args{ + memberID: "c1", + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + "c2": []topicPartitionAssignment{}, + }, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + "c2": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1", "c2"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c1", "c2"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c2"}, + }, + }, + want: true, + }, + { + name: "Consumer has been assigned all available partitions", + args: args{ + memberID: "c1", + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + }, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c1"}, + }, + }, + want: false, + }, + { + name: "Consumer has not been assigned all available partitions", + args: args{ + memberID: "c1", + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + }, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + }, + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: []string{"c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 1}: []string{"c1"}, + topicPartitionAssignment{Topic: "t1", Partition: 2}: []string{"c1"}, + }, + }, + want: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := canConsumerParticipateInReassignment(tt.args.memberID, tt.args.currentAssignment, tt.args.consumer2AllPotentialPartitions, tt.args.partition2AllPotentialConsumers); got != tt.want { + t.Errorf("canConsumerParticipateInReassignment() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_removeTopicPartitionFromMemberAssignments(t *testing.T) { + type args struct { + assignments []topicPartitionAssignment + topic topicPartitionAssignment + } + tests := []struct { + name string + args args + want []topicPartitionAssignment + }{ + { + name: "Empty", + args: args{ + assignments: make([]topicPartitionAssignment, 0), + topic: topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + want: make([]topicPartitionAssignment, 0), + }, + { + name: "Remove first entry", + args: args{ + assignments: []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + topic: topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + want: []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + }, + { + name: "Remove middle entry", + args: args{ + assignments: []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + topic: topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + want: []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + }, + { + name: "Remove last entry", + args: args{ + assignments: []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + topic: topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + want: []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := removeTopicPartitionFromMemberAssignments(tt.args.assignments, tt.args.topic); !reflect.DeepEqual(got, tt.want) { + t.Errorf("removeTopicPartitionFromMemberAssignments() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_removeIndexFromStringSlice(t *testing.T) { + type args struct { + s []string + i int + } + tests := []struct { + name string + args args + want []string + }{ + { + name: "Empty slice", + args: args{ + s: make([]string, 0), + i: 0, + }, + want: make([]string, 0), + }, + { + name: "Slice with single entry", + args: args{ + s: []string{"foo"}, + i: 0, + }, + want: make([]string, 0), + }, + { + name: "Slice with multiple entries", + args: args{ + s: []string{"a", "b", "c"}, + i: 0, + }, + want: []string{"b", "c"}, + }, + { + name: "Slice with multiple entries and index is in the middle", + args: args{ + s: []string{"a", "b", "c"}, + i: 1, + }, + want: []string{"a", "c"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := removeIndexFromStringSlice(tt.args.s, tt.args.i); !reflect.DeepEqual(got, tt.want) { + t.Errorf("removeIndexFromSlice() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_removeValueFromStringSlice(t *testing.T) { + type args struct { + s []string + e string + } + tests := []struct { + name string + args args + want []string + }{ + { + name: "Empty input slice", + args: args{ + s: []string{}, + e: "", + }, + want: []string{}, + }, + { + name: "Input slice with one entry that doesn't match", + args: args{ + s: []string{"a"}, + e: "b", + }, + want: []string{"a"}, + }, + { + name: "Input slice with multiple entries and a positive match", + args: args{ + s: []string{"a", "b", "c"}, + e: "b", + }, + want: []string{"a", "c"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := removeValueFromStringSlice(tt.args.s, tt.args.e); !reflect.DeepEqual(got, tt.want) { + t.Errorf("removeValueFromSlice() = %v, want %v", got, tt.want) + } + }) + } +} + +func Test_assignPartition(t *testing.T) { + type args struct { + partition topicPartitionAssignment + sortedCurrentSubscriptions []string + currentAssignment map[string][]topicPartitionAssignment + consumer2AllPotentialPartitions map[string][]topicPartitionAssignment + currentPartitionConsumer map[topicPartitionAssignment]string + } + tests := []struct { + name string + args args + want []string + wantCurrentAssignment map[string][]topicPartitionAssignment + wantCurrentPartitionConsumer map[topicPartitionAssignment]string + }{ + { + name: "Base", + args: args{ + partition: topicPartitionAssignment{Topic: "t1", Partition: 2}, + sortedCurrentSubscriptions: []string{"c3", "c1", "c2"}, + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + "c2": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + "c3": []topicPartitionAssignment{}, + }, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + "c2": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + "c3": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + }, + currentPartitionConsumer: map[topicPartitionAssignment]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: "c1", + topicPartitionAssignment{Topic: "t1", Partition: 1}: "c2", + }, + }, + want: []string{"c1", "c2", "c3"}, + wantCurrentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + "c2": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + "c3": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + }, + wantCurrentPartitionConsumer: map[topicPartitionAssignment]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: "c1", + topicPartitionAssignment{Topic: "t1", Partition: 1}: "c2", + topicPartitionAssignment{Topic: "t1", Partition: 2}: "c3", + }, + }, + { + name: "Unassignable Partition", + args: args{ + partition: topicPartitionAssignment{Topic: "t1", Partition: 3}, + sortedCurrentSubscriptions: []string{"c3", "c1", "c2"}, + currentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + "c2": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + "c3": []topicPartitionAssignment{}, + }, + consumer2AllPotentialPartitions: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + "c2": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + "c3": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 2}, + }, + }, + currentPartitionConsumer: map[topicPartitionAssignment]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: "c1", + topicPartitionAssignment{Topic: "t1", Partition: 1}: "c2", + }, + }, + want: []string{"c3", "c1", "c2"}, + wantCurrentAssignment: map[string][]topicPartitionAssignment{ + "c1": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 0}, + }, + "c2": []topicPartitionAssignment{ + topicPartitionAssignment{Topic: "t1", Partition: 1}, + }, + "c3": []topicPartitionAssignment{}, + }, + wantCurrentPartitionConsumer: map[topicPartitionAssignment]string{ + topicPartitionAssignment{Topic: "t1", Partition: 0}: "c1", + topicPartitionAssignment{Topic: "t1", Partition: 1}: "c2", + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := assignPartition(tt.args.partition, tt.args.sortedCurrentSubscriptions, tt.args.currentAssignment, tt.args.consumer2AllPotentialPartitions, tt.args.currentPartitionConsumer); !reflect.DeepEqual(got, tt.want) { + t.Errorf("assignPartition() = %v, want %v", got, tt.want) + } + if !reflect.DeepEqual(tt.args.currentAssignment, tt.wantCurrentAssignment) { + t.Errorf("assignPartition() currentAssignment = %v, want %v", tt.args.currentAssignment, tt.wantCurrentAssignment) + } + if !reflect.DeepEqual(tt.args.currentPartitionConsumer, tt.wantCurrentPartitionConsumer) { + t.Errorf("assignPartition() currentPartitionConsumer = %v, want %v", tt.args.currentPartitionConsumer, tt.wantCurrentPartitionConsumer) + } + }) + } +} + +func Test_stickyBalanceStrategy_Plan(t *testing.T) { + type args struct { + members map[string]ConsumerGroupMemberMetadata + topics map[string][]int32 + } + tests := []struct { + name string + s *stickyBalanceStrategy + args args + }{ + { + name: "One consumer with no topics", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer": ConsumerGroupMemberMetadata{}, + }, + topics: make(map[string][]int32), + }, + }, + { + name: "One consumer with non-existent topic", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer": ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + }, + }, + topics: map[string][]int32{ + "topic": make([]int32, 0), + }, + }, + }, + { + name: "One consumer with one topic", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer": ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + }, + }, + topics: map[string][]int32{ + "topic": []int32{0, 1, 2}, + }, + }, + }, + { + name: "Only assigns partitions from subscribed topics", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer": ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + }, + }, + topics: map[string][]int32{ + "topic": []int32{0, 1, 2}, + "other": []int32{0, 1, 2}, + }, + }, + }, + { + name: "One consumer with multiple topics", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + }, + }, + topics: map[string][]int32{ + "topic1": []int32{0}, + "topic2": []int32{0, 1}, + }, + }, + }, + { + name: "Two consumers with one topic and one partition", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + }, + }, + topics: map[string][]int32{ + "topic": []int32{0}, + }, + }, + }, + { + name: "Two consumers with one topic and two partitions", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + }, + }, + topics: map[string][]int32{ + "topic": []int32{0, 1}, + }, + }, + }, + { + name: "Multiple consumers with mixed topic subscriptions", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + }, + "consumer3": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + }, + }, + topics: map[string][]int32{ + "topic1": []int32{0, 1, 2}, + "topic2": []int32{0, 1}, + }, + }, + }, + { + name: "Two consumers with two topics and six partitions", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + }, + }, + topics: map[string][]int32{ + "topic1": []int32{0, 1, 2}, + "topic2": []int32{0, 1, 2}, + }, + }, + }, + { + name: "Three consumers (two old, one new) with one topic and twelve partitions", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{4, 11, 8, 5, 9, 2}}, 1), + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{1, 3, 0, 7, 10, 6}}, 1), + }, + "consumer3": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + }, + }, + topics: map[string][]int32{ + "topic1": []int32{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, + }, + }, + }, + { + name: "Three consumers (two old, one new) with one topic and 13 partitions", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{4, 11, 8, 5, 9, 2, 6}}, 1), + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{1, 3, 0, 7, 10, 12}}, 1), + }, + "consumer3": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + }, + }, + topics: map[string][]int32{ + "topic1": []int32{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, + }, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := &stickyBalanceStrategy{} + plan, err := s.Plan(tt.args.members, tt.args.topics) + verifyPlanIsBalancedAndSticky(t, s, tt.args.members, plan, err) + verifyFullyBalanced(t, plan) + }) + } +} + +func Test_stickyBalanceStrategy_Plan_KIP54_ExampleOne(t *testing.T) { + s := &stickyBalanceStrategy{} + + // PLAN 1 + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3", "topic4"}, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3", "topic4"}, + }, + "consumer3": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3", "topic4"}, + }, + } + topics := map[string][]int32{ + "topic1": []int32{0, 1}, + "topic2": []int32{0, 1}, + "topic3": []int32{0, 1}, + "topic4": []int32{0, 1}, + } + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + verifyFullyBalanced(t, plan1) + + // PLAN 2 + delete(members, "consumer1") + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3", "topic4"}, + UserData: encodeSubscriberPlan(t, plan1["consumer2"]), + } + members["consumer3"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3", "topic4"}, + UserData: encodeSubscriberPlan(t, plan1["consumer3"]), + } + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) + verifyFullyBalanced(t, plan2) +} + +func Test_stickyBalanceStrategy_Plan_KIP54_ExampleTwo(t *testing.T) { + s := &stickyBalanceStrategy{} + + // PLAN 1 + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + }, + "consumer3": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3"}, + }, + } + topics := map[string][]int32{ + "topic1": []int32{0}, + "topic2": []int32{0, 1}, + "topic3": []int32{0, 1, 2}, + } + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + if len(plan1["consumer1"]["topic1"]) != 1 || len(plan1["consumer2"]["topic2"]) != 2 || len(plan1["consumer3"]["topic3"]) != 3 { + t.Error("Incorrect distribution of topic partition assignments") + } + + // PLAN 2 + delete(members, "consumer1") + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: members["consumer2"].Topics, + UserData: encodeSubscriberPlan(t, plan1["consumer2"]), + } + members["consumer3"] = ConsumerGroupMemberMetadata{ + Topics: members["consumer3"].Topics, + UserData: encodeSubscriberPlan(t, plan1["consumer3"]), + } + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) + verifyFullyBalanced(t, plan2) + if len(plan2["consumer2"]["topic1"]) != 1 || len(plan2["consumer2"]["topic2"]) != 2 || len(plan2["consumer3"]["topic3"]) != 3 { + t.Error("Incorrect distribution of topic partition assignments") + } +} + +func Test_stickyBalanceStrategy_Plan_KIP54_ExampleThree(t *testing.T) { + s := &stickyBalanceStrategy{} + topicNames := []string{"topic1", "topic2"} + + // PLAN 1 + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: topicNames, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: topicNames, + }, + } + topics := map[string][]int32{ + "topic1": []int32{0, 1}, + "topic2": []int32{0, 1}, + } + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + + // PLAN 2 + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: topicNames, + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: topicNames, + UserData: encodeSubscriberPlan(t, plan1["consumer2"]), + } + members["consumer3"] = ConsumerGroupMemberMetadata{ + Topics: topicNames, + UserData: encodeSubscriberPlan(t, plan1["consumer3"]), + } + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) + verifyFullyBalanced(t, plan2) +} + +func Test_stickyBalanceStrategy_Plan_AddRemoveConsumerOneTopic(t *testing.T) { + s := &stickyBalanceStrategy{} + + // PLAN 1 + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + }, + } + topics := map[string][]int32{ + "topic": []int32{0, 1, 2}, + } + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + + // PLAN 2 + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + UserData: encodeSubscriberPlan(t, plan1["consumer1"]), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + } + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) + + // PLAN 3 + delete(members, "consumer1") + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic"}, + UserData: encodeSubscriberPlan(t, plan2["consumer2"]), + } + plan3, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan3, err) +} + +func Test_stickyBalanceStrategy_Plan_PoorRoundRobinAssignmentScenario(t *testing.T) { + s := &stickyBalanceStrategy{} + + // PLAN 1 + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3", "topic4", "topic5"}, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic3", "topic5"}, + }, + "consumer3": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic3", "topic5"}, + }, + "consumer4": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3", "topic4", "topic5"}, + }, + } + topics := make(map[string][]int32, 5) + for i := 1; i <= 5; i++ { + partitions := make([]int32, i%2+1) + for j := 0; j < i%2+1; j++ { + partitions[j] = int32(j) + } + topics[fmt.Sprintf("topic%d", i)] = partitions + } + + plan, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan, err) +} + +func Test_stickyBalanceStrategy_Plan_AddRemoveTopicTwoConsumers(t *testing.T) { + s := &stickyBalanceStrategy{} + + // PLAN 1 + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + }, + } + topics := map[string][]int32{ + "topic1": []int32{0, 1, 2}, + } + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + verifyFullyBalanced(t, plan1) + + // PLAN 2 + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + UserData: encodeSubscriberPlan(t, plan1["consumer1"]), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + UserData: encodeSubscriberPlan(t, plan1["consumer2"]), + } + topics["topic2"] = []int32{0, 1, 2} + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) + verifyFullyBalanced(t, plan2) + + // PLAN 3 + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + UserData: encodeSubscriberPlan(t, plan2["consumer1"]), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + UserData: encodeSubscriberPlan(t, plan2["consumer2"]), + } + delete(topics, "topic1") + + plan3, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan3, err) + verifyFullyBalanced(t, plan3) +} + +func Test_stickyBalanceStrategy_Plan_ReassignmentAfterOneConsumerLeaves(t *testing.T) { + s := &stickyBalanceStrategy{} + + // PLAN 1 + members := make(map[string]ConsumerGroupMemberMetadata, 20) + for i := 0; i < 20; i++ { + topics := make([]string, 20) + for j := 0; j < 20; j++ { + topics[j] = fmt.Sprintf("topic%d", j) + } + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{Topics: topics} + } + topics := make(map[string][]int32, 20) + for i := 0; i < 20; i++ { + partitions := make([]int32, 20) + for j := 0; j < 20; j++ { + partitions[j] = int32(j) + } + topics[fmt.Sprintf("topic%d", i)] = partitions + } + + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + + for i := 0; i < 20; i++ { + topics := make([]string, 20) + for j := 0; j < 20; j++ { + topics[j] = fmt.Sprintf("topic%d", j) + } + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{ + Topics: members[fmt.Sprintf("consumer%d", i)].Topics, + UserData: encodeSubscriberPlan(t, plan1[fmt.Sprintf("consumer%d", i)]), + } + } + delete(members, "consumer10") + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) +} + +func Test_stickyBalanceStrategy_Plan_ReassignmentAfterOneConsumerAdded(t *testing.T) { + s := &stickyBalanceStrategy{} + + // PLAN 1 + members := make(map[string]ConsumerGroupMemberMetadata) + for i := 0; i < 10; i++ { + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{Topics: []string{"topic1"}} + } + partitions := make([]int32, 20) + for j := 0; j < 20; j++ { + partitions[j] = int32(j) + } + topics := map[string][]int32{"topic1": partitions} + + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + + // add a new consumer + members["consumer10"] = ConsumerGroupMemberMetadata{Topics: []string{"topic1"}} + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) +} + +func Test_stickyBalanceStrategy_Plan_SameSubscriptions(t *testing.T) { + s := &stickyBalanceStrategy{} + + // PLAN 1 + members := make(map[string]ConsumerGroupMemberMetadata, 20) + for i := 0; i < 9; i++ { + topics := make([]string, 15) + for j := 0; j < 15; j++ { + topics[j] = fmt.Sprintf("topic%d", j) + } + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{Topics: topics} + } + topics := make(map[string][]int32, 15) + for i := 0; i < 15; i++ { + partitions := make([]int32, i) + for j := 0; j < i; j++ { + partitions[j] = int32(j) + } + topics[fmt.Sprintf("topic%d", i)] = partitions + } + + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + + // PLAN 2 + for i := 0; i < 9; i++ { + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{ + Topics: members[fmt.Sprintf("consumer%d", i)].Topics, + UserData: encodeSubscriberPlan(t, plan1[fmt.Sprintf("consumer%d", i)]), + } + } + delete(members, "consumer5") + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) +} + +func Test_stickyBalanceStrategy_Plan_LargeAssignmentWithMultipleConsumersLeaving(t *testing.T) { + s := &stickyBalanceStrategy{} + r := rand.New(rand.NewSource(time.Now().UnixNano())) + + // PLAN 1 + members := make(map[string]ConsumerGroupMemberMetadata, 20) + for i := 0; i < 200; i++ { + topics := make([]string, 200) + for j := 0; j < 200; j++ { + topics[j] = fmt.Sprintf("topic%d", j) + } + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{Topics: topics} + } + topics := make(map[string][]int32, 40) + for i := 0; i < 40; i++ { + partitionCount := r.Intn(20) + partitions := make([]int32, partitionCount) + for j := 0; j < partitionCount; j++ { + partitions[j] = int32(j) + } + topics[fmt.Sprintf("topic%d", i)] = partitions + } + + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + + for i := 0; i < 200; i++ { + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{ + Topics: members[fmt.Sprintf("consumer%d", i)].Topics, + UserData: encodeSubscriberPlan(t, plan1[fmt.Sprintf("consumer%d", i)]), + } + } + for i := 0; i < 50; i++ { + delete(members, fmt.Sprintf("consumer%d", i)) + } + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) +} + +func Test_stickyBalanceStrategy_Plan_NewSubscription(t *testing.T) { + s := &stickyBalanceStrategy{} + + members := make(map[string]ConsumerGroupMemberMetadata, 20) + for i := 0; i < 3; i++ { + topics := make([]string, 0) + for j := i; j <= 3*i-2; j++ { + topics = append(topics, fmt.Sprintf("topic%d", j)) + } + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{Topics: topics} + } + topics := make(map[string][]int32, 5) + for i := 1; i < 5; i++ { + topics[fmt.Sprintf("topic%d", i)] = []int32{0} + } + + plan1, err := s.Plan(members, topics) + if err != nil { + t.Errorf("stickyBalanceStrategy.Plan() error = %v", err) + return + } + verifyValidityAndBalance(t, members, plan1) + + members["consumer0"] = ConsumerGroupMemberMetadata{Topics: []string{"topic1"}} + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) +} + +func Test_stickyBalanceStrategy_Plan_ReassignmentWithRandomSubscriptionsAndChanges(t *testing.T) { + r := rand.New(rand.NewSource(time.Now().UnixNano())) + + minNumConsumers := 20 + maxNumConsumers := 40 + minNumTopics := 10 + maxNumTopics := 20 + + for round := 0; round < 100; round++ { + numTopics := minNumTopics + r.Intn(maxNumTopics-minNumTopics) + topics := make([]string, numTopics) + partitionsPerTopic := make(map[string][]int32, numTopics) + for i := 0; i < numTopics; i++ { + topicName := fmt.Sprintf("topic%d", i) + topics[i] = topicName + partitions := make([]int32, maxNumTopics) + for j := 0; j < maxNumTopics; j++ { + partitions[j] = int32(j) + } + partitionsPerTopic[topicName] = partitions + } + + numConsumers := minNumConsumers + r.Intn(maxNumConsumers-minNumConsumers) + members := make(map[string]ConsumerGroupMemberMetadata, numConsumers) + for i := 0; i < numConsumers; i++ { + sub := getRandomSublist(r, topics) + sort.Strings(sub) + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{Topics: sub} + } + + s := &stickyBalanceStrategy{} + plan, err := s.Plan(members, partitionsPerTopic) + verifyPlanIsBalancedAndSticky(t, s, members, plan, err) + + // PLAN 2 + membersPlan2 := make(map[string]ConsumerGroupMemberMetadata, numConsumers) + for i := 0; i < numConsumers; i++ { + sub := getRandomSublist(r, topics) + sort.Strings(sub) + membersPlan2[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{ + Topics: sub, + UserData: encodeSubscriberPlan(t, plan[fmt.Sprintf("consumer%d", i)]), + } + } + plan2, err := s.Plan(membersPlan2, partitionsPerTopic) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) + } +} + +func Test_stickyBalanceStrategy_Plan_MoveExistingAssignments(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := make(map[string][]int32, 6) + for i := 1; i <= 6; i++ { + topics[fmt.Sprintf("topic%d", i)] = []int32{0} + } + members := make(map[string]ConsumerGroupMemberMetadata, 3) + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + UserData: encodeSubscriberPlan(t, map[string][]int32{"topic1": []int32{0}}), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2", "topic3", "topic4"}, + UserData: encodeSubscriberPlan(t, map[string][]int32{"topic2": []int32{0}, "topic3": []int32{0}}), + } + members["consumer3"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic2", "topic3", "topic4", "topic5", "topic6"}, + UserData: encodeSubscriberPlan(t, map[string][]int32{"topic4": []int32{0}, "topic5": []int32{0}, "topic6": []int32{0}}), + } + + plan, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan, err) +} + +func Test_stickyBalanceStrategy_Plan_Stickiness(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := map[string][]int32{"topic1": []int32{0, 1, 2}} + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + "consumer2": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + "consumer3": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + "consumer4": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + } + + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + + // PLAN 2 + // remove the potential group leader + delete(members, "consumer1") + for i := 2; i <= 4; i++ { + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlan(t, plan1[fmt.Sprintf("consumer%d", i)]), + } + } + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) +} + +func Test_stickyBalanceStrategy_Plan_AssignmentUpdatedForDeletedTopic(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := make(map[string][]int32, 2) + topics["topic1"] = []int32{0} + topics["topic3"] = make([]int32, 100) + for i := 0; i < 100; i++ { + topics["topic3"][i] = int32(i) + } + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{Topics: []string{"topic1", "topic2", "topic3"}}, + } + + plan, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan, err) + verifyFullyBalanced(t, plan) + if (len(plan["consumer1"]["topic1"]) + len(plan["consumer1"]["topic3"])) != 101 { + t.Error("Incorrect number of partitions assigned") + return + } +} + +func Test_stickyBalanceStrategy_Plan_NoExceptionRaisedWhenOnlySubscribedTopicDeleted(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := map[string][]int32{"topic1": []int32{0, 1, 2}} + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + } + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + + // PLAN 2 + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: members["consumer1"].Topics, + UserData: encodeSubscriberPlan(t, plan1["consumer1"]), + } + + plan2, err := s.Plan(members, map[string][]int32{}) + if len(plan2) != 1 { + t.Error("Incorrect number of consumers") + return + } + if len(plan2["consumer1"]) != 0 { + t.Error("Incorrect number of consumer topic assignments") + return + } + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) +} + +func Test_stickyBalanceStrategy_Plan_AssignmentWithMultipleGenerations1(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := map[string][]int32{"topic1": []int32{0, 1, 2, 3, 4, 5}} + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + "consumer2": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + "consumer3": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + } + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + verifyFullyBalanced(t, plan1) + + // PLAN 2 + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, plan1["consumer1"], 1), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, plan1["consumer2"], 1), + } + delete(members, "consumer3") + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) + verifyFullyBalanced(t, plan2) + if len(intersection(plan1["consumer1"]["topic1"], plan2["consumer1"]["topic1"])) != 2 { + t.Error("stickyBalanceStrategy.Plan() consumer1 didn't maintain partitions across reassignment") + } + if len(intersection(plan1["consumer2"]["topic1"], plan2["consumer2"]["topic1"])) != 2 { + t.Error("stickyBalanceStrategy.Plan() consumer1 didn't maintain partitions across reassignment") + } + + // PLAN 3 + delete(members, "consumer1") + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, plan2["consumer2"], 2), + } + members["consumer3"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, plan1["consumer3"], 1), + } + + plan3, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan3, err) + verifyFullyBalanced(t, plan3) +} + +func Test_stickyBalanceStrategy_Plan_AssignmentWithMultipleGenerations2(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := map[string][]int32{"topic1": []int32{0, 1, 2, 3, 4, 5}} + members := map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + "consumer2": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + "consumer3": ConsumerGroupMemberMetadata{Topics: []string{"topic1"}}, + } + plan1, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan1, err) + verifyFullyBalanced(t, plan1) + + // PLAN 2 + delete(members, "consumer1") + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, plan1["consumer2"], 1), + } + delete(members, "consumer3") + + plan2, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan2, err) + verifyFullyBalanced(t, plan2) + if len(intersection(plan1["consumer2"]["topic1"], plan2["consumer2"]["topic1"])) != 2 { + t.Error("stickyBalanceStrategy.Plan() consumer1 didn't maintain partitions across reassignment") + } + + // PLAN 3 + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, plan1["consumer1"], 1), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, plan2["consumer2"], 2), + } + members["consumer3"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, plan1["consumer3"], 1), + } + plan3, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan3, err) + verifyFullyBalanced(t, plan3) +} +func Test_stickyBalanceStrategy_Plan_AssignmentWithConflictingPreviousGenerations(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := map[string][]int32{"topic1": []int32{0, 1, 2, 3, 4, 5}} + members := make(map[string]ConsumerGroupMemberMetadata, 3) + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{0, 1, 4}}, 1), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{0, 2, 3}}, 1), + } + members["consumer3"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{3, 4, 5}}, 2), + } + + plan, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan, err) + verifyFullyBalanced(t, plan) +} + +func Test_stickyBalanceStrategy_Plan_SchemaBackwardCompatibility(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := map[string][]int32{"topic1": []int32{0, 1, 2}} + members := make(map[string]ConsumerGroupMemberMetadata, 3) + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{0, 2}}, 1), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithOldSchema(t, map[string][]int32{"topic1": []int32{1}}), + } + members["consumer3"] = ConsumerGroupMemberMetadata{Topics: []string{"topic1"}} + + plan, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan, err) + verifyFullyBalanced(t, plan) +} + +func Test_stickyBalanceStrategy_Plan_ConflictingPreviousAssignments(t *testing.T) { + s := &stickyBalanceStrategy{} + + topics := map[string][]int32{"topic1": []int32{0, 1}} + members := make(map[string]ConsumerGroupMemberMetadata, 2) + members["consumer1"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{0, 1}}, 1), + } + members["consumer2"] = ConsumerGroupMemberMetadata{ + Topics: []string{"topic1"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{0, 1}}, 1), + } + + plan, err := s.Plan(members, topics) + verifyPlanIsBalancedAndSticky(t, s, members, plan, err) + verifyFullyBalanced(t, plan) +} + +func BenchmarkStickAssignmentWithLargeNumberOfConsumersAndTopics(b *testing.B) { + s := &stickyBalanceStrategy{} + r := rand.New(rand.NewSource(time.Now().UnixNano())) + members := make(map[string]ConsumerGroupMemberMetadata, 20) + for i := 0; i < 200; i++ { + topics := make([]string, 200) + for j := 0; j < 200; j++ { + topics[j] = fmt.Sprintf("topic%d", j) + } + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{Topics: topics} + } + topics := make(map[string][]int32, 40) + for i := 0; i < 40; i++ { + partitionCount := r.Intn(20) + partitions := make([]int32, partitionCount) + for j := 0; j < partitionCount; j++ { + partitions[j] = int32(j) + } + topics[fmt.Sprintf("topic%d", i)] = partitions + } + b.ResetTimer() + + for n := 0; n < b.N; n++ { + if _, err := s.Plan(members, topics); err != nil { + b.Errorf("Error building plan in benchmark: %v", err) + } + } +} + +func BenchmarkStickAssignmentWithLargeNumberOfConsumersAndTopicsAndExistingAssignments(b *testing.B) { + s := &stickyBalanceStrategy{} + r := rand.New(rand.NewSource(time.Now().UnixNano())) + members := make(map[string]ConsumerGroupMemberMetadata, 20) + for i := 0; i < 200; i++ { + topics := make([]string, 200) + for j := 0; j < 200; j++ { + topics[j] = fmt.Sprintf("topic%d", j) + } + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{Topics: topics} + } + topics := make(map[string][]int32, 40) + for i := 0; i < 40; i++ { + partitionCount := r.Intn(20) + partitions := make([]int32, partitionCount) + for j := 0; j < partitionCount; j++ { + partitions[j] = int32(j) + } + topics[fmt.Sprintf("topic%d", i)] = partitions + } + plan, _ := s.Plan(members, topics) + + for i := 0; i < 200; i++ { + members[fmt.Sprintf("consumer%d", i)] = ConsumerGroupMemberMetadata{ + Topics: members[fmt.Sprintf("consumer%d", i)].Topics, + UserData: encodeSubscriberPlanWithGenerationForBenchmark(b, plan[fmt.Sprintf("consumer%d", i)], 1), + } + } + for i := 0; i < 1; i++ { + delete(members, fmt.Sprintf("consumer%d", i)) + } + b.ResetTimer() + + for n := 0; n < b.N; n++ { + if _, err := s.Plan(members, topics); err != nil { + b.Errorf("Error building plan in benchmark: %v", err) + } + } +} + +func verifyPlanIsBalancedAndSticky(t *testing.T, s *stickyBalanceStrategy, members map[string]ConsumerGroupMemberMetadata, plan BalanceStrategyPlan, err error) { + if err != nil { + t.Errorf("stickyBalanceStrategy.Plan() error = %v", err) + return + } + if !s.movements.isSticky() { + t.Error("stickyBalanceStrategy.Plan() not sticky") + return + } + verifyValidityAndBalance(t, members, plan) +} + +func verifyValidityAndBalance(t *testing.T, consumers map[string]ConsumerGroupMemberMetadata, plan BalanceStrategyPlan) { + size := len(consumers) + if size != len(plan) { + t.Errorf("Subscription size (%d) not equal to plan size (%d)", size, len(plan)) + t.FailNow() + } + + members := make([]string, size) + i := 0 + for memberID := range consumers { + members[i] = memberID + i++ + } + sort.Strings(members) + + for i, memberID := range members { + for assignedTopic := range plan[memberID] { + found := false + for _, assignableTopic := range consumers[memberID].Topics { + if assignableTopic == assignableTopic { + found = true + break + } + } + if !found { + t.Errorf("Consumer %s had assigned topic %s that wasn't in the list of assignable topics", memberID, assignedTopic) + t.FailNow() + } + } + + // skip last consumer + if i == len(members)-1 { + continue + } + + consumerAssignments := make([]topicPartitionAssignment, 0) + for topic, partitions := range plan[memberID] { + for _, partition := range partitions { + consumerAssignments = append(consumerAssignments, topicPartitionAssignment{Topic: topic, Partition: partition}) + } + } + + for j := i + 1; j < size; j++ { + otherConsumer := members[j] + otherConsumerAssignments := make([]topicPartitionAssignment, 0) + for topic, partitions := range plan[otherConsumer] { + for _, partition := range partitions { + otherConsumerAssignments = append(otherConsumerAssignments, topicPartitionAssignment{Topic: topic, Partition: partition}) + } + } + assignmentsIntersection := intersection(consumerAssignments, otherConsumerAssignments) + if len(assignmentsIntersection) > 0 { + t.Errorf("Consumers %s and %s have common partitions assigned to them: %v", memberID, otherConsumer, assignmentsIntersection) + t.FailNow() + } + + if math.Abs(float64(len(consumerAssignments)-len(otherConsumerAssignments))) <= 1 { + continue + } + + if len(consumerAssignments) > len(otherConsumerAssignments) { + for _, topic := range consumerAssignments { + if _, exists := plan[otherConsumer][topic.Topic]; exists { + t.Errorf("Some partitions can be moved from %s to %s to achieve a better balance, %s has %d assignments, and %s has %d assignments", otherConsumer, memberID, memberID, len(consumerAssignments), otherConsumer, len(otherConsumerAssignments)) + t.FailNow() + } + } + } + + if len(otherConsumerAssignments) > len(consumerAssignments) { + for _, topic := range otherConsumerAssignments { + if _, exists := plan[memberID][topic.Topic]; exists { + t.Errorf("Some partitions can be moved from %s to %s to achieve a better balance, %s has %d assignments, and %s has %d assignments", memberID, otherConsumer, otherConsumer, len(otherConsumerAssignments), memberID, len(consumerAssignments)) + t.FailNow() + } + } + } + } + } +} + +// Produces the intersection of two slices +// From https://github.com/juliangruber/go-intersect +func intersection(a interface{}, b interface{}) []interface{} { + set := make([]interface{}, 0) + hash := make(map[interface{}]bool) + av := reflect.ValueOf(a) + bv := reflect.ValueOf(b) + + for i := 0; i < av.Len(); i++ { + el := av.Index(i).Interface() + hash[el] = true + } + + for i := 0; i < bv.Len(); i++ { + el := bv.Index(i).Interface() + if _, found := hash[el]; found { + set = append(set, el) + } + } + + return set +} + +func encodeSubscriberPlan(t *testing.T, assignments map[string][]int32) []byte { + return encodeSubscriberPlanWithGeneration(t, assignments, defaultGeneration) +} + +func encodeSubscriberPlanWithGeneration(t *testing.T, assignments map[string][]int32, generation int32) []byte { + userDataBytes, err := encode(&StickyAssignorUserDataV1{ + Topics: assignments, + Generation: generation, + }, nil) + if err != nil { + t.Errorf("encodeSubscriberPlan error = %v", err) + t.FailNow() + } + return userDataBytes +} + +func encodeSubscriberPlanWithGenerationForBenchmark(b *testing.B, assignments map[string][]int32, generation int32) []byte { + userDataBytes, err := encode(&StickyAssignorUserDataV1{ + Topics: assignments, + Generation: generation, + }, nil) + if err != nil { + b.Errorf("encodeSubscriberPlan error = %v", err) + b.FailNow() + } + return userDataBytes +} + +func encodeSubscriberPlanWithOldSchema(t *testing.T, assignments map[string][]int32) []byte { + userDataBytes, err := encode(&StickyAssignorUserDataV0{ + Topics: assignments, + }, nil) + if err != nil { + t.Errorf("encodeSubscriberPlan error = %v", err) + t.FailNow() + } + return userDataBytes +} + +// verify that the plan is fully balanced, assumes that all consumers can +// consume from the same set of topics +func verifyFullyBalanced(t *testing.T, plan BalanceStrategyPlan) { + min := math.MaxInt32 + max := math.MinInt32 + for _, topics := range plan { + assignedPartitionsCount := 0 + for _, partitions := range topics { + assignedPartitionsCount += len(partitions) + } + if assignedPartitionsCount < min { + min = assignedPartitionsCount + } + if assignedPartitionsCount > max { + max = assignedPartitionsCount + } + } + if (max - min) > 1 { + t.Errorf("Plan partition assignment is not fully balanced: min=%d, max=%d", min, max) + } +} + +func getRandomSublist(r *rand.Rand, s []string) []string { + howManyToRemove := r.Intn(len(s)) + allEntriesMap := make(map[int]string) + for i, s := range s { + allEntriesMap[i] = s + } + for i := 0; i < howManyToRemove; i++ { + delete(allEntriesMap, r.Intn(len(allEntriesMap))) + } + + subList := make([]string, len(allEntriesMap)) + i := 0 + for _, s := range allEntriesMap { + subList[i] = s + i++ + } + return subList +} diff --git a/consumer_group.go b/consumer_group.go index 8de95137e..df0d737ad 100644 --- a/consumer_group.go +++ b/consumer_group.go @@ -63,6 +63,8 @@ type consumerGroup struct { lock sync.Mutex closed chan none closeOnce sync.Once + + userData []byte } // NewConsumerGroup creates a new consumer group the given broker addresses and configuration. @@ -282,6 +284,7 @@ func (c *consumerGroup) newSession(ctx context.Context, topics []string, handler return nil, err } claims = members.Topics + c.userData = members.UserData for _, partitions := range claims { sort.Sort(int32Slice(partitions)) @@ -303,9 +306,14 @@ func (c *consumerGroup) joinGroupRequest(coordinator *Broker, topics []string) ( req.RebalanceTimeout = int32(c.config.Consumer.Group.Rebalance.Timeout / time.Millisecond) } + // use static user-data if configured, otherwise use consumer-group userdata from the last sync + userData := c.config.Consumer.Group.Member.UserData + if len(userData) == 0 { + userData = c.userData + } meta := &ConsumerGroupMemberMetadata{ Topics: topics, - UserData: c.config.Consumer.Group.Member.UserData, + UserData: userData, } strategy := c.config.Consumer.Group.Rebalance.Strategy if err := req.AddGroupProtocolMetadata(strategy.Name(), meta); err != nil { @@ -322,10 +330,20 @@ func (c *consumerGroup) syncGroupRequest(coordinator *Broker, plan BalanceStrate GenerationId: generationID, } for memberID, topics := range plan { - err := req.AddGroupAssignmentMember(memberID, &ConsumerGroupMemberAssignment{ - Topics: topics, - }) - if err != nil { + assignment := &ConsumerGroupMemberAssignment{Topics: topics} + + // Include topic assignments in group-assignment userdata for each consumer-group member + if c.config.Consumer.Group.Rebalance.Strategy == BalanceStrategySticky { + userDataBytes, err := encode(&StickyAssignorUserDataV1{ + Topics: topics, + Generation: generationID, + }, nil) + if err != nil { + return nil, err + } + assignment.UserData = userDataBytes + } + if err := req.AddGroupAssignmentMember(memberID, assignment); err != nil { return nil, err } } diff --git a/examples/consumergroup/main.go b/examples/consumergroup/main.go index 54db82b42..06777e564 100644 --- a/examples/consumergroup/main.go +++ b/examples/consumergroup/main.go @@ -15,12 +15,13 @@ import ( // Sarma configuration options var ( - brokers = "" - version = "" - group = "" - topics = "" - oldest = true - verbose = false + brokers = "" + version = "" + group = "" + topics = "" + assignor = "" + oldest = true + verbose = false ) func init() { @@ -28,6 +29,7 @@ func init() { flag.StringVar(&group, "group", "", "Kafka consumer group definition") flag.StringVar(&version, "version", "2.1.1", "Kafka cluster version") flag.StringVar(&topics, "topics", "", "Kafka topics to be consumed, as a comma seperated list") + flag.StringVar(&assignor, "assignor", "range", "Consumer group partition assignment strategy (range, roundrobin, sticky)") flag.BoolVar(&oldest, "oldest", true, "Kafka consumer consume initial ofset from oldest") flag.BoolVar(&verbose, "verbose", false, "Sarama logging") flag.Parse() @@ -64,6 +66,17 @@ func main() { config := sarama.NewConfig() config.Version = version + switch assignor { + case "sticky": + config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategySticky + case "roundrobin": + config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRoundRobin + case "range": + config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRange + default: + log.Panicf("Unrecognized consumer group partition assignor: %s", assignor) + } + if oldest { config.Consumer.Offsets.Initial = sarama.OffsetOldest } diff --git a/sticky_assignor_user_data.go b/sticky_assignor_user_data.go new file mode 100644 index 000000000..bb0c82c34 --- /dev/null +++ b/sticky_assignor_user_data.go @@ -0,0 +1,124 @@ +package sarama + +type topicPartitionAssignment struct { + Topic string + Partition int32 +} + +type StickyAssignorUserData interface { + partitions() []topicPartitionAssignment + hasGeneration() bool + generation() int +} + +//StickyAssignorUserDataV0 holds topic partition information for an assignment +type StickyAssignorUserDataV0 struct { + Topics map[string][]int32 + + topicPartitions []topicPartitionAssignment +} + +func (m *StickyAssignorUserDataV0) encode(pe packetEncoder) error { + if err := pe.putArrayLength(len(m.Topics)); err != nil { + return err + } + + for topic, partitions := range m.Topics { + if err := pe.putString(topic); err != nil { + return err + } + if err := pe.putInt32Array(partitions); err != nil { + return err + } + } + return nil +} + +func (m *StickyAssignorUserDataV0) decode(pd packetDecoder) (err error) { + var topicLen int + if topicLen, err = pd.getArrayLength(); err != nil { + return + } + + m.Topics = make(map[string][]int32, topicLen) + for i := 0; i < topicLen; i++ { + var topic string + if topic, err = pd.getString(); err != nil { + return + } + if m.Topics[topic], err = pd.getInt32Array(); err != nil { + return + } + } + m.topicPartitions = populateTopicPartitions(m.Topics) + return nil +} + +func (m *StickyAssignorUserDataV0) partitions() []topicPartitionAssignment { return m.topicPartitions } +func (m *StickyAssignorUserDataV0) hasGeneration() bool { return false } +func (m *StickyAssignorUserDataV0) generation() int { return defaultGeneration } + +//StickyAssignorUserDataV1 holds topic partition information for an assignment +type StickyAssignorUserDataV1 struct { + Topics map[string][]int32 + Generation int32 + + topicPartitions []topicPartitionAssignment +} + +func (m *StickyAssignorUserDataV1) encode(pe packetEncoder) error { + if err := pe.putArrayLength(len(m.Topics)); err != nil { + return err + } + + for topic, partitions := range m.Topics { + if err := pe.putString(topic); err != nil { + return err + } + if err := pe.putInt32Array(partitions); err != nil { + return err + } + } + + pe.putInt32(m.Generation) + return nil +} + +func (m *StickyAssignorUserDataV1) decode(pd packetDecoder) (err error) { + var topicLen int + if topicLen, err = pd.getArrayLength(); err != nil { + return + } + + m.Topics = make(map[string][]int32, topicLen) + for i := 0; i < topicLen; i++ { + var topic string + if topic, err = pd.getString(); err != nil { + return + } + if m.Topics[topic], err = pd.getInt32Array(); err != nil { + return + } + } + + m.Generation, err = pd.getInt32() + if err != nil { + return err + } + m.topicPartitions = populateTopicPartitions(m.Topics) + return nil +} + +func (m *StickyAssignorUserDataV1) partitions() []topicPartitionAssignment { return m.topicPartitions } +func (m *StickyAssignorUserDataV1) hasGeneration() bool { return true } +func (m *StickyAssignorUserDataV1) generation() int { return int(m.Generation) } + +func populateTopicPartitions(topics map[string][]int32) []topicPartitionAssignment { + topicPartitions := make([]topicPartitionAssignment, 0) + for topic, partitions := range topics { + for _, partition := range partitions { + topicPartitions = append(topicPartitions, topicPartitionAssignment{Topic: topic, Partition: partition}) + } + } + return topicPartitions +} diff --git a/sticky_assignor_user_data_test.go b/sticky_assignor_user_data_test.go new file mode 100644 index 000000000..9eb09a615 --- /dev/null +++ b/sticky_assignor_user_data_test.go @@ -0,0 +1,51 @@ +package sarama + +import ( + "encoding/base64" + "testing" +) + +func TestStickyAssignorUserDataV0(t *testing.T) { + // Single topic with deterministic ordering across encode-decode + req := &StickyAssignorUserDataV0{} + data := decodeUserDataBytes(t, "AAAAAQADdDAzAAAAAQAAAAU=") + testDecodable(t, "", req, data) + testEncodable(t, "", req, data) + + // Multiple partitions + req = &StickyAssignorUserDataV0{} + data = decodeUserDataBytes(t, "AAAAAQADdDE4AAAAEgAAAAAAAAABAAAAAgAAAAMAAAAEAAAABQAAAAYAAAAHAAAACAAAAAkAAAAKAAAACwAAAAwAAAANAAAADgAAAA8AAAAQAAAAEQ==") + testDecodable(t, "", req, data) + + // Multiple topics and partitions + req = &StickyAssignorUserDataV0{} + data = decodeUserDataBytes(t, "AAAABQADdDEyAAAAAgAAAAIAAAAKAAN0MTEAAAABAAAABAADdDE0AAAAAQAAAAgAA3QxMwAAAAEAAAANAAN0MDkAAAABAAAABQ==") + testDecodable(t, "", req, data) +} + +func TestStickyAssignorUserDataV1(t *testing.T) { + // Single topic with deterministic ordering across encode-decode + req := &StickyAssignorUserDataV1{} + data := decodeUserDataBytes(t, "AAAAAQADdDA2AAAAAgAAAAAAAAAE/////w==") + testDecodable(t, "", req, data) + testEncodable(t, "", req, data) + + // Multiple topics and partitions + req = &StickyAssignorUserDataV1{} + data = decodeUserDataBytes(t, "AAAABgADdDEwAAAAAgAAAAIAAAAJAAN0MTIAAAACAAAAAwAAAAsAA3QxNAAAAAEAAAAEAAN0MTMAAAABAAAACwADdDE1AAAAAQAAAAwAA3QwOQAAAAEAAAAG/////w==") + testDecodable(t, "", req, data) + + // Generation is populated + req = &StickyAssignorUserDataV1{} + data = decodeUserDataBytes(t, "AAAAAQAHdG9waWMwMQAAAAMAAAAAAAAAAQAAAAIAAAAB") + testDecodable(t, "", req, data) +} + +func decodeUserDataBytes(t *testing.T, base64Data string) []byte { + data, err := base64.StdEncoding.DecodeString(base64Data) + if err != nil { + t.Errorf("Error decoding data: %v", err) + t.FailNow() + } + return data +} From 78d3678f40b6b6144aeddc0bc2605fc9637f81fc Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Thu, 11 Jul 2019 07:41:10 -0700 Subject: [PATCH 2/9] Make the partition strategy assignment names constants, and compare the active strategy name against the sticky balance strategy name when handling a sync group request to allow for easy extension of the sticky balance strategy --- balance_strategy.go | 15 ++++++++++++--- consumer_group.go | 2 +- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/balance_strategy.go b/balance_strategy.go index 7392cc738..75f41cf6c 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -6,6 +6,15 @@ import ( ) const ( + // RangeBalanceStrategyName identifies strategies that use the range partition assignment strategy + RangeBalanceStrategyName = "range" + + // RoundRobinBalanceStrategyName identifies strategies that use the round-robin partition assignment strategy + RoundRobinBalanceStrategyName = "roundrobin" + + // StickyBalanceStrategyName identifies strategies that use the sticky-partition assignment strategy + StickyBalanceStrategyName = "sticky" + defaultGeneration = -1 ) @@ -45,7 +54,7 @@ type BalanceStrategy interface { // M1: {T: [0, 1, 2]} // M2: {T: [3, 4, 5]} var BalanceStrategyRange = &balanceStrategy{ - name: "range", + name: RangeBalanceStrategyName, coreFn: func(plan BalanceStrategyPlan, memberIDs []string, topic string, partitions []int32) { step := float64(len(partitions)) / float64(len(memberIDs)) @@ -63,7 +72,7 @@ var BalanceStrategyRange = &balanceStrategy{ // M1: {T: [0, 2, 4]} // M2: {T: [1, 3, 5]} var BalanceStrategyRoundRobin = &balanceStrategy{ - name: "roundrobin", + name: RoundRobinBalanceStrategyName, coreFn: func(plan BalanceStrategyPlan, memberIDs []string, topic string, partitions []int32) { for i, part := range partitions { memberID := memberIDs[i%len(memberIDs)] @@ -150,7 +159,7 @@ type stickyBalanceStrategy struct { } // Name implements BalanceStrategy. -func (s *stickyBalanceStrategy) Name() string { return "sticky" } +func (s *stickyBalanceStrategy) Name() string { return StickyBalanceStrategyName } // Plan implements BalanceStrategy. func (s *stickyBalanceStrategy) Plan(members map[string]ConsumerGroupMemberMetadata, topics map[string][]int32) (BalanceStrategyPlan, error) { diff --git a/consumer_group.go b/consumer_group.go index df0d737ad..db7172a93 100644 --- a/consumer_group.go +++ b/consumer_group.go @@ -333,7 +333,7 @@ func (c *consumerGroup) syncGroupRequest(coordinator *Broker, plan BalanceStrate assignment := &ConsumerGroupMemberAssignment{Topics: topics} // Include topic assignments in group-assignment userdata for each consumer-group member - if c.config.Consumer.Group.Rebalance.Strategy == BalanceStrategySticky { + if c.config.Consumer.Group.Rebalance.Strategy.Name() == StickyBalanceStrategyName { userDataBytes, err := encode(&StickyAssignorUserDataV1{ Topics: topics, Generation: generationID, From 7f26806906465ec6a2ba2109244bc6e19609a83a Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Mon, 12 Aug 2019 12:32:50 -0700 Subject: [PATCH 3/9] Use priority queue to sort consumer group members by number of partition assignments --- balance_strategy.go | 117 +++++++++++++++++++++++++++++--------------- 1 file changed, 78 insertions(+), 39 deletions(-) diff --git a/balance_strategy.go b/balance_strategy.go index 75f41cf6c..dac689f49 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -1,6 +1,7 @@ package sarama import ( + "container/heap" "math" "sort" ) @@ -507,22 +508,14 @@ func canTopicPartitionParticipateInReassignment(partition topicPartitionAssignme // The assignment should improve the overall balance of the partition assignments to consumers. func assignPartition(partition topicPartitionAssignment, sortedCurrentSubscriptions []string, currentAssignment map[string][]topicPartitionAssignment, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment, currentPartitionConsumer map[topicPartitionAssignment]string) []string { - updatedSubscriptions := make([]string, len(sortedCurrentSubscriptions)) - for i, s := range sortedCurrentSubscriptions { - updatedSubscriptions[i] = s - } - i := 0 for _, memberID := range sortedCurrentSubscriptions { if memberAssignmentsIncludeTopicPartition(consumer2AllPotentialPartitions[memberID], partition) { - updatedSubscriptions = removeIndexFromStringSlice(updatedSubscriptions, i) currentAssignment[memberID] = append(currentAssignment[memberID], partition) currentPartitionConsumer[partition] = memberID - updatedSubscriptions = append(updatedSubscriptions, memberID) break } - i++ } - return updatedSubscriptions + return sortMemberIDsByPartitionAssignments(currentAssignment) } // Deserialize topic partition assignment data to aid with creation of a sticky assignment. @@ -599,44 +592,52 @@ func sortPartitions(currentAssignment map[string][]topicPartitionAssignment, par // most assigned partitions to those with least) assignments := filterAssignedPartitions(currentAssignment, partition2AllPotentialConsumers) - // sortedMemberIDs contains a descending-sorted list of consumers based on how many valid partitions are currently assigned to them - sortedMemberIDs := sortMemberIDsByPartitionAssignments(assignments) - for i := len(sortedMemberIDs)/2 - 1; i >= 0; i-- { - opp := len(sortedMemberIDs) - 1 - i - sortedMemberIDs[i], sortedMemberIDs[opp] = sortedMemberIDs[opp], sortedMemberIDs[i] + // use priority-queue to evaluate consumer group members in descending-order based on + // the number of topic partition assignments (i.e. consumers with most assignments first) + pq := make(assignmentPriorityQueue, len(assignments)) + i := 0 + for consumerID, consumerAssignments := range assignments { + pq[i] = &consumerGroupMember{ + id: consumerID, + assignments: consumerAssignments, + index: i, + } + i++ } + heap.Init(&pq) + for { // loop until no consumer-group members remain - if len(sortedMemberIDs) == 0 { + if pq.Len() == 0 { break } - updatedMemberIDs := make([]string, 0) - for _, memberID := range sortedMemberIDs { - // partitions that were assigned to a different consumer last time - prevPartitions := make([]topicPartitionAssignment, 0) - for partition := range partitionsWithADifferentPreviousAssignment { - // from partitions that had a different consumer before, keep only those that are assigned to this consumer now - if memberAssignmentsIncludeTopicPartition(assignments[memberID], partition) { - prevPartitions = append(prevPartitions, partition) - } + member := pq[0] + + // partitions that were assigned to a different consumer last time + prevPartitions := make([]topicPartitionAssignment, 0) + for partition := range partitionsWithADifferentPreviousAssignment { + // from partitions that had a different consumer before, keep only those that are assigned to this consumer now + if memberAssignmentsIncludeTopicPartition(member.assignments, partition) { + prevPartitions = append(prevPartitions, partition) } + } - if len(prevPartitions) > 0 { - // if there is a partition on this consumer that was assigned to another consumer before mark it as good options for reassignment - partition := prevPartitions[0] - prevPartitions = append(prevPartitions[:0], prevPartitions[1:]...) - assignments[memberID] = removeTopicPartitionFromMemberAssignments(assignments[memberID], partition) - sortedPartitions = append(sortedPartitions, partition) - updatedMemberIDs = append(updatedMemberIDs, memberID) - } else if len(assignments[memberID]) > 0 { - // otherwise, mark any other one of the current partitions as a reassignment candidate - partition := assignments[memberID][0] - assignments[memberID] = append(assignments[memberID][:0], assignments[memberID][1:]...) - sortedPartitions = append(sortedPartitions, partition) - updatedMemberIDs = append(updatedMemberIDs, memberID) - } + if len(prevPartitions) > 0 { + // if there is a partition on this consumer that was assigned to another consumer before mark it as good options for reassignment + partition := prevPartitions[0] + prevPartitions = append(prevPartitions[:0], prevPartitions[1:]...) + member.assignments = removeTopicPartitionFromMemberAssignments(member.assignments, partition) + sortedPartitions = append(sortedPartitions, partition) + heap.Fix(&pq, member.index) + } else if len(member.assignments) > 0 { + // otherwise, mark any other one of the current partitions as a reassignment candidate + partition := member.assignments[0] + member.assignments = append(member.assignments[:0], member.assignments[1:]...) + sortedPartitions = append(sortedPartitions, partition) + heap.Fix(&pq, member.index) + } else { + heap.Remove(&pq, 0) } - sortedMemberIDs = updatedMemberIDs } for partition := range partition2AllPotentialConsumers { @@ -1034,3 +1035,41 @@ nextCand: } return -1 } + +type consumerGroupMember struct { + index int // the index of the item in the heap + id string + assignments []topicPartitionAssignment +} + +// A assignmentPriorityQueue implements heap.Interface and holds Items. +type assignmentPriorityQueue []*consumerGroupMember + +func (pq assignmentPriorityQueue) Len() int { return len(pq) } + +func (pq assignmentPriorityQueue) Less(i, j int) bool { + // We want Pop to give us the highest, not lowest, priority so we use greater than here. + return len(pq[i].assignments) > len(pq[j].assignments) +} + +func (pq assignmentPriorityQueue) Swap(i, j int) { + pq[i], pq[j] = pq[j], pq[i] + pq[i].index = i + pq[j].index = j +} + +func (pq *assignmentPriorityQueue) Push(x interface{}) { + n := len(*pq) + member := x.(*consumerGroupMember) + member.index = n + *pq = append(*pq, member) +} + +func (pq *assignmentPriorityQueue) Pop() interface{} { + old := *pq + n := len(old) + member := old[n-1] + member.index = -1 // for safety + *pq = old[0 : n-1] + return member +} From df5d0c0ba128eabc57273f36c9ff8f46a52b5609 Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Mon, 12 Aug 2019 14:31:30 -0700 Subject: [PATCH 4/9] Sort the assignment priority queue in descending order using member-id if the assignment counts are identical --- balance_strategy.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/balance_strategy.go b/balance_strategy.go index dac689f49..1a1ad91b4 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -4,6 +4,7 @@ import ( "container/heap" "math" "sort" + "strings" ) const ( @@ -1048,7 +1049,10 @@ type assignmentPriorityQueue []*consumerGroupMember func (pq assignmentPriorityQueue) Len() int { return len(pq) } func (pq assignmentPriorityQueue) Less(i, j int) bool { - // We want Pop to give us the highest, not lowest, priority so we use greater than here. + // order asssignment priority queue in descending order using assignment-count/member-id + if len(pq[i].assignments) == len(pq[j].assignments) { + return strings.Compare(pq[i].id, pq[j].id) > 0 + } return len(pq[i].assignments) > len(pq[j].assignments) } From ae8e47f3da108cd10be062561870bdba545fd437 Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Mon, 12 Aug 2019 15:03:33 -0700 Subject: [PATCH 5/9] Fix comments --- balance_strategy.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/balance_strategy.go b/balance_strategy.go index 1a1ad91b4..bd86d5293 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -1038,12 +1038,13 @@ nextCand: } type consumerGroupMember struct { - index int // the index of the item in the heap + index int // the index of the consumer group member on the heap id string assignments []topicPartitionAssignment } -// A assignmentPriorityQueue implements heap.Interface and holds Items. +// assignmentPriorityQueue is a priority-queue of consumer group members that is sorted +// in descending order (most assignments to least assignments). type assignmentPriorityQueue []*consumerGroupMember func (pq assignmentPriorityQueue) Len() int { return len(pq) } From e52742d7deaf175c0a8499e4a947183d6ae3a55c Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Tue, 13 Aug 2019 07:16:51 -0700 Subject: [PATCH 6/9] PR feedback --- balance_strategy.go | 116 ++++++++++--------------- balance_strategy_test.go | 182 ++++++++++++++++++--------------------- 2 files changed, 133 insertions(+), 165 deletions(-) diff --git a/balance_strategy.go b/balance_strategy.go index bd86d5293..26ac514bb 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -444,7 +444,6 @@ func (s *stickyBalanceStrategy) performReassignments(reassignablePartitions []to modified = true break } - } } if !modified { @@ -550,23 +549,6 @@ func filterAssignedPartitions(currentAssignment map[string][]topicPartitionAssig return assignments } -func removeValueFromStringSlice(s []string, e string) []string { - for i, v := range s { - if v == e { - s = append(s[:i], s[i+1:]...) - break - } - } - return s -} - -func removeIndexFromStringSlice(s []string, i int) []string { - if len(s) == 0 { - return s - } - return append(s[:i], s[i+1:]...) -} - func removeTopicPartitionFromMemberAssignments(assignments []topicPartitionAssignment, topic topicPartitionAssignment) []topicPartitionAssignment { for i, assignment := range assignments { if assignment == topic { @@ -586,6 +568,11 @@ func memberAssignmentsIncludeTopicPartition(assignments []topicPartitionAssignme } func sortPartitions(currentAssignment map[string][]topicPartitionAssignment, partitionsWithADifferentPreviousAssignment map[topicPartitionAssignment]consumerGenerationPair, isFreshAssignment bool, partition2AllPotentialConsumers map[topicPartitionAssignment][]string, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment) []topicPartitionAssignment { + unassignedPartitions := make(map[topicPartitionAssignment]bool, len(partition2AllPotentialConsumers)) + for partition := range partition2AllPotentialConsumers { + unassignedPartitions[partition] = true + } + sortedPartitions := make([]topicPartitionAssignment, 0) if !isFreshAssignment && areSubscriptionsIdentical(partition2AllPotentialConsumers, consumer2AllPotentialPartitions) { // if this is a reassignment and the subscriptions are identical (all consumers can consumer from all topics) @@ -601,7 +588,6 @@ func sortPartitions(currentAssignment map[string][]topicPartitionAssignment, par pq[i] = &consumerGroupMember{ id: consumerID, assignments: consumerAssignments, - index: i, } i++ } @@ -615,43 +601,36 @@ func sortPartitions(currentAssignment map[string][]topicPartitionAssignment, par member := pq[0] // partitions that were assigned to a different consumer last time - prevPartitions := make([]topicPartitionAssignment, 0) - for partition := range partitionsWithADifferentPreviousAssignment { - // from partitions that had a different consumer before, keep only those that are assigned to this consumer now - if memberAssignmentsIncludeTopicPartition(member.assignments, partition) { - prevPartitions = append(prevPartitions, partition) + var prevPartition *topicPartitionAssignment + var prevPartitionIndex int + for i, partition := range member.assignments { + if _, exists := partitionsWithADifferentPreviousAssignment[partition]; exists { + prevPartition = &partition + prevPartitionIndex = i + break } } - if len(prevPartitions) > 0 { + if prevPartition != nil { // if there is a partition on this consumer that was assigned to another consumer before mark it as good options for reassignment - partition := prevPartitions[0] - prevPartitions = append(prevPartitions[:0], prevPartitions[1:]...) - member.assignments = removeTopicPartitionFromMemberAssignments(member.assignments, partition) - sortedPartitions = append(sortedPartitions, partition) - heap.Fix(&pq, member.index) + member.assignments = append(member.assignments[:prevPartitionIndex], member.assignments[prevPartitionIndex+1:]...) + sortedPartitions = append(sortedPartitions, *prevPartition) + delete(unassignedPartitions, *prevPartition) + heap.Fix(&pq, 0) } else if len(member.assignments) > 0 { // otherwise, mark any other one of the current partitions as a reassignment candidate partition := member.assignments[0] member.assignments = append(member.assignments[:0], member.assignments[1:]...) sortedPartitions = append(sortedPartitions, partition) - heap.Fix(&pq, member.index) + delete(unassignedPartitions, partition) + heap.Fix(&pq, 0) } else { - heap.Remove(&pq, 0) + heap.Pop(&pq) } } - for partition := range partition2AllPotentialConsumers { - found := false - for _, p := range sortedPartitions { - if partition == p { - found = true - break - } - } - if !found { - sortedPartitions = append(sortedPartitions, partition) - } + for partition := range unassignedPartitions { + sortedPartitions = append(sortedPartitions, partition) } } else { // an ascending sorted set of topic partitions based on how many consumers can potentially use them @@ -678,12 +657,21 @@ func sortMemberIDsByPartitionAssignments(assignments map[string][]topicPartition func sortPartitionsByPotentialConsumerAssignments(partition2AllPotentialConsumers map[topicPartitionAssignment][]string) []topicPartitionAssignment { // sort the members by the number of partition assignments in descending order - sortedPartionIDs := make([]topicPartitionAssignment, 0, len(partition2AllPotentialConsumers)) + sortedPartionIDs := make([]topicPartitionAssignment, len(partition2AllPotentialConsumers)) + i := 0 for partition := range partition2AllPotentialConsumers { - sortedPartionIDs = append(sortedPartionIDs, partition) - } - sort.SliceStable(sortedPartionIDs, func(i, j int) bool { - return len(partition2AllPotentialConsumers[sortedPartionIDs[i]]) > len(partition2AllPotentialConsumers[sortedPartionIDs[j]]) + sortedPartionIDs[i] = partition + i++ + } + sort.Slice(sortedPartionIDs, func(i, j int) bool { + if len(partition2AllPotentialConsumers[sortedPartionIDs[i]]) == len(partition2AllPotentialConsumers[sortedPartionIDs[j]]) { + ret := strings.Compare(sortedPartionIDs[i].Topic, sortedPartionIDs[j].Topic) + if ret == 0 { + return sortedPartionIDs[i].Partition < sortedPartionIDs[j].Partition + } + return ret < 0 + } + return len(partition2AllPotentialConsumers[sortedPartionIDs[i]]) < len(partition2AllPotentialConsumers[sortedPartionIDs[j]]) }) return sortedPartionIDs } @@ -805,22 +793,18 @@ func prepopulateCurrentAssignments(members map[string]ConsumerGroupMemberMetadat } sort.Sort(sort.Reverse(sort.IntSlice(generations))) - if len(generations) > 0 { - consumer := consumers[generations[0]] - currentConsumerAssignments, exists := currentAssignment[consumer] - if !exists { - currentConsumerAssignments = []topicPartitionAssignment{partition} - } else { - currentConsumerAssignments = append(currentConsumerAssignments, partition) - } - currentAssignment[consumer] = currentConsumerAssignments + consumer := consumers[generations[0]] + if _, exists := currentAssignment[consumer]; !exists { + currentAssignment[consumer] = []topicPartitionAssignment{partition} + } else { + currentAssignment[consumer] = append(currentAssignment[consumer], partition) + } - // check for previous assignment, if any - if len(generations) > 1 { - prevAssignment[partition] = consumerGenerationPair{ - MemberID: consumers[generations[1]], - Generation: generations[1], - } + // check for previous assignment, if any + if len(generations) > 1 { + prevAssignment[partition] = consumerGenerationPair{ + MemberID: consumers[generations[1]], + Generation: generations[1], } } } @@ -1038,7 +1022,6 @@ nextCand: } type consumerGroupMember struct { - index int // the index of the consumer group member on the heap id string assignments []topicPartitionAssignment } @@ -1059,14 +1042,10 @@ func (pq assignmentPriorityQueue) Less(i, j int) bool { func (pq assignmentPriorityQueue) Swap(i, j int) { pq[i], pq[j] = pq[j], pq[i] - pq[i].index = i - pq[j].index = j } func (pq *assignmentPriorityQueue) Push(x interface{}) { - n := len(*pq) member := x.(*consumerGroupMember) - member.index = n *pq = append(*pq, member) } @@ -1074,7 +1053,6 @@ func (pq *assignmentPriorityQueue) Pop() interface{} { old := *pq n := len(old) member := old[n-1] - member.index = -1 // for safety *pq = old[0 : n-1] return member } diff --git a/balance_strategy_test.go b/balance_strategy_test.go index ef23d0ebf..72b0492dd 100644 --- a/balance_strategy_test.go +++ b/balance_strategy_test.go @@ -880,102 +880,6 @@ func Test_removeTopicPartitionFromMemberAssignments(t *testing.T) { } } -func Test_removeIndexFromStringSlice(t *testing.T) { - type args struct { - s []string - i int - } - tests := []struct { - name string - args args - want []string - }{ - { - name: "Empty slice", - args: args{ - s: make([]string, 0), - i: 0, - }, - want: make([]string, 0), - }, - { - name: "Slice with single entry", - args: args{ - s: []string{"foo"}, - i: 0, - }, - want: make([]string, 0), - }, - { - name: "Slice with multiple entries", - args: args{ - s: []string{"a", "b", "c"}, - i: 0, - }, - want: []string{"b", "c"}, - }, - { - name: "Slice with multiple entries and index is in the middle", - args: args{ - s: []string{"a", "b", "c"}, - i: 1, - }, - want: []string{"a", "c"}, - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := removeIndexFromStringSlice(tt.args.s, tt.args.i); !reflect.DeepEqual(got, tt.want) { - t.Errorf("removeIndexFromSlice() = %v, want %v", got, tt.want) - } - }) - } -} - -func Test_removeValueFromStringSlice(t *testing.T) { - type args struct { - s []string - e string - } - tests := []struct { - name string - args args - want []string - }{ - { - name: "Empty input slice", - args: args{ - s: []string{}, - e: "", - }, - want: []string{}, - }, - { - name: "Input slice with one entry that doesn't match", - args: args{ - s: []string{"a"}, - e: "b", - }, - want: []string{"a"}, - }, - { - name: "Input slice with multiple entries and a positive match", - args: args{ - s: []string{"a", "b", "c"}, - e: "b", - }, - want: []string{"a", "c"}, - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := removeValueFromStringSlice(tt.args.s, tt.args.e); !reflect.DeepEqual(got, tt.want) { - t.Errorf("removeValueFromSlice() = %v, want %v", got, tt.want) - } - }) - } -} - func Test_assignPartition(t *testing.T) { type args struct { partition topicPartitionAssignment @@ -2276,3 +2180,89 @@ func getRandomSublist(r *rand.Rand, s []string) []string { } return subList } + +func Test_sortPartitionsByPotentialConsumerAssignments(t *testing.T) { + type args struct { + partition2AllPotentialConsumers map[topicPartitionAssignment][]string + } + tests := []struct { + name string + args args + want []topicPartitionAssignment + }{ + { + name: "Single topic partition", + args: args{ + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{ + Topic: "t1", + Partition: 0, + }: []string{"c1", "c2"}, + }, + }, + want: []topicPartitionAssignment{ + topicPartitionAssignment{ + Topic: "t1", + Partition: 0, + }, + }, + }, + { + name: "Multiple topic partitions with the same number of consumers but different topic names", + args: args{ + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{ + Topic: "t1", + Partition: 0, + }: []string{"c1", "c2"}, + topicPartitionAssignment{ + Topic: "t2", + Partition: 0, + }: []string{"c1", "c2"}, + }, + }, + want: []topicPartitionAssignment{ + topicPartitionAssignment{ + Topic: "t1", + Partition: 0, + }, + topicPartitionAssignment{ + Topic: "t2", + Partition: 0, + }, + }, + }, + { + name: "Multiple topic partitions with the same number of consumers and topic names", + args: args{ + partition2AllPotentialConsumers: map[topicPartitionAssignment][]string{ + topicPartitionAssignment{ + Topic: "t1", + Partition: 0, + }: []string{"c1", "c2"}, + topicPartitionAssignment{ + Topic: "t1", + Partition: 1, + }: []string{"c1", "c2"}, + }, + }, + want: []topicPartitionAssignment{ + topicPartitionAssignment{ + Topic: "t1", + Partition: 0, + }, + topicPartitionAssignment{ + Topic: "t1", + Partition: 1, + }, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := sortPartitionsByPotentialConsumerAssignments(tt.args.partition2AllPotentialConsumers); !reflect.DeepEqual(got, tt.want) { + t.Errorf("sortPartitionsByPotentialConsumerAssignments() = %v, want %v", got, tt.want) + } + }) + } +} From 79f734c100417addd882b16a6f7fa0d899126ff3 Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Tue, 13 Aug 2019 11:30:03 -0700 Subject: [PATCH 7/9] Initialize slice of consumers for a topic partition using an empty string slice literal --- balance_strategy.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/balance_strategy.go b/balance_strategy.go index 26ac514bb..2f405100c 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -187,7 +187,7 @@ func (s *stickyBalanceStrategy) Plan(members map[string]ConsumerGroupMemberMetad partition2AllPotentialConsumers := make(map[topicPartitionAssignment][]string) for topic, partitions := range topics { for _, partition := range partitions { - partition2AllPotentialConsumers[topicPartitionAssignment{Topic: topic, Partition: partition}] = make([]string, 0) + partition2AllPotentialConsumers[topicPartitionAssignment{Topic: topic, Partition: partition}] = []string{} } } From e46c30deb27ff01c05f5865afc611a0cf503f679 Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Wed, 14 Aug 2019 09:21:00 -0700 Subject: [PATCH 8/9] Correctly identify consumers that are no longer subscribed to topics --- balance_strategy.go | 87 +++++++++++++++++++--------------------- balance_strategy_test.go | 34 ++++++++++++++++ 2 files changed, 75 insertions(+), 46 deletions(-) diff --git a/balance_strategy.go b/balance_strategy.go index 2f405100c..d6604843d 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -214,54 +214,45 @@ func (s *stickyBalanceStrategy) Plan(members map[string]ConsumerGroupMemberMetad } // create a mapping of each partition to its current consumer, where possible - currentPartitionConsumer := make(map[topicPartitionAssignment]string, len(currentAssignment)) - for memberID, partitions := range currentAssignment { - for _, partition := range partitions { - currentPartitionConsumer[partition] = memberID - } + currentPartitionConsumers := make(map[topicPartitionAssignment]string, len(currentAssignment)) + unvisitedPartitions := make(map[topicPartitionAssignment][]string, len(partition2AllPotentialConsumers)) + for partition := range partition2AllPotentialConsumers { + unvisitedPartitions[partition] = []string{} } - - // sort the topic partitions in order of priority for reassignment - sortedPartitions := sortPartitions(currentAssignment, prevAssignment, isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions) - unassignedPartitions := deepCopyPartitions(sortedPartitions) + var unassignedPartitions []topicPartitionAssignment for memberID, partitions := range currentAssignment { - // if a consumer that existed before (and had some partition assignments) is now removed, remove it from currentAssignment - if _, exists := members[memberID]; !exists { - for _, partition := range partitions { - delete(currentPartitionConsumer, partition) - } - delete(currentAssignment, memberID) - continue - } - - // otherwise (the consumer still exists) - updatedPartitions := deepCopyPartitions(partitions) + var keepPartitions []topicPartitionAssignment for _, partition := range partitions { + // If this partition no longer exists at all, likely due to the + // topic being deleted, we remove the partition from the member. if _, exists := partition2AllPotentialConsumers[partition]; !exists { - // if this topic partition of this consumer no longer exists remove it from currentAssignment of the consumer - updatedPartitions = removeTopicPartitionFromMemberAssignments(updatedPartitions, partition) - delete(currentPartitionConsumer, partition) - } else if _, exists := topics[partition.Topic]; !exists { - // if this partition cannot remain assigned to its current consumer because the consumer - // is no longer subscribed to its topic remove it from currentAssignment of the consumer - updatedPartitions = removeTopicPartitionFromMemberAssignments(updatedPartitions, partition) - } else { - // otherwise, remove the topic partition from those that need to be assigned only if - // its current consumer is still subscribed to its topic (because it is already assigned - // and we would want to preserve that assignment as much as possible) - unassignedPartitions = removeTopicPartitionFromMemberAssignments(unassignedPartitions, partition) + continue } + delete(unvisitedPartitions, partition) + currentPartitionConsumers[partition] = memberID + + if !strsContains(members[memberID].Topics, partition.Topic) { + unassignedPartitions = append(unassignedPartitions, partition) + continue + } + keepPartitions = append(keepPartitions, partition) } - currentAssignment[memberID] = updatedPartitions + currentAssignment[memberID] = keepPartitions + } + for unvisited := range unvisitedPartitions { + unassignedPartitions = append(unassignedPartitions, unvisited) } + // sort the topic partitions in order of priority for reassignment + sortedPartitions := sortPartitions(currentAssignment, prevAssignment, isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions) + // at this point we have preserved all valid topic partition to consumer assignments and removed // all invalid topic partitions and invalid consumers. Now we need to assign unassignedPartitions // to consumers so that the topic partition assignments are as balanced as possible. // an ascending sorted set of consumers based on how many topic partitions are already assigned to them sortedCurrentSubscriptions := sortMemberIDsByPartitionAssignments(currentAssignment) - s.balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer) + s.balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumers) // Assemble plan plan := make(BalanceStrategyPlan, len(currentAssignment)) @@ -277,6 +268,15 @@ func (s *stickyBalanceStrategy) Plan(members map[string]ConsumerGroupMemberMetad return plan, nil } +func strsContains(s []string, value string) bool { + for _, entry := range s { + if entry == value { + return true + } + } + return false +} + // Balance assignments across consumers for maximum fairness and stickiness. func (s *stickyBalanceStrategy) balance(currentAssignment map[string][]topicPartitionAssignment, prevAssignment map[topicPartitionAssignment]consumerGenerationPair, sortedPartitions []topicPartitionAssignment, unassignedPartitions []topicPartitionAssignment, sortedCurrentSubscriptions []string, consumer2AllPotentialPartitions map[string][]topicPartitionAssignment, partition2AllPotentialConsumers map[topicPartitionAssignment][]string, currentPartitionConsumer map[topicPartitionAssignment]string) { initializing := false @@ -601,28 +601,23 @@ func sortPartitions(currentAssignment map[string][]topicPartitionAssignment, par member := pq[0] // partitions that were assigned to a different consumer last time - var prevPartition *topicPartitionAssignment var prevPartitionIndex int for i, partition := range member.assignments { if _, exists := partitionsWithADifferentPreviousAssignment[partition]; exists { - prevPartition = &partition prevPartitionIndex = i break } } - if prevPartition != nil { - // if there is a partition on this consumer that was assigned to another consumer before mark it as good options for reassignment - member.assignments = append(member.assignments[:prevPartitionIndex], member.assignments[prevPartitionIndex+1:]...) - sortedPartitions = append(sortedPartitions, *prevPartition) - delete(unassignedPartitions, *prevPartition) - heap.Fix(&pq, 0) - } else if len(member.assignments) > 0 { - // otherwise, mark any other one of the current partitions as a reassignment candidate - partition := member.assignments[0] - member.assignments = append(member.assignments[:0], member.assignments[1:]...) + if len(member.assignments) > 0 { + partition := member.assignments[prevPartitionIndex] sortedPartitions = append(sortedPartitions, partition) delete(unassignedPartitions, partition) + if prevPartitionIndex == 0 { + member.assignments = member.assignments[1:] + } else { + member.assignments = append(member.assignments[:prevPartitionIndex], member.assignments[prevPartitionIndex+1:]...) + } heap.Fix(&pq, 0) } else { heap.Pop(&pq) diff --git a/balance_strategy_test.go b/balance_strategy_test.go index 72b0492dd..ab48af39a 100644 --- a/balance_strategy_test.go +++ b/balance_strategy_test.go @@ -1188,6 +1188,40 @@ func Test_stickyBalanceStrategy_Plan(t *testing.T) { }, }, }, + { + name: "One consumer that is no longer subscribed to a topic that it had previously been consuming from", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic2"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{0}}, 1), + }, + }, + topics: map[string][]int32{ + "topic1": []int32{0}, + "topic2": []int32{0}, + }, + }, + }, + { + name: "Two consumers where one is no longer interested in consuming from a topic that it had been consuming from", + args: args{ + members: map[string]ConsumerGroupMemberMetadata{ + "consumer1": ConsumerGroupMemberMetadata{ + Topics: []string{"topic2"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{0}}, 1), + }, + "consumer2": ConsumerGroupMemberMetadata{ + Topics: []string{"topic1", "topic2"}, + UserData: encodeSubscriberPlanWithGeneration(t, map[string][]int32{"topic1": []int32{1}}, 1), + }, + }, + topics: map[string][]int32{ + "topic1": []int32{0, 1}, + "topic2": []int32{0, 1}, + }, + }, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { From b29086bdaae0da7ce71eae3f854d50685fd6b631 Mon Sep 17 00:00:00 2001 From: Scott Kidder Date: Wed, 14 Aug 2019 10:23:29 -0700 Subject: [PATCH 9/9] Change type of unvisitedPartitions from []string to bool to economize --- balance_strategy.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/balance_strategy.go b/balance_strategy.go index d6604843d..67c4d96d0 100644 --- a/balance_strategy.go +++ b/balance_strategy.go @@ -215,9 +215,9 @@ func (s *stickyBalanceStrategy) Plan(members map[string]ConsumerGroupMemberMetad // create a mapping of each partition to its current consumer, where possible currentPartitionConsumers := make(map[topicPartitionAssignment]string, len(currentAssignment)) - unvisitedPartitions := make(map[topicPartitionAssignment][]string, len(partition2AllPotentialConsumers)) + unvisitedPartitions := make(map[topicPartitionAssignment]bool, len(partition2AllPotentialConsumers)) for partition := range partition2AllPotentialConsumers { - unvisitedPartitions[partition] = []string{} + unvisitedPartitions[partition] = true } var unassignedPartitions []topicPartitionAssignment for memberID, partitions := range currentAssignment {