rajinisivaram commented on code in PR #12990:
URL: https://github.com/apache/kafka/pull/12990#discussion_r1112013804
##########
clients/src/main/java/org/apache/kafka/clients/consumer/RangeAssignor.java:
##########
@@ -76,43 +99,185 @@ private Map<String, List<MemberInfo>>
consumersPerTopic(Map<String, Subscription
Map<String, List<MemberInfo>> topicToConsumers = new HashMap<>();
for (Map.Entry<String, Subscription> subscriptionEntry :
consumerMetadata.entrySet()) {
String consumerId = subscriptionEntry.getKey();
- MemberInfo memberInfo = new MemberInfo(consumerId,
subscriptionEntry.getValue().groupInstanceId());
- for (String topic : subscriptionEntry.getValue().topics()) {
+ Subscription subscription = subscriptionEntry.getValue();
+ MemberInfo memberInfo = new MemberInfo(consumerId,
subscription.groupInstanceId(), subscription.rackId());
+ for (String topic : subscription.topics()) {
put(topicToConsumers, topic, memberInfo);
}
}
return topicToConsumers;
}
@Override
- public Map<String, List<TopicPartition>> assign(Map<String, Integer>
partitionsPerTopic,
- Map<String, Subscription>
subscriptions) {
+ public Map<String, List<TopicPartition>> assignPartitions(Map<String,
List<PartitionInfo>> partitionsPerTopic,
+ Map<String,
Subscription> subscriptions) {
Map<String, List<MemberInfo>> consumersPerTopic =
consumersPerTopic(subscriptions);
+ List<TopicAssignmentState> topicAssignmentStates =
partitionsPerTopic.entrySet().stream()
+ .filter(e -> !e.getValue().isEmpty())
+ .map(e -> new TopicAssignmentState(e.getKey(), e.getValue(),
consumersPerTopic.get(e.getKey())))
+ .collect(Collectors.toList());
Map<String, List<TopicPartition>> assignment = new HashMap<>();
for (String memberId : subscriptions.keySet())
assignment.put(memberId, new ArrayList<>());
- for (Map.Entry<String, List<MemberInfo>> topicEntry :
consumersPerTopic.entrySet()) {
- String topic = topicEntry.getKey();
- List<MemberInfo> consumersForTopic = topicEntry.getValue();
+ boolean useRackAware = topicAssignmentStates.stream().anyMatch(t ->
t.needsRackAwareAssignment);
+ if (useRackAware)
+ assignWithRackMatching(topicAssignmentStates, assignment);
+
+ topicAssignmentStates.forEach(t -> assignRanges(t, (c, tp) -> true,
assignment));
+
+ if (useRackAware)
+ assignment.values().forEach(list ->
list.sort(PARTITION_COMPARATOR));
+ return assignment;
+ }
+
+ // This method is not used, but retained for compatibility with any custom
assignors that extend this class.
+ @Override
+ public Map<String, List<TopicPartition>> assign(Map<String, Integer>
partitionsPerTopic,
+ Map<String, Subscription>
subscriptions) {
+ return
assignPartitions(partitionInfosWithoutRacks(partitionsPerTopic), subscriptions);
+ }
+
+ private void assignRanges(TopicAssignmentState assignmentState,
+ BiFunction<String, TopicPartition, Boolean>
mayAssign,
+ Map<String, List<TopicPartition>> assignment) {
+ for (String consumer : assignmentState.consumers) {
+ if (assignmentState.unassignedPartitions.isEmpty())
+ break;
+ List<TopicPartition> assignablePartitions =
assignmentState.unassignedPartitions.stream()
+ .filter(tp -> mayAssign.apply(consumer, tp))
+ .collect(Collectors.toList());
- Integer numPartitionsForTopic = partitionsPerTopic.get(topic);
- if (numPartitionsForTopic == null)
+ int maxAssignable =
Math.min(assignmentState.maxAssignable(consumer), assignablePartitions.size());
+ if (maxAssignable <= 0)
continue;
- Collections.sort(consumersForTopic);
+ assign(consumer, assignablePartitions.subList(0, maxAssignable),
assignmentState, assignment);
+ }
+ }
+
+ private void assignWithRackMatching(Collection<TopicAssignmentState>
assignmentStates,
+ Map<String, List<TopicPartition>>
assignment) {
- int numPartitionsPerConsumer = numPartitionsForTopic /
consumersForTopic.size();
- int consumersWithExtraPartition = numPartitionsForTopic %
consumersForTopic.size();
+ assignmentStates.stream().collect(Collectors.groupingBy(t ->
t.consumers)).forEach((consumers, states) -> {
+ states.stream().collect(Collectors.groupingBy(t ->
t.partitionRacks.size())).forEach((numPartitions, coPartitionedStates) -> {
+ if (coPartitionedStates.size() > 1)
+ assignCoPartitionedWithRackMatching(consumers,
numPartitions, states, assignment);
+ else {
+ TopicAssignmentState state = coPartitionedStates.get(0);
+ if (state.needsRackAwareAssignment)
+ assignRanges(state, state::racksMatch, assignment);
+ }
+ });
+ });
+ }
+
+ private void assignCoPartitionedWithRackMatching(List<String> consumers,
+ int numPartitions,
+
Collection<TopicAssignmentState> assignmentStates,
+ Map<String,
List<TopicPartition>> assignment) {
+
+ List<String> remainingConsumers = new LinkedList<>(consumers);
+ for (int i = 0; i < numPartitions; i++) {
+ int p = i;
- List<TopicPartition> partitions =
AbstractPartitionAssignor.partitions(topic, numPartitionsForTopic);
- for (int i = 0, n = consumersForTopic.size(); i < n; i++) {
- int start = numPartitionsPerConsumer * i + Math.min(i,
consumersWithExtraPartition);
- int length = numPartitionsPerConsumer + (i + 1 >
consumersWithExtraPartition ? 0 : 1);
-
assignment.get(consumersForTopic.get(i).memberId).addAll(partitions.subList(start,
start + length));
+ Optional<String> matchingConsumer = remainingConsumers.stream()
+ .filter(c -> assignmentStates.stream().allMatch(t ->
t.racksMatch(c, new TopicPartition(t.topic, p)) && t.maxAssignable(c) > 0))
+ .findFirst();
+ if (matchingConsumer.isPresent()) {
+ String consumer = matchingConsumer.get();
+ assignmentStates.forEach(t -> assign(consumer,
Collections.singletonList(new TopicPartition(t.topic, p)), t, assignment));
+
+ if (assignmentStates.stream().noneMatch(t ->
t.maxAssignable(consumer) > 0)) {
+ remainingConsumers.remove(consumer);
+ if (remainingConsumers.isEmpty())
+ break;
+ }
}
}
- return assignment;
+ }
+
+ private void assign(String consumer, List<TopicPartition> partitions,
TopicAssignmentState assignmentState, Map<String, List<TopicPartition>>
assignment) {
+ assignment.get(consumer).addAll(partitions);
+ assignmentState.onAssigned(consumer, partitions);
+ }
+
+ private class TopicAssignmentState {
+ private final String topic;
+ private final List<String> consumers;
+ private final boolean needsRackAwareAssignment;
+ private final Map<TopicPartition, Set<String>> partitionRacks;
+ private final Map<String, String> consumerRacks;
Review Comment:
It was useful to have the consumer racks along with the partition racks for
each topic for matching racks. But we don't need to populate it for each topic.
I have removed this map and updated the list of consumers of the topic to
include the rack.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscr...@kafka.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
