org.apache.kafka.streams.processor.internals.assignment.CopartitionedTopicsEnforcer Maven / Gradle / Ivy
/*
* 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.assignment;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.streams.errors.TopologyException;
import org.apache.kafka.streams.processor.internals.InternalTopicConfig;
import org.slf4j.Logger;
import java.util.Collection;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Set;
import java.util.TreeMap;
import java.util.function.Supplier;
import java.util.stream.Collectors;
public class CopartitionedTopicsEnforcer {
private final String logPrefix;
private final Logger log;
public CopartitionedTopicsEnforcer(final String logPrefix) {
this.logPrefix = logPrefix;
final LogContext logContext = new LogContext(logPrefix);
log = logContext.logger(getClass());
}
public void enforce(final Set copartitionGroup,
final Map allRepartitionTopicsNumPartitions,
final Cluster metadata) {
if (copartitionGroup.isEmpty()) {
return;
}
final Map repartitionTopicConfigs =
copartitionGroup.stream()
.filter(allRepartitionTopicsNumPartitions::containsKey)
.collect(Collectors.toMap(topic -> topic, allRepartitionTopicsNumPartitions::get));
final Map nonRepartitionTopicPartitions =
copartitionGroup.stream().filter(topic -> !allRepartitionTopicsNumPartitions.containsKey(topic))
.collect(Collectors.toMap(topic -> topic, topic -> {
final Integer partitions = metadata.partitionCountForTopic(topic);
if (partitions == null) {
final String str = String.format("%sTopic not found: %s", logPrefix, topic);
log.error(str);
throw new IllegalStateException(str);
} else {
return partitions;
}
}));
final int numPartitionsToUseForRepartitionTopics;
final Collection internalTopicConfigs = repartitionTopicConfigs.values();
if (copartitionGroup.equals(repartitionTopicConfigs.keySet())) {
final Collection internalTopicConfigsWithEnforcedNumberOfPartitions = internalTopicConfigs
.stream()
.filter(InternalTopicConfig::hasEnforcedNumberOfPartitions)
.collect(Collectors.toList());
// if there's at least one repartition topic with enforced number of partitions
// validate that they all have same number of partitions
if (!internalTopicConfigsWithEnforcedNumberOfPartitions.isEmpty()) {
numPartitionsToUseForRepartitionTopics = validateAndGetNumOfPartitions(
repartitionTopicConfigs,
internalTopicConfigsWithEnforcedNumberOfPartitions
);
} else {
// If all topics for this co-partition group are repartition topics,
// then set the number of partitions to be the maximum of the number of partitions.
numPartitionsToUseForRepartitionTopics = getMaxPartitions(repartitionTopicConfigs);
}
} else {
// Otherwise, use the number of partitions from external topics (which must all be the same)
numPartitionsToUseForRepartitionTopics = getSamePartitions(nonRepartitionTopicPartitions);
}
// coerce all the repartition topics to use the decided number of partitions.
for (final InternalTopicConfig config : internalTopicConfigs) {
maybeSetNumberOfPartitionsForInternalTopic(numPartitionsToUseForRepartitionTopics, config);
final int numberOfPartitionsOfInternalTopic = config
.numberOfPartitions()
.orElseThrow(emptyNumberOfPartitionsExceptionSupplier(config.name()));
if (numberOfPartitionsOfInternalTopic != numPartitionsToUseForRepartitionTopics) {
final String msg = String.format("%sNumber of partitions [%d] of repartition topic [%s] " +
"doesn't match number of partitions [%d] of the source topic.",
logPrefix,
numberOfPartitionsOfInternalTopic,
config.name(),
numPartitionsToUseForRepartitionTopics);
throw new TopologyException(msg);
}
}
}
private static void maybeSetNumberOfPartitionsForInternalTopic(final int numPartitionsToUseForRepartitionTopics,
final InternalTopicConfig config) {
if (!config.hasEnforcedNumberOfPartitions()) {
config.setNumberOfPartitions(numPartitionsToUseForRepartitionTopics);
}
}
private int validateAndGetNumOfPartitions(final Map repartitionTopicConfigs,
final Collection internalTopicConfigs) {
final InternalTopicConfig firstInternalTopicConfig = internalTopicConfigs.iterator().next();
final int firstNumberOfPartitionsOfInternalTopic = firstInternalTopicConfig
.numberOfPartitions()
.orElseThrow(emptyNumberOfPartitionsExceptionSupplier(firstInternalTopicConfig.name()));
for (final InternalTopicConfig internalTopicConfig : internalTopicConfigs) {
final Integer numberOfPartitions = internalTopicConfig
.numberOfPartitions()
.orElseThrow(emptyNumberOfPartitionsExceptionSupplier(internalTopicConfig.name()));
if (numberOfPartitions != firstNumberOfPartitionsOfInternalTopic) {
final Map repartitionTopics = repartitionTopicConfigs
.entrySet()
.stream()
.collect(Collectors.toMap(Entry::getKey, entry -> entry.getValue().numberOfPartitions().get()));
final String msg = String.format("%sFollowing topics do not have the same number of partitions: [%s]",
logPrefix,
new TreeMap<>(repartitionTopics));
throw new TopologyException(msg);
}
}
return firstNumberOfPartitionsOfInternalTopic;
}
private static Supplier emptyNumberOfPartitionsExceptionSupplier(final String topic) {
return () -> new TopologyException("Number of partitions is not set for topic: " + topic);
}
private int getSamePartitions(final Map nonRepartitionTopicsInCopartitionGroup) {
final int partitions = nonRepartitionTopicsInCopartitionGroup.values().iterator().next();
for (final Entry entry : nonRepartitionTopicsInCopartitionGroup.entrySet()) {
if (entry.getValue() != partitions) {
final TreeMap sorted = new TreeMap<>(nonRepartitionTopicsInCopartitionGroup);
throw new TopologyException(
String.format("%sTopics not co-partitioned: [%s]",
logPrefix, sorted)
);
}
}
return partitions;
}
private int getMaxPartitions(final Map repartitionTopicsInCopartitionGroup) {
int maxPartitions = 0;
for (final InternalTopicConfig config : repartitionTopicsInCopartitionGroup.values()) {
final Optional partitions = config.numberOfPartitions();
maxPartitions = Integer.max(maxPartitions, partitions.orElse(maxPartitions));
}
if (maxPartitions <= 0) {
throw new IllegalStateException(logPrefix + "Could not validate the copartitioning of topics: " + repartitionTopicsInCopartitionGroup.keySet());
}
return maxPartitions;
}
}