feyman2016 commented on a change in pull request #8832:
URL: https://github.com/apache/kafka/pull/8832#discussion_r437532739
##########
File path:
streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
##########
@@ -513,66 +553,94 @@ private boolean checkMetadataVersions(final int
minReceivedMetadataVersion,
*/
private void setRepartitionTopicMetadataNumberOfPartitions(final
Map<String, InternalTopicConfig> repartitionTopicMetadata,
final
Map<Integer, TopicsInfo> topicGroups,
- final Cluster
metadata) {
- boolean numPartitionsNeeded;
- do {
- numPartitionsNeeded = false;
-
- for (final TopicsInfo topicsInfo : topicGroups.values()) {
- for (final String topicName :
topicsInfo.repartitionSourceTopics.keySet()) {
- final Optional<Integer> maybeNumPartitions =
repartitionTopicMetadata.get(topicName)
-
.numberOfPartitions();
- Integer numPartitions = null;
-
- if (!maybeNumPartitions.isPresent()) {
- // try set the number of partitions for this
repartition topic if it is not set yet
- for (final TopicsInfo otherTopicsInfo :
topicGroups.values()) {
- final Set<String> otherSinkTopics =
otherTopicsInfo.sinkTopics;
-
- if (otherSinkTopics.contains(topicName)) {
- // if this topic is one of the sink topics of
this topology,
- // use the maximum of all its source topic
partitions as the number of partitions
- for (final String sourceTopicName :
otherTopicsInfo.sourceTopics) {
- Integer numPartitionsCandidate = null;
- // It is possible the sourceTopic is
another internal topic, i.e,
- // map().join().join(map())
- if
(repartitionTopicMetadata.containsKey(sourceTopicName)) {
- if
(repartitionTopicMetadata.get(sourceTopicName).numberOfPartitions().isPresent())
{
- numPartitionsCandidate =
-
repartitionTopicMetadata.get(sourceTopicName).numberOfPartitions().get();
- }
- } else {
- final Integer count =
metadata.partitionCountForTopic(sourceTopicName);
- if (count == null) {
- throw new IllegalStateException(
- "No partition count found for
source topic "
- + sourceTopicName
- + ", but it should have
been."
- );
- }
- numPartitionsCandidate = count;
- }
-
- if (numPartitionsCandidate != null) {
- if (numPartitions == null ||
numPartitionsCandidate > numPartitions) {
- numPartitions =
numPartitionsCandidate;
- }
- }
- }
- }
- }
-
- // if we still have not found the right number of
partitions,
- // another iteration is needed
- if (numPartitions == null) {
- numPartitionsNeeded = true;
- } else {
-
repartitionTopicMetadata.get(topicName).setNumberOfPartitions(numPartitions);
- }
- }
+ final Cluster
metadata) {
+ final Set<String> allRepartitionSourceTopics = new HashSet<>();
+ final Map<String, TopicNode> builtTopicNodes = new HashMap<>();
+ // 1. Build a graph contains the TopicsInfo and TopicsNode
+ for (final TopicsInfo topicsInfo : topicGroups.values()) {
+ for (final String topicName :
topicsInfo.repartitionSourceTopics.keySet()) {
+ allRepartitionSourceTopics.add(topicName);
+ }
+ for (final String sourceTopic : topicsInfo.sourceTopics) {
+ if (builtTopicNodes.containsKey(sourceTopic)) {
+
builtTopicNodes.get(sourceTopic).addDownStreamTopicsInfo(topicsInfo);
+ } else {
+ final TopicNode topicNode = new TopicNode(sourceTopic);
+ topicNode.addDownStreamTopicsInfo(topicsInfo);
+ builtTopicNodes.put(sourceTopic, topicNode);
+ }
+ }
+
+ for (final String sinkTopic : topicsInfo.sinkTopics) {
+ if (builtTopicNodes.containsKey(sinkTopic)) {
+
builtTopicNodes.get(sinkTopic).addUpStreamTopicsInfo(topicsInfo);
+ } else {
+ final TopicNode topicNode = new TopicNode(sinkTopic);
+ topicNode.addUpStreamTopicsInfo(topicsInfo);
+ builtTopicNodes.put(sinkTopic, topicNode);
}
}
- } while (numPartitionsNeeded);
+ }
+
+ // 2. Use DFS along with memoization to calc repartition number of all
repartitionSourceTopics
+ for (final String topic : allRepartitionSourceTopics) {
+ calcRepartitionNumForTopic(topic, repartitionTopicMetadata,
metadata, builtTopicNodes, new HashMap<TopicsInfo, Integer>());
+ }
+ }
+
+ private int calcRepartitionNumForTopic(final String topic,
+ final Map<String,
InternalTopicConfig> repartitionTopicMetadata,
+ final Cluster metadata,
+ final Map<String, TopicNode>
builtTopicNodes,
+ final Map<TopicsInfo, Integer>
topicsInfoNumberOfPartitions) {
+
+ if (repartitionTopicMetadata.containsKey(topic)) {
+ final Optional<Integer> maybeNumberPartitions =
repartitionTopicMetadata.get(topic).numberOfPartitions();
+ // if numberOfPartitions already calculated, return directly
+ if (maybeNumberPartitions.isPresent()) {
+ return maybeNumberPartitions.get();
+ } else {
+ // else calculate the max numRepartitions of its upStream
TopicsInfo and set the repartitionTopicMetadata for memoization before return
+ final TopicNode topicNode = builtTopicNodes.get(topic);
+ Integer maxNumberPartitions = 0;
+ for (final TopicsInfo upstream : topicNode.upStreams) {
+ final Integer upStreamRepartitionNum =
calcRepartitionNumForTopicInfo(upstream, repartitionTopicMetadata,
+ metadata, builtTopicNodes,
topicsInfoNumberOfPartitions);
+ maxNumberPartitions = upStreamRepartitionNum >
maxNumberPartitions ? upStreamRepartitionNum : maxNumberPartitions;
+ }
+
repartitionTopicMetadata.get(topic).setNumberOfPartitions(maxNumberPartitions);
+ return maxNumberPartitions;
+ }
+ } else {
+ final Integer count = metadata.partitionCountForTopic(topic);
+ if (count == null) {
+ throw new IllegalStateException(
+ "No partition count found for source topic "
+ + topic
+ + ", but it should have been."
+ );
+ }
+ return count;
+ }
+ }
+
+ private int calcRepartitionNumForTopicInfo(final TopicsInfo topicsInfo,
Review comment:
Yeah, seems to be the case, but in the mutual recursive function, it
seems acceptable...
Anyway, I can rename them later if think the algo is finalized
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
