diff --git a/clients/src/main/java/org/apache/kafka/common/PartitionInfo.java b/clients/src/main/java/org/apache/kafka/common/PartitionInfo.java index 7ed54f098a3c8..29bedff0cf19b 100644 --- a/clients/src/main/java/org/apache/kafka/common/PartitionInfo.java +++ b/clients/src/main/java/org/apache/kafka/common/PartitionInfo.java @@ -27,7 +27,6 @@ public class PartitionInfo { private final Node[] inSyncReplicas; private final Node[] offlineReplicas; - // Used only by tests public PartitionInfo(String topic, int partition, Node leader, Node[] replicas, Node[] inSyncReplicas) { this(topic, partition, leader, replicas, inSyncReplicas, new Node[0]); } @@ -113,5 +112,4 @@ private String formatNodeIds(Node[] nodes) { b.append("]"); return b.toString(); } - } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java index 622bf0887319e..bb18f92ef23d2 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java @@ -1787,7 +1787,7 @@ public int hashCode() { } public static class TopicsInfo { - final Set sinkTopics; + public final Set sinkTopics; public final Set sourceTopics; public final Map stateChangelogTopics; public final Map repartitionSourceTopics; diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RepartitionTopics.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RepartitionTopics.java new file mode 100644 index 0000000000000..889ee7c2df135 --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RepartitionTopics.java @@ -0,0 +1,212 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.streams.processor.internals; + +import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.PartitionInfo; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.utils.LogContext; +import org.apache.kafka.streams.errors.MissingSourceTopicException; +import org.apache.kafka.streams.errors.TaskAssignmentException; +import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder.TopicsInfo; +import org.apache.kafka.streams.processor.internals.assignment.CopartitionedTopicsEnforcer; +import org.slf4j.Logger; + +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +public class RepartitionTopics { + + private final InternalTopicManager internalTopicManager; + private final InternalTopologyBuilder internalTopologyBuilder; + private final Cluster clusterMetadata; + private final CopartitionedTopicsEnforcer copartitionedTopicsEnforcer; + private final Logger log; + private final Map topicPartitionInfos = new HashMap<>(); + + public RepartitionTopics(final InternalTopologyBuilder internalTopologyBuilder, + final InternalTopicManager internalTopicManager, + final CopartitionedTopicsEnforcer copartitionedTopicsEnforcer, + final Cluster clusterMetadata, + final String logPrefix) { + this.internalTopologyBuilder = internalTopologyBuilder; + this.internalTopicManager = internalTopicManager; + this.clusterMetadata = clusterMetadata; + this.copartitionedTopicsEnforcer = copartitionedTopicsEnforcer; + final LogContext logContext = new LogContext(logPrefix); + log = logContext.logger(getClass()); + } + + public void setup() { + final Map topicGroups = internalTopologyBuilder.topicGroups(); + final Map repartitionTopicMetadata = computeRepartitionTopicConfig(topicGroups, clusterMetadata); + + // ensure the co-partitioning topics within the group have the same number of partitions, + // and enforce the number of partitions for those repartition topics to be the same if they + // are co-partitioned as well. + ensureCopartitioning(internalTopologyBuilder.copartitionGroups(), repartitionTopicMetadata, clusterMetadata); + + // make sure the repartition source topics exist with the right number of partitions, + // create these topics if necessary + internalTopicManager.makeReady(repartitionTopicMetadata); + + // augment the metadata with the newly computed number of partitions for all the + // repartition source topics + for (final Map.Entry entry : repartitionTopicMetadata.entrySet()) { + final String topic = entry.getKey(); + final int numPartitions = entry.getValue().numberOfPartitions().orElse(-1); + + for (int partition = 0; partition < numPartitions; partition++) { + topicPartitionInfos.put( + new TopicPartition(topic, partition), + new PartitionInfo(topic, partition, null, new Node[0], new Node[0]) + ); + } + } + } + + public Map topicPartitionsInfo() { + return Collections.unmodifiableMap(topicPartitionInfos); + } + + private Map computeRepartitionTopicConfig(final Map topicGroups, + final Cluster clusterMetadata) { + + final Map repartitionTopicConfigs = new HashMap<>(); + for (final TopicsInfo topicsInfo : topicGroups.values()) { + checkIfExternalSourceTopicsExist(topicsInfo, clusterMetadata); + repartitionTopicConfigs.putAll(topicsInfo.repartitionSourceTopics.values().stream() + .collect(Collectors.toMap(InternalTopicConfig::name, topicConfig -> topicConfig))); + } + setRepartitionSourceTopicPartitionCount(repartitionTopicConfigs, topicGroups, clusterMetadata); + + return repartitionTopicConfigs; + } + + private void ensureCopartitioning(final Collection> copartitionGroups, + final Map repartitionTopicMetadata, + final Cluster clusterMetadata) { + for (final Set copartitionGroup : copartitionGroups) { + copartitionedTopicsEnforcer.enforce(copartitionGroup, repartitionTopicMetadata, clusterMetadata); + } + } + + private void checkIfExternalSourceTopicsExist(final TopicsInfo topicsInfo, + final Cluster clusterMetadata) { + final Set missingExternalSourceTopics = new HashSet<>(topicsInfo.sourceTopics); + missingExternalSourceTopics.removeAll(topicsInfo.repartitionSourceTopics.keySet()); + missingExternalSourceTopics.removeAll(clusterMetadata.topics()); + if (!missingExternalSourceTopics.isEmpty()) { + log.error("The following source topics are missing/unknown: {}. Please make sure all source topics " + + "have been pre-created before starting the Streams application. ", + missingExternalSourceTopics); + throw new MissingSourceTopicException("Missing source topics."); + } + } + + /** + * Computes the number of partitions and sets it for each repartition topic in repartitionTopicMetadata + */ + private void setRepartitionSourceTopicPartitionCount(final Map repartitionTopicMetadata, + final Map topicGroups, + final Cluster clusterMetadata) { + boolean partitionCountNeeded; + do { + partitionCountNeeded = false; + boolean progressMadeThisIteration = false; // avoid infinitely looping without making any progress on unknown repartitions + + for (final TopicsInfo topicsInfo : topicGroups.values()) { + for (final String repartitionSourceTopic : topicsInfo.repartitionSourceTopics.keySet()) { + final Optional repartitionSourceTopicPartitionCount = + repartitionTopicMetadata.get(repartitionSourceTopic).numberOfPartitions(); + + if (!repartitionSourceTopicPartitionCount.isPresent()) { + final Integer numPartitions = computePartitionCount( + repartitionTopicMetadata, + topicGroups, + clusterMetadata, + repartitionSourceTopic + ); + + if (numPartitions == null) { + partitionCountNeeded = true; + log.trace("Unable to determine number of partitions for {}, another iteration is needed", + repartitionSourceTopic); + } else { + repartitionTopicMetadata.get(repartitionSourceTopic).setNumberOfPartitions(numPartitions); + progressMadeThisIteration = true; + } + } + } + } + if (!progressMadeThisIteration && partitionCountNeeded) { + throw new TaskAssignmentException("Failed to compute number of partitions for all repartition topics"); + } + } while (partitionCountNeeded); + } + + private Integer computePartitionCount(final Map repartitionTopicMetadata, + final Map topicGroups, + final Cluster clusterMetadata, + final String repartitionSourceTopic) { + Integer partitionCount = null; + // try set the number of partitions for this repartition topic if it is not set yet + for (final TopicsInfo topicsInfo : topicGroups.values()) { + final Set sinkTopics = topicsInfo.sinkTopics; + + if (sinkTopics.contains(repartitionSourceTopic)) { + // 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 upstreamSourceTopic : topicsInfo.sourceTopics) { + Integer numPartitionsCandidate = null; + // It is possible the sourceTopic is another internal topic, i.e, + // map().join().join(map()) + if (repartitionTopicMetadata.containsKey(upstreamSourceTopic)) { + if (repartitionTopicMetadata.get(upstreamSourceTopic).numberOfPartitions().isPresent()) { + numPartitionsCandidate = + repartitionTopicMetadata.get(upstreamSourceTopic).numberOfPartitions().get(); + } + } else { + final Integer count = clusterMetadata.partitionCountForTopic(upstreamSourceTopic); + if (count == null) { + throw new TaskAssignmentException( + "No partition count found for source topic " + + upstreamSourceTopic + + ", but it should have been." + ); + } + numPartitionsCandidate = count; + } + + if (numPartitionsCandidate != null) { + if (partitionCount == null || numPartitionsCandidate > partitionCount) { + partitionCount = numPartitionsCandidate; + } + } + } + } + } + return partitionCount; + } +} diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 1065b0297f0d9..fd4b7b705ae63 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -64,7 +64,6 @@ import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.PriorityQueue; import java.util.Queue; import java.util.Set; @@ -363,18 +362,15 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr // parse the topology to determine the repartition source topics, // making sure they are created with the number of partitions as // the maximum of the depending sub-topologies source topics' number of partitions - final Map topicGroups = taskManager.builder().topicGroups(); - - final Map allRepartitionTopicPartitions = prepareRepartitionTopics(topicGroups, metadata); - + final Map allRepartitionTopicPartitions = prepareRepartitionTopics(metadata); final Cluster fullMetadata = metadata.withPartitions(allRepartitionTopicPartitions); - log.debug("Created repartition topics {} from the parsed topology.", allRepartitionTopicPartitions.values()); // ---------------- Step Two ---------------- // // construct the assignment of tasks to clients + final Map topicGroups = taskManager.builder().topicGroups(); final Set allSourceTopics = new HashSet<>(); final Map> sourceTopicsByGroup = new HashMap<>(); for (final Map.Entry entry : topicGroups.entrySet()) { @@ -471,135 +467,22 @@ private boolean checkMetadataVersions(final int minReceivedMetadataVersion, return versionProbing; } - /** - * @return a map of repartition topics and their metadata - */ - private Map computeRepartitionTopicMetadata(final Map topicGroups, - final Cluster metadata) { - final Map repartitionTopicMetadata = new HashMap<>(); - for (final TopicsInfo topicsInfo : topicGroups.values()) { - for (final String topic : topicsInfo.sourceTopics) { - if (!topicsInfo.repartitionSourceTopics.containsKey(topic) && !metadata.topics().contains(topic)) { - log.error("Source topic {} is missing/unknown during rebalance, please make sure all source topics " + - "have been pre-created before starting the Streams application. Returning error {}", - topic, AssignorError.INCOMPLETE_SOURCE_TOPIC_METADATA.name()); - throw new MissingSourceTopicException("Missing source topic during assignment."); - } - } - for (final InternalTopicConfig topic : topicsInfo.repartitionSourceTopics.values()) { - repartitionTopicMetadata.put(topic.name(), topic); - } - } - return repartitionTopicMetadata; - } - /** * Computes and assembles all repartition topic metadata then creates the topics if necessary. * * @return map from repartition topic to its partition info */ - private Map prepareRepartitionTopics(final Map topicGroups, - final Cluster metadata) { - final Map repartitionTopicMetadata = computeRepartitionTopicMetadata(topicGroups, metadata); - - setRepartitionTopicMetadataNumberOfPartitions(repartitionTopicMetadata, topicGroups, metadata); - - // ensure the co-partitioning topics within the group have the same number of partitions, - // and enforce the number of partitions for those repartition topics to be the same if they - // are co-partitioned as well. - ensureCopartitioning(taskManager.builder().copartitionGroups(), repartitionTopicMetadata, metadata); - - // make sure the repartition source topics exist with the right number of partitions, - // create these topics if necessary - internalTopicManager.makeReady(repartitionTopicMetadata); - - // augment the metadata with the newly computed number of partitions for all the - // repartition source topics - final Map allRepartitionTopicPartitions = new HashMap<>(); - for (final Map.Entry entry : repartitionTopicMetadata.entrySet()) { - final String topic = entry.getKey(); - final int numPartitions = entry.getValue().numberOfPartitions().orElse(-1); - - for (int partition = 0; partition < numPartitions; partition++) { - allRepartitionTopicPartitions.put( - new TopicPartition(topic, partition), - new PartitionInfo(topic, partition, null, new Node[0], new Node[0]) - ); - } - } - return allRepartitionTopicPartitions; - } - - /** - * Computes the number of partitions and sets it for each repartition topic in repartitionTopicMetadata - */ - private void setRepartitionTopicMetadataNumberOfPartitions(final Map repartitionTopicMetadata, - final Map topicGroups, - final Cluster metadata) { - boolean numPartitionsNeeded; - do { - numPartitionsNeeded = false; - boolean progressMadeThisIteration = false; // avoid infinitely looping without making any progress on unknown repartitions - - for (final TopicsInfo topicsInfo : topicGroups.values()) { - for (final String repartitionSourceTopic : topicsInfo.repartitionSourceTopics.keySet()) { - final Optional maybeNumPartitions = repartitionTopicMetadata.get(repartitionSourceTopic) - .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 otherSinkTopics = otherTopicsInfo.sinkTopics; - - if (otherSinkTopics.contains(repartitionSourceTopic)) { - // 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 upstreamSourceTopic : otherTopicsInfo.sourceTopics) { - Integer numPartitionsCandidate = null; - // It is possible the sourceTopic is another internal topic, i.e, - // map().join().join(map()) - if (repartitionTopicMetadata.containsKey(upstreamSourceTopic)) { - if (repartitionTopicMetadata.get(upstreamSourceTopic).numberOfPartitions().isPresent()) { - numPartitionsCandidate = - repartitionTopicMetadata.get(upstreamSourceTopic).numberOfPartitions().get(); - } - } else { - final Integer count = metadata.partitionCountForTopic(upstreamSourceTopic); - if (count == null) { - throw new TaskAssignmentException( - "No partition count found for source topic " - + upstreamSourceTopic - + ", but it should have been." - ); - } - numPartitionsCandidate = count; - } - - if (numPartitionsCandidate != null) { - if (numPartitions == null || numPartitionsCandidate > numPartitions) { - numPartitions = numPartitionsCandidate; - } - } - } - } - } - - if (numPartitions == null) { - numPartitionsNeeded = true; - log.trace("Unable to determine number of partitions for {}, another iteration is needed", - repartitionSourceTopic); - } else { - repartitionTopicMetadata.get(repartitionSourceTopic).setNumberOfPartitions(numPartitions); - progressMadeThisIteration = true; - } - } - } - } - if (!progressMadeThisIteration && numPartitionsNeeded) { - throw new TaskAssignmentException("Failed to compute number of partitions for all repartition topics"); - } - } while (numPartitionsNeeded); + private Map prepareRepartitionTopics(final Cluster metadata) { + + final RepartitionTopics repartitionTopics = new RepartitionTopics( + taskManager.builder(), + internalTopicManager, + copartitionedTopicsEnforcer, + metadata, + logPrefix + ); + repartitionTopics.setup(); + return repartitionTopics.topicPartitionsInfo(); } /** @@ -1568,14 +1451,6 @@ private static void validateActiveTaskEncoding(final List partit } } - private void ensureCopartitioning(final Collection> copartitionGroups, - final Map allRepartitionTopicsNumPartitions, - final Cluster metadata) { - for (final Set copartitionGroup : copartitionGroups) { - copartitionedTopicsEnforcer.enforce(copartitionGroup, allRepartitionTopicsNumPartitions, metadata); - } - } - private int updateMinReceivedVersion(final int usedVersion, final int minReceivedMetadataVersion) { return Math.min(usedVersion, minReceivedMetadataVersion); } diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/RepartitionTopicsTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/RepartitionTopicsTest.java new file mode 100644 index 0000000000000..61a3716fef43a --- /dev/null +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/RepartitionTopicsTest.java @@ -0,0 +1,429 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.streams.processor.internals; + +import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.PartitionInfo; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.streams.errors.MissingSourceTopicException; +import org.apache.kafka.streams.errors.TaskAssignmentException; +import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder.TopicsInfo; +import org.apache.kafka.streams.processor.internals.assignment.CopartitionedTopicsEnforcer; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.powermock.core.classloader.annotations.PrepareForTest; +import org.powermock.modules.junit4.PowerMockRunner; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.apache.kafka.common.utils.Utils.mkEntry; +import static org.apache.kafka.common.utils.Utils.mkMap; +import static org.apache.kafka.common.utils.Utils.mkSet; +import static org.easymock.EasyMock.anyObject; +import static org.easymock.EasyMock.eq; +import static org.easymock.EasyMock.expect; +import static org.easymock.EasyMock.mock; +import static org.easymock.EasyMock.niceMock; +import static org.easymock.EasyMock.replay; +import static org.easymock.EasyMock.verify; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; + +@RunWith(PowerMockRunner.class) +@PrepareForTest({Cluster.class}) +public class RepartitionTopicsTest { + + private static final String SOURCE_TOPIC_NAME1 = "source1"; + private static final String SOURCE_TOPIC_NAME2 = "source2"; + private static final String SOURCE_TOPIC_NAME3 = "source3"; + private static final String SINK_TOPIC_NAME1 = "sink1"; + private static final String SINK_TOPIC_NAME2 = "sink2"; + private static final String REPARTITION_TOPIC_NAME1 = "repartition1"; + private static final String REPARTITION_TOPIC_NAME2 = "repartition2"; + private static final String REPARTITION_TOPIC_NAME3 = "repartition3"; + private static final String REPARTITION_TOPIC_NAME4 = "repartition4"; + private static final String REPARTITION_WITHOUT_PARTITION_COUNT = "repartitionWithoutPartitionCount"; + private static final String SOME_OTHER_TOPIC = "someOtherTopic"; + private static final Map TOPIC_CONFIG1 = Collections.singletonMap("config1", "val1"); + private static final Map TOPIC_CONFIG2 = Collections.singletonMap("config2", "val2"); + private static final Map TOPIC_CONFIG5 = Collections.singletonMap("config5", "val5"); + private static final RepartitionTopicConfig REPARTITION_TOPIC_CONFIG1 = + new RepartitionTopicConfig(REPARTITION_TOPIC_NAME1, TOPIC_CONFIG1, 4, true); + private static final RepartitionTopicConfig REPARTITION_TOPIC_CONFIG2 = + new RepartitionTopicConfig(REPARTITION_TOPIC_NAME2, TOPIC_CONFIG2, 2, true); + private static final TopicsInfo TOPICS_INFO1 = new TopicsInfo( + mkSet(REPARTITION_TOPIC_NAME1), + mkSet(SOURCE_TOPIC_NAME1, SOURCE_TOPIC_NAME2), + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1), + mkEntry(REPARTITION_TOPIC_NAME2, REPARTITION_TOPIC_CONFIG2) + ), + Collections.emptyMap() + ); + private static final TopicsInfo TOPICS_INFO2 = new TopicsInfo( + mkSet(SINK_TOPIC_NAME1), + mkSet(REPARTITION_TOPIC_NAME1), + mkMap(mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1)), + Collections.emptyMap() + ); + + final InternalTopologyBuilder internalTopologyBuilder = mock(InternalTopologyBuilder.class); + final InternalTopicManager internalTopicManager = mock(InternalTopicManager.class); + final CopartitionedTopicsEnforcer copartitionedTopicsEnforcer = mock(CopartitionedTopicsEnforcer.class); + final Cluster clusterMetadata = niceMock(Cluster.class); + + @Test + public void shouldSetupRepartitionTopics() { + expect(internalTopologyBuilder.topicGroups()) + .andReturn(mkMap(mkEntry(0, TOPICS_INFO1), mkEntry(1, TOPICS_INFO2))); + final Set coPartitionGroup1 = mkSet(SOURCE_TOPIC_NAME1, SOURCE_TOPIC_NAME2); + final Set coPartitionGroup2 = mkSet(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_NAME2); + final List> coPartitionGroups = Arrays.asList(coPartitionGroup1, coPartitionGroup2); + expect(internalTopologyBuilder.copartitionGroups()).andReturn(coPartitionGroups); + copartitionedTopicsEnforcer.enforce(eq(coPartitionGroup1), anyObject(), eq(clusterMetadata)); + copartitionedTopicsEnforcer.enforce(eq(coPartitionGroup2), anyObject(), eq(clusterMetadata)); + expect(internalTopicManager.makeReady( + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1), + mkEntry(REPARTITION_TOPIC_NAME2, REPARTITION_TOPIC_CONFIG2) + )) + ).andReturn(Collections.emptySet()); + setupCluster(); + replay(internalTopicManager, internalTopologyBuilder, clusterMetadata); + final RepartitionTopics repartitionTopics = new RepartitionTopics( + internalTopologyBuilder, + internalTopicManager, + copartitionedTopicsEnforcer, + clusterMetadata, + "[test] " + ); + + repartitionTopics.setup(); + + verify(internalTopicManager, internalTopologyBuilder); + final Map topicPartitionsInfo = repartitionTopics.topicPartitionsInfo(); + assertThat(topicPartitionsInfo.size(), is(6)); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 0); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 1); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 2); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 3); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME2, 0); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME2, 1); + } + + @Test + public void shouldThrowMissingSourceTopicException() { + expect(internalTopologyBuilder.topicGroups()) + .andReturn(mkMap(mkEntry(0, TOPICS_INFO1), mkEntry(1, TOPICS_INFO2))); + expect(internalTopologyBuilder.copartitionGroups()).andReturn(Collections.emptyList()); + copartitionedTopicsEnforcer.enforce(eq(Collections.emptySet()), anyObject(), eq(clusterMetadata)); + expect(internalTopicManager.makeReady( + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1) + )) + ).andReturn(Collections.emptySet()); + setupClusterWithMissingTopics(mkSet(SOURCE_TOPIC_NAME1)); + replay(internalTopicManager, internalTopologyBuilder, clusterMetadata); + final RepartitionTopics repartitionTopics = new RepartitionTopics( + internalTopologyBuilder, + internalTopicManager, + copartitionedTopicsEnforcer, + clusterMetadata, + "[test] " + ); + + assertThrows(MissingSourceTopicException.class, repartitionTopics::setup); + } + + @Test + public void shouldThrowTaskAssignmentExceptionIfPartitionCountCannotBeComputedForAllRepartitionTopics() { + final RepartitionTopicConfig repartitionTopicConfigWithoutPartitionCount = + new RepartitionTopicConfig(REPARTITION_WITHOUT_PARTITION_COUNT, TOPIC_CONFIG5); + expect(internalTopologyBuilder.topicGroups()) + .andReturn(mkMap( + mkEntry(0, TOPICS_INFO1), + mkEntry(1, setupTopicInfoWithRepartitionTopicWithoutPartitionCount(repartitionTopicConfigWithoutPartitionCount)) + )); + expect(internalTopologyBuilder.copartitionGroups()).andReturn(Collections.emptyList()); + copartitionedTopicsEnforcer.enforce(eq(Collections.emptySet()), anyObject(), eq(clusterMetadata)); + expect(internalTopicManager.makeReady( + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1) + )) + ).andReturn(Collections.emptySet()); + setupCluster(); + replay(internalTopicManager, internalTopologyBuilder, clusterMetadata); + final RepartitionTopics repartitionTopics = new RepartitionTopics( + internalTopologyBuilder, + internalTopicManager, + copartitionedTopicsEnforcer, + clusterMetadata, + "[test] " + ); + + final TaskAssignmentException exception = assertThrows(TaskAssignmentException.class, repartitionTopics::setup); + assertThat(exception.getMessage(), is("Failed to compute number of partitions for all repartition topics")); + } + + @Test + public void shouldThrowTaskAssignmentExceptionIfSourceTopicHasNoPartitionCount() { + final RepartitionTopicConfig repartitionTopicConfigWithoutPartitionCount = + new RepartitionTopicConfig(REPARTITION_WITHOUT_PARTITION_COUNT, TOPIC_CONFIG5); + final TopicsInfo topicsInfo = new TopicsInfo( + mkSet(REPARTITION_WITHOUT_PARTITION_COUNT), + mkSet(SOURCE_TOPIC_NAME1), + mkMap( + mkEntry(REPARTITION_WITHOUT_PARTITION_COUNT, repartitionTopicConfigWithoutPartitionCount) + ), + Collections.emptyMap() + ); + expect(internalTopologyBuilder.topicGroups()) + .andReturn(mkMap( + mkEntry(0, topicsInfo), + mkEntry(1, setupTopicInfoWithRepartitionTopicWithoutPartitionCount(repartitionTopicConfigWithoutPartitionCount)) + )); + expect(internalTopologyBuilder.copartitionGroups()).andReturn(Collections.emptyList()); + copartitionedTopicsEnforcer.enforce(eq(Collections.emptySet()), anyObject(), eq(clusterMetadata)); + expect(internalTopicManager.makeReady( + mkMap( + mkEntry(REPARTITION_WITHOUT_PARTITION_COUNT, repartitionTopicConfigWithoutPartitionCount) + )) + ).andReturn(Collections.emptySet()); + setupClusterWithMissingPartitionCounts(mkSet(SOURCE_TOPIC_NAME1)); + replay(internalTopicManager, internalTopologyBuilder, clusterMetadata); + final RepartitionTopics repartitionTopics = new RepartitionTopics( + internalTopologyBuilder, + internalTopicManager, + copartitionedTopicsEnforcer, + clusterMetadata, + "[test] " + ); + + final TaskAssignmentException exception = assertThrows(TaskAssignmentException.class, repartitionTopics::setup); + assertThat( + exception.getMessage(), + is("No partition count found for source topic " + SOURCE_TOPIC_NAME1 + ", but it should have been.") + ); + } + + @Test + public void shouldSetRepartitionTopicPartitionCountFromUpstreamExternalSourceTopic() { + final RepartitionTopicConfig repartitionTopicConfigWithoutPartitionCount = + new RepartitionTopicConfig(REPARTITION_WITHOUT_PARTITION_COUNT, TOPIC_CONFIG5); + final TopicsInfo topicsInfo = new TopicsInfo( + mkSet(REPARTITION_TOPIC_NAME1, REPARTITION_WITHOUT_PARTITION_COUNT), + mkSet(SOURCE_TOPIC_NAME1, REPARTITION_TOPIC_NAME2), + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1), + mkEntry(REPARTITION_TOPIC_NAME2, REPARTITION_TOPIC_CONFIG2), + mkEntry(REPARTITION_WITHOUT_PARTITION_COUNT, repartitionTopicConfigWithoutPartitionCount) + ), + Collections.emptyMap() + ); + expect(internalTopologyBuilder.topicGroups()) + .andReturn(mkMap( + mkEntry(0, topicsInfo), + mkEntry(1, setupTopicInfoWithRepartitionTopicWithoutPartitionCount(repartitionTopicConfigWithoutPartitionCount)) + )); + expect(internalTopologyBuilder.copartitionGroups()).andReturn(Collections.emptyList()); + copartitionedTopicsEnforcer.enforce(eq(Collections.emptySet()), anyObject(), eq(clusterMetadata)); + expect(internalTopicManager.makeReady( + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1), + mkEntry(REPARTITION_TOPIC_NAME2, REPARTITION_TOPIC_CONFIG2), + mkEntry(REPARTITION_WITHOUT_PARTITION_COUNT, repartitionTopicConfigWithoutPartitionCount) + )) + ).andReturn(Collections.emptySet()); + setupCluster(); + replay(internalTopicManager, internalTopologyBuilder, clusterMetadata); + final RepartitionTopics repartitionTopics = new RepartitionTopics( + internalTopologyBuilder, + internalTopicManager, + copartitionedTopicsEnforcer, + clusterMetadata, + "[test] " + ); + + repartitionTopics.setup(); + + verify(internalTopicManager, internalTopologyBuilder); + final Map topicPartitionsInfo = repartitionTopics.topicPartitionsInfo(); + assertThat(topicPartitionsInfo.size(), is(9)); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 0); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 1); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 2); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 3); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME2, 0); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME2, 1); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_WITHOUT_PARTITION_COUNT, 0); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_WITHOUT_PARTITION_COUNT, 1); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_WITHOUT_PARTITION_COUNT, 2); + } + + @Test + public void shouldSetRepartitionTopicPartitionCountFromUpstreamInternalRepartitionSourceTopic() { + final RepartitionTopicConfig repartitionTopicConfigWithoutPartitionCount = + new RepartitionTopicConfig(REPARTITION_WITHOUT_PARTITION_COUNT, TOPIC_CONFIG5); + final TopicsInfo topicsInfo = new TopicsInfo( + mkSet(REPARTITION_TOPIC_NAME2, REPARTITION_WITHOUT_PARTITION_COUNT), + mkSet(SOURCE_TOPIC_NAME1, REPARTITION_TOPIC_NAME1), + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1), + mkEntry(REPARTITION_TOPIC_NAME2, REPARTITION_TOPIC_CONFIG2), + mkEntry(REPARTITION_WITHOUT_PARTITION_COUNT, repartitionTopicConfigWithoutPartitionCount) + ), + Collections.emptyMap() + ); + expect(internalTopologyBuilder.topicGroups()) + .andReturn(mkMap( + mkEntry(0, topicsInfo), + mkEntry(1, setupTopicInfoWithRepartitionTopicWithoutPartitionCount(repartitionTopicConfigWithoutPartitionCount)) + )); + expect(internalTopologyBuilder.copartitionGroups()).andReturn(Collections.emptyList()); + copartitionedTopicsEnforcer.enforce(eq(Collections.emptySet()), anyObject(), eq(clusterMetadata)); + expect(internalTopicManager.makeReady( + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1), + mkEntry(REPARTITION_TOPIC_NAME2, REPARTITION_TOPIC_CONFIG2), + mkEntry(REPARTITION_WITHOUT_PARTITION_COUNT, repartitionTopicConfigWithoutPartitionCount) + )) + ).andReturn(Collections.emptySet()); + setupCluster(); + replay(internalTopicManager, internalTopologyBuilder, clusterMetadata); + final RepartitionTopics repartitionTopics = new RepartitionTopics( + internalTopologyBuilder, + internalTopicManager, + copartitionedTopicsEnforcer, + clusterMetadata, + "[test] " + ); + + repartitionTopics.setup(); + + verify(internalTopicManager, internalTopologyBuilder); + final Map topicPartitionsInfo = repartitionTopics.topicPartitionsInfo(); + assertThat(topicPartitionsInfo.size(), is(10)); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 0); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 1); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 2); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME1, 3); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME2, 0); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_TOPIC_NAME2, 1); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_WITHOUT_PARTITION_COUNT, 0); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_WITHOUT_PARTITION_COUNT, 1); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_WITHOUT_PARTITION_COUNT, 2); + verifyRepartitionTopicPartitionInfo(topicPartitionsInfo, REPARTITION_WITHOUT_PARTITION_COUNT, 3); + } + + @Test + public void shouldNotSetupRepartitionTopicsWhenTopologyDoesNotContainAnyRepartitionTopics() { + final TopicsInfo topicsInfo = new TopicsInfo( + mkSet(SINK_TOPIC_NAME1), + mkSet(SOURCE_TOPIC_NAME1), + Collections.emptyMap(), + Collections.emptyMap() + ); + expect(internalTopologyBuilder.topicGroups()) + .andReturn(mkMap(mkEntry(0, topicsInfo))); + expect(internalTopologyBuilder.copartitionGroups()).andReturn(Collections.emptySet()); + expect(internalTopicManager.makeReady(Collections.emptyMap())).andReturn(Collections.emptySet()); + setupCluster(); + replay(internalTopicManager, internalTopologyBuilder, clusterMetadata); + final RepartitionTopics repartitionTopics = new RepartitionTopics( + internalTopologyBuilder, + internalTopicManager, + copartitionedTopicsEnforcer, + clusterMetadata, + "[test] " + ); + + repartitionTopics.setup(); + + verify(internalTopicManager, internalTopologyBuilder); + final Map topicPartitionsInfo = repartitionTopics.topicPartitionsInfo(); + assertThat(topicPartitionsInfo, is(Collections.emptyMap())); + } + + private void verifyRepartitionTopicPartitionInfo(final Map topicPartitionsInfo, + final String topic, + final int partition) { + final TopicPartition repartitionTopicPartition = new TopicPartition(topic, partition); + assertThat(topicPartitionsInfo.containsKey(repartitionTopicPartition), is(true)); + final PartitionInfo repartitionTopicInfo = topicPartitionsInfo.get(repartitionTopicPartition); + assertThat(repartitionTopicInfo.topic(), is(topic)); + assertThat(repartitionTopicInfo.partition(), is(partition)); + assertThat(repartitionTopicInfo.inSyncReplicas(), is(new Node[0])); + assertThat(repartitionTopicInfo.leader(), nullValue()); + assertThat(repartitionTopicInfo.offlineReplicas(), is(new Node[0])); + assertThat(repartitionTopicInfo.replicas(), is(new Node[0])); + } + + private void setupCluster() { + setupClusterWithMissingTopicsAndMissingPartitionCounts(Collections.emptySet(), Collections.emptySet()); + } + + private void setupClusterWithMissingTopics(final Set missingTopics) { + setupClusterWithMissingTopicsAndMissingPartitionCounts(missingTopics, Collections.emptySet()); + } + + private void setupClusterWithMissingPartitionCounts(final Set topicsWithMissingPartitionCounts) { + setupClusterWithMissingTopicsAndMissingPartitionCounts(Collections.emptySet(), topicsWithMissingPartitionCounts); + } + + private void setupClusterWithMissingTopicsAndMissingPartitionCounts(final Set missingTopics, + final Set topicsWithMissingPartitionCounts) { + final Set topics = mkSet( + SOURCE_TOPIC_NAME1, + SOURCE_TOPIC_NAME2, + SOURCE_TOPIC_NAME3, + SINK_TOPIC_NAME1, + SINK_TOPIC_NAME2, + REPARTITION_TOPIC_NAME1, + REPARTITION_TOPIC_NAME2, + REPARTITION_TOPIC_NAME3, + REPARTITION_TOPIC_NAME4, + SOME_OTHER_TOPIC + ); + topics.removeAll(missingTopics); + expect(clusterMetadata.topics()).andStubReturn(topics); + expect(clusterMetadata.partitionCountForTopic(SOURCE_TOPIC_NAME1)) + .andStubReturn(topicsWithMissingPartitionCounts.contains(SOURCE_TOPIC_NAME1) ? null : 3); + expect(clusterMetadata.partitionCountForTopic(SOURCE_TOPIC_NAME2)) + .andStubReturn(topicsWithMissingPartitionCounts.contains(SOURCE_TOPIC_NAME2) ? null : 1); + expect(clusterMetadata.partitionCountForTopic(SOURCE_TOPIC_NAME3)) + .andStubReturn(topicsWithMissingPartitionCounts.contains(SOURCE_TOPIC_NAME3) ? null : 2); + } + + private TopicsInfo setupTopicInfoWithRepartitionTopicWithoutPartitionCount(final RepartitionTopicConfig repartitionTopicConfigWithoutPartitionCount) { + return new TopicsInfo( + mkSet(SINK_TOPIC_NAME2), + mkSet(REPARTITION_TOPIC_NAME1, REPARTITION_WITHOUT_PARTITION_COUNT), + mkMap( + mkEntry(REPARTITION_TOPIC_NAME1, REPARTITION_TOPIC_CONFIG1), + mkEntry(REPARTITION_WITHOUT_PARTITION_COUNT, repartitionTopicConfigWithoutPartitionCount) + ), + Collections.emptyMap() + ); + } +} \ No newline at end of file