content: /*
   * 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.coordinator.group.streams.topics;
  
  import org.apache.kafka.common.errors.StreamsInvalidTopologyException;
  import org.apache.kafka.common.utils.LogContext;
  
  import org.slf4j.Logger;
  
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.OptionalInt;
  import java.util.Set;
  import java.util.TreeMap;
  import java.util.function.Function;
  import java.util.stream.Collectors;
  
  /**
   * This class is responsible for enforcing the number of partitions in copartitioned topics. For each copartition group, it checks whether
   * the number of partitions for all repartition topics is the same, and enforces copartitioning for repartition topics whose number of
   * partitions is not enforced by the topology.
   */
  public class CopartitionedTopicsEnforcer {
  
      private final Logger log;
      private final Function<String, OptionalInt> topicPartitionCountProvider;
  
      /**
       * The constructor for the class.
       *
       * @param logContext                  The context for emitting log messages.
       * @param topicPartitionCountProvider Returns the number of partitions for a given topic, representing the current state of the broker
       *                                    as well as any partition number decisions that have already been made. In particular, we expect
       *                                    the number of partitions for all topics in all co-partitions groups to be defined.
       */
      public CopartitionedTopicsEnforcer(final LogContext logContext,
                                         final Function<String, OptionalInt> topicPartitionCountProvider) {
          this.log = logContext.logger(getClass());
          this.topicPartitionCountProvider = topicPartitionCountProvider;
      }
  
      /**
       * Enforces the number of partitions for copartitioned topics.
       *
       * @param copartitionedTopics          The set of copartitioned topics (external source topics and repartition topics).
       * @param fixedRepartitionTopics       The set of repartition topics whose partition count is fixed by the topology sent by the
       *                                     client (in particular, when the user uses `repartition` in the DSL).
       * @param flexibleRepartitionTopics    The set of repartition topics whose partition count is flexible, and can be changed.
       *
       * @throws IllegalStateException If the partition count for any topic in copartitionedTopics is not defined by
       *                               topicPartitionCountProvider.
       *
       * @return A map from all repartition topics in copartitionedTopics to their updated partition counts.
       */
      public Map<String, Integer> enforce(final Set<String> copartitionedTopics,
                                          final Set<String> fixedRepartitionTopics,
                                          final Set<String> flexibleRepartitionTopics) throws StreamsInvalidTopologyException {
          if (copartitionedTopics.isEmpty()) {
              log.debug("Ignoring unexpected empty copartitioned topics set.");
              return Map.of();
          }
          final Map<String, Integer> returnedPartitionCounts = new HashMap<>();
  
          final Map<String, Integer> repartitionTopicPartitionCounts =
              copartitionedTopics.stream()
                  .filter(x -> fixedRepartitionTopics.contains(x) || flexibleRepartitionTopics.contains(x))
                  .collect(Collectors.toMap(topic -> topic, this::getPartitionCount));
  
          final Map<String, Integer> nonRepartitionTopicPartitions =
              copartitionedTopics.stream().filter(topic -> !repartitionTopicPartitionCounts.containsKey(topic))
                  .collect(Collectors.toMap(topic -> topic, this::getPartitionCount));
  
          final int numPartitionsToUseForRepartitionTopics;
  
          if (copartitionedTopics.equals(repartitionTopicPartitionCounts.keySet())) {
  
              // if there's at least one repartition topic with fixed number of partitions
              // validate that they all have same number of partitions
              if (!fixedRepartitionTopics.isEmpty()) {
                  numPartitionsToUseForRepartitionTopics = validateAndGetNumOfPartitions(
                      repartitionTopicPartitionCounts,
                      fixedRepartitionTopics
                  );
              } 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(repartitionTopicPartitionCounts);
              }
          } 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 Entry<String, Integer> repartitionTopic : repartitionTopicPartitionCounts.entrySet()) {
              returnedPartitionCounts.put(repartitionTopic.getKey(), numPartitionsToUseForRepartitionTopics);
              if (fixedRepartitionTopics.contains(repartitionTopic.getKey())
                  && repartitionTopic.getValue() != numPartitionsToUseForRepartitionTopics) {
                  final String msg = String.format("Number of partitions [%d] of repartition topic [%s] " +
                          "doesn't match number of partitions [%d] of the source topic.",
                      repartitionTopic.getValue(),
                      repartitionTopic.getKey(),
                      numPartitionsToUseForRepartitionTopics);
                  throw TopicConfigurationException.incorrectlyPartitionedTopics(msg);
              }
          }
  
          return returnedPartitionCounts;
      }
  
      private int getPartitionCount(final String topicName) {
          OptionalInt partitions = topicPartitionCountProvider.apply(topicName);
          if (partitions.isPresent()) {
              return partitions.getAsInt();
          } else {
              throw new IllegalStateException("Number of partitions is not set for topic: " + topicName);
          }
      }
  
      private int validateAndGetNumOfPartitions(final Map<String, Integer> repartitionTopics,
                                                final Collection<String> fixedRepartitionTopics) {
          final String firstTopicName = fixedRepartitionTopics.iterator().next();
  
          final int firstNumberOfPartitionsOfInternalTopic = getPartitionCount(firstTopicName);
  
          for (final String topicName : fixedRepartitionTopics) {
              final int numberOfPartitions = getPartitionCount(topicName);
  
              if (numberOfPartitions != firstNumberOfPartitionsOfInternalTopic) {
                  final String msg = String.format("Following topics do not have the same number of partitions: [%s]",
                      new TreeMap<>(repartitionTopics));
                  throw TopicConfigurationException.incorrectlyPartitionedTopics(msg);
              }
          }
  
          return firstNumberOfPartitionsOfInternalTopic;
      }
  
      private int getSamePartitions(final Map<String, Integer> nonRepartitionTopicsInCopartitionGroup) {
          final int partitions = nonRepartitionTopicsInCopartitionGroup.values().iterator().next();
          for (final Entry<String, Integer> entry : nonRepartitionTopicsInCopartitionGroup.entrySet()) {
              if (entry.getValue() != partitions) {
                  final TreeMap<String, Integer> sorted = new TreeMap<>(nonRepartitionTopicsInCopartitionGroup);
                  throw TopicConfigurationException.incorrectlyPartitionedTopics(
                      String.format("Following topics do not have the same number of partitions: [%s]", sorted));
              }
          }
          return partitions;
      }
  
      private int getMaxPartitions(final Map<String, Integer> repartitionTopicsInCopartitionGroup) {
          int maxPartitions = 0;
  
          for (final Integer numPartitions : repartitionTopicsInCopartitionGroup.values()) {
              maxPartitions = Integer.max(maxPartitions, numPartitions);
          }
          if (maxPartitions == 0) {
              throw new StreamsInvalidTopologyException("All topics in the copartition group had undefined partition number: " +
                  repartitionTopicsInCopartitionGroup.keySet());
          }
          return maxPartitions;
      }
  
  }
  