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.streams.processor.internals.assignment;
  
  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.TopicPartitionInfo;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.streams.KeyValue;
  import org.apache.kafka.streams.StreamsConfig;
  import org.apache.kafka.streams.processor.TaskId;
  import org.apache.kafka.streams.processor.assignment.AssignmentConfigs;
  import org.apache.kafka.streams.processor.assignment.ProcessId;
  import org.apache.kafka.streams.processor.internals.InternalTopicManager;
  import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  import java.util.SortedMap;
  import java.util.SortedSet;
  import java.util.TreeSet;
  import java.util.function.BiFunction;
  import java.util.function.BiPredicate;
  import java.util.stream.Collectors;
  import java.util.stream.Stream;
  
  public class RackAwareTaskAssignor {
  
      @FunctionalInterface
      public interface MoveStandbyTaskPredicate {
          boolean canMove(final ClientState source,
                          final ClientState destination,
                          final TaskId taskId,
                          final Map<ProcessId, ClientState> clientStateMap);
      }
  
      @FunctionalInterface
      public interface CostFunction {
          int getCost(final TaskId taskId,
                      final ProcessId processId,
                      final boolean inCurrentAssignment,
                      final int trafficCost,
                      final int nonOverlapCost,
                      final boolean isStandby);
      }
  
      // For stateless tasks, it's ok to move them around. So we have 0 non_overlap_cost
      public static final int STATELESS_TRAFFIC_COST = 1;
      public static final int STATELESS_NON_OVERLAP_COST = 0;
  
      private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
  
      // This is number is picked based on testing. Usually the optimization for standby assignment
      // stops after 3 rounds
      public static final int STANDBY_OPTIMIZER_MAX_ITERATION = 4;
  
      private final Cluster fullMetadata;
      private final Map<TaskId, Set<TopicPartition>> partitionsForTask;
      private final Map<TaskId, Set<TopicPartition>> changelogPartitionsForTask;
      private final Map<Subtopology, Set<TaskId>> tasksForTopicGroup;
      private final AssignmentConfigs assignmentConfigs;
      private final Map<TopicPartition, Set<String>> racksForPartition;
      private final Map<ProcessId, String> rackForProcess;
      private final InternalTopicManager internalTopicManager;
      private final boolean validClientRack;
      private final Time time;
      private Boolean canEnable = null;
  
      public RackAwareTaskAssignor(final Cluster fullMetadata,
                                   final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                   final Map<TaskId, Set<TopicPartition>> changelogPartitionsForTask,
                                   final Map<Subtopology, Set<TaskId>> tasksForTopicGroup,
                                   final Map<ProcessId, Map<String, Optional<String>>> racksForProcessConsumer,
                                   final InternalTopicManager internalTopicManager,
                                   final AssignmentConfigs assignmentConfigs,
                                   final Time time) {
          this.fullMetadata = fullMetadata;
          this.partitionsForTask = partitionsForTask;
          this.changelogPartitionsForTask = changelogPartitionsForTask;
          this.tasksForTopicGroup = tasksForTopicGroup;
          this.internalTopicManager = internalTopicManager;
          this.assignmentConfigs = assignmentConfigs;
          this.racksForPartition = new HashMap<>();
          this.rackForProcess = new HashMap<>();
          this.time = Objects.requireNonNull(time, "Time was not specified");
          validClientRack = validateClientRack(racksForProcessConsumer, assignmentConfigs,
              rackForProcess
          );
      }
  
      public boolean validClientRack() {
          return validClientRack;
      }
  
      public synchronized boolean canEnableRackAwareAssignor() {
          if (StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE.equals(assignmentConfigs.rackAwareAssignmentStrategy())) {
              return false;
          }
          if (canEnable != null) {
              return canEnable;
          }
          canEnable = validClientRack && validateTopicPartitionRack(false);
          if (assignmentConfigs.numStandbyReplicas() == 0 || !canEnable) {
              return canEnable;
          }
  
          canEnable = validateTopicPartitionRack(true);
          return canEnable;
      }
  
      // Visible for testing. This method also checks if all TopicPartitions exist in cluster
      boolean populateTopicsToDescribe(final Set<String> topicsToDescribe, final boolean changelog) {
          if (changelog) {
              // Changelog topics are not in metadata, we need to describe them
              changelogPartitionsForTask.values().stream().flatMap(Collection::stream).forEach(tp -> topicsToDescribe.add(tp.topic()));
              return true;
          }
  
          // Make sure rackId exist for all TopicPartitions needed
          for (final Set<TopicPartition> topicPartitions : partitionsForTask.values()) {
              for (final TopicPartition topicPartition : topicPartitions) {
                  final PartitionInfo partitionInfo = fullMetadata.partition(topicPartition);
                  if (partitionInfo == null) {
                      log.error("TopicPartition {} doesn't exist in cluster", topicPartition);
                      return false;
                  }
                  final Node[] replica = partitionInfo.replicas();
                  if (replica == null || replica.length == 0) {
                      topicsToDescribe.add(topicPartition.topic());
                      continue;
                  }
                  for (final Node node : replica) {
                      if (node.hasRack()) {
                          racksForPartition.computeIfAbsent(topicPartition, k -> new HashSet<>()).add(node.rack());
                      } else {
                          log.warn("Node {} for topic partition {} doesn't have rack", node, topicPartition);
                          return false;
                      }
                  }
              }
          }
          return true;
      }
  
      /**
       * This function populates the {@param racksForPartition} parameter passed into the function by using both
       * the {@code Cluster} metadata as well as the {@param internalTopicManager} for topics that have stale
       * information.
       *
       * @return whether the operation successfully completed and the rack information is valid.
       */
      private boolean validateTopicPartitionRack(final boolean changelogTopics) {
          // Make sure rackId exist for all TopicPartitions needed
          final Set<String> topicsToDescribe = new HashSet<>();
          if (!populateTopicsToDescribe(topicsToDescribe, changelogTopics)) {
              return false;
          }
  
          if (!topicsToDescribe.isEmpty()) {
              log.info("Fetching PartitionInfo for topics {}", topicsToDescribe);
              try {
                  final Map<String, List<TopicPartitionInfo>> topicPartitionInfo = internalTopicManager.getTopicPartitionInfo(topicsToDescribe);
                  if (topicsToDescribe.size() > topicPartitionInfo.size()) {
                      topicsToDescribe.removeAll(topicPartitionInfo.keySet());
                      log.error("Failed to describe topic for {}", topicsToDescribe);
                      return false;
                  }
                  for (final Map.Entry<String, List<TopicPartitionInfo>> entry : topicPartitionInfo.entrySet()) {
                      final List<TopicPartitionInfo> partitionInfos = entry.getValue();
                      for (final TopicPartitionInfo partitionInfo : partitionInfos) {
                          final int partition = partitionInfo.partition();
                          final List<Node> replicas = partitionInfo.replicas();
                          if (replicas == null || replicas.isEmpty()) {
                              log.error("No replicas found for topic partition {}: {}", entry.getKey(), partition);
                              return false;
                          }
                          final TopicPartition topicPartition = new TopicPartition(entry.getKey(), partition);
                          for (final Node node : replicas) {
                              if (node.hasRack()) {
                                  racksForPartition.computeIfAbsent(topicPartition, k -> new HashSet<>()).add(node.rack());
                              } else {
                                  return false;
                              }
                          }
                      }
                  }
              } catch (final Exception e) {
                  log.error("Failed to describe topics {}", topicsToDescribe, e);
                  return false;
              }
          }
          return true;
      }
  
      /**
       * Verifies that within the {@param racksForProcessConsumer}, a single process is only running on a single
       * rack. This function mutates the parameter {@param racksForProcess} to contain the resulting process
       * to rack mapping.
       *
       * @return true if the validation was successful, and false otherwise.
       */
      public static boolean validateClientRack(final Map<ProcessId, Map<String, Optional<String>>> racksForProcessConsumer,
                                               final AssignmentConfigs assignmentConfigs,
                                               final Map<ProcessId, String> rackForProcess) {
          if (racksForProcessConsumer == null) {
              return false;
          }
          /*
           * Check rack information is populated correctly in clients
           * 1. RackId exist for all clients
           * 2. Different consumerId for same process should have same rackId
           */
          for (final Map.Entry<ProcessId, Map<String, Optional<String>>> entry : racksForProcessConsumer.entrySet()) {
              final ProcessId processId = entry.getKey();
              KeyValue<String, String> previousRackInfo = null;
              for (final Map.Entry<String, Optional<String>> rackEntry : entry.getValue().entrySet()) {
                  if (rackEntry.getValue().isEmpty()) {
                      if (!StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE.equals(assignmentConfigs.rackAwareAssignmentStrategy())) {
                          log.error(
                              String.format("RackId doesn't exist for process %s and consumer %s",
                                  processId, rackEntry.getKey()));
                      }
                      return false;
                  }
                  if (previousRackInfo == null) {
                      previousRackInfo = KeyValue.pair(rackEntry.getKey(), rackEntry.getValue().get());
                  } else if (!previousRackInfo.value.equals(rackEntry.getValue().get())) {
                      log.error(
                          String.format("Consumers %s and %s for same process %s has different rackId %s and %s. File a ticket for this bug",
                              previousRackInfo.key,
                              rackEntry.getKey(),
                              entry.getKey(),
                              previousRackInfo.value,
                              rackEntry.getValue().get()
                          )
                      );
                      return false;
                  }
              }
              if (previousRackInfo == null) {
                  if (!StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE.equals(assignmentConfigs.rackAwareAssignmentStrategy())) {
                      log.error(String.format("RackId doesn't exist for process %s", processId));
                  }
                  return false;
              }
              rackForProcess.put(entry.getKey(), previousRackInfo.value);
          }
  
          return true;
      }
  
      public Map<ProcessId, String> racksForProcess() {
          return Collections.unmodifiableMap(rackForProcess);
      }
  
      public Map<TopicPartition, Set<String>> racksForPartition() {
          return Collections.unmodifiableMap(racksForPartition);
      }
  
      private int getCost(final TaskId taskId, final ProcessId processId, final boolean inCurrentAssignment, final int trafficCost, final int nonOverlapCost, final boolean isStandby) {
          final String clientRack = rackForProcess.get(processId);
          if (clientRack == null) {
              throw new IllegalStateException("Client " + processId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
          }
  
          final Set<TopicPartition> topicPartitions = isStandby ? changelogPartitionsForTask.get(taskId) : partitionsForTask.get(taskId);
          if (topicPartitions == null || topicPartitions.isEmpty()) {
              throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
          }
  
          int cost = 0;
          for (final TopicPartition tp : topicPartitions) {
              final Set<String> tpRacks = racksForPartition.get(tp);
              if (tpRacks == null || tpRacks.isEmpty()) {
                  throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
              }
              if (!tpRacks.contains(clientRack)) {
                  cost += trafficCost;
              }
          }
  
          if (!inCurrentAssignment) {
              cost += nonOverlapCost;
          }
  
          return cost;
      }
  
      /**
       * Compute the cost for the provided {@code activeTasks}. The passed in active tasks must be contained in {@code clientState}.
       */
      long activeTasksCost(final SortedSet<TaskId> activeTasks,
                           final SortedMap<ProcessId, ClientState> clientStates,
                           final int trafficCost,
                           final int nonOverlapCost) {
          return tasksCost(activeTasks, clientStates, trafficCost, nonOverlapCost, ClientState::hasActiveTask, false, false);
      }
  
      /**
       * Compute the cost for the provided {@code standbyTasks}. The passed in standby tasks must be contained in {@code clientState}.
       */
      long standByTasksCost(final SortedSet<TaskId> standbyTasks,
                            final SortedMap<ProcessId, ClientState> clientStates,
                            final int trafficCost,
                            final int nonOverlapCost) {
          return tasksCost(standbyTasks, clientStates, trafficCost, nonOverlapCost, ClientState::hasStandbyTask, true, true);
      }
  
      private long tasksCost(final SortedSet<TaskId> tasks,
                             final SortedMap<ProcessId, ClientState> clientStates,
                             final int trafficCost,
                             final int nonOverlapCost,
                             final BiPredicate<ClientState, TaskId> hasAssignedTask,
                             final boolean hasReplica,
                             final boolean isStandby) {
          if (tasks.isEmpty()) {
              return 0;
          }
          final List<ProcessId> clientList = new ArrayList<>(clientStates.keySet());
          final List<TaskId> taskIdList = new ArrayList<>(tasks);
          final Graph<Integer> graph = new MinTrafficGraphConstructor<ClientState>()
              .constructTaskGraph(
                  clientList,
                  taskIdList,
                  clientStates,
                  new HashMap<>(),
                  new HashMap<>(),
                  hasAssignedTask,
                  this::getCost,
                  trafficCost,
                  nonOverlapCost,
                  hasReplica,
                  isStandby
              );
          return graph.totalCost();
      }
  
      /**
       * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first.
       * {@code trafficCost} and {@code nonOverlapCost} balance cross rack traffic optimization and task movement.
       * If we set {@code trafficCost} to a larger number, we are more likely to compute an assignment with less
       * cross rack traffic. However, tasks may be shuffled a lot across clients. If we set {@code nonOverlapCost}
       * to a larger number, we are more likely to compute an assignment with similar to input assignment. However,
       * cross rack traffic can be higher. In extreme case, if we set {@code nonOverlapCost} to 0 and @{code trafficCost}
       * to a positive value, the computed assignment will be minimum for cross rack traffic. If we set {@code trafficCost} to 0,
       * and {@code nonOverlapCost} to a positive value, the computed assignment should be the same as input
       * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
       * @param clientStates Client states
       * @param trafficCost Cost of cross rack traffic for each TopicPartition
       * @param nonOverlapCost Cost of assign a task to a different client
       * @return Total cost after optimization
       */
      public long optimizeActiveTasks(final SortedSet<TaskId> activeTasks,
                                      final SortedMap<ProcessId, ClientState> clientStates,
                                      final int trafficCost,
                                      final int nonOverlapCost) {
          if (activeTasks.isEmpty()) {
              return 0;
          }
  
          log.info("Assignment before active task optimization is {}\n with cost {}", clientStates,
              activeTasksCost(activeTasks, clientStates, trafficCost, nonOverlapCost));
  
          final long startTime = time.milliseconds();
          final List<ProcessId> clientList = new ArrayList<>(clientStates.keySet());
          final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
          final Map<TaskId, ProcessId> taskClientMap = new HashMap<>();
          final Map<ProcessId, Integer> originalAssignedTaskNumber = new HashMap<>();
          final RackAwareGraphConstructor<ClientState> graphConstructor = RackAwareGraphConstructorFactory.create(assignmentConfigs, tasksForTopicGroup);
          final Graph<Integer> graph = graphConstructor.constructTaskGraph(
              clientList,
              taskIdList,
              clientStates,
              taskClientMap,
              originalAssignedTaskNumber,
              ClientState::hasActiveTask,
              this::getCost,
              trafficCost,
              nonOverlapCost,
              false,
              false
          );
  
          graph.solveMinCostFlow();
          final long cost = graph.totalCost();
  
          graphConstructor.assignTaskFromMinCostFlow(graph, clientList, taskIdList, clientStates, originalAssignedTaskNumber,
              taskClientMap, ClientState::assignActive, ClientState::unassignActive, ClientState::hasActiveTask);
  
          final long duration = time.milliseconds() - startTime;
          log.info("Assignment after {} milliseconds for active task optimization is {}\n with cost {}", duration, clientStates, cost);
          return cost;
      }
  
      public long optimizeStandbyTasks(final SortedMap<ProcessId, ClientState> clientStates,
                                       final int trafficCost,
                                       final int nonOverlapCost,
                                       final MoveStandbyTaskPredicate moveStandbyTask) {
          final BiFunction<ClientState, ClientState, List<TaskId>> getMovableTasks = (source, destination) -> source.standbyTasks().stream()
              .filter(task -> !destination.hasAssignedTask(task))
              .filter(task -> moveStandbyTask.canMove(source, destination, task, clientStates))
              .sorted()
              .collect(Collectors.toList());
  
          final long startTime = time.milliseconds();
          final List<ProcessId> clientList = new ArrayList<>(clientStates.keySet());
          final SortedSet<TaskId> standbyTasks = new TreeSet<>();
          clientStates.values().forEach(clientState -> standbyTasks.addAll(clientState.standbyTasks()));
  
          log.info("Assignment before standby task optimization is {}\n with cost {}", clientStates,
              standByTasksCost(standbyTasks, clientStates, trafficCost, nonOverlapCost));
  
          boolean taskMoved = true;
          int round = 0;
          final RackAwareGraphConstructor<ClientState> graphConstructor = new MinTrafficGraphConstructor<>();
          while (taskMoved && round < STANDBY_OPTIMIZER_MAX_ITERATION) {
              taskMoved = false;
              round++;
              for (int i = 0; i < clientList.size(); i++) {
                  final ClientState clientState1 = clientStates.get(clientList.get(i));
                  for (int j = i + 1; j < clientList.size(); j++) {
                      final ClientState clientState2 = clientStates.get(clientList.get(j));
  
                      final String rack1 = rackForProcess.get(clientState1.processId());
                      final String rack2 = rackForProcess.get(clientState2.processId());
                      // Cross rack traffic can not be reduced if racks are the same
                      if (rack1.equals(rack2)) {
                          continue;
                      }
  
                      final List<TaskId> movable1 = getMovableTasks.apply(clientState1, clientState2);
                      final List<TaskId> movable2 = getMovableTasks.apply(clientState2, clientState1);
  
                      // There's no needed to optimize if one is empty because the optimization
                      // can only swap tasks to keep the client's load balanced
                      if (movable1.isEmpty() || movable2.isEmpty()) {
                          continue;
                      }
  
                      final List<TaskId> taskIdList = Stream.concat(movable1.stream(),
                              movable2.stream())
                          .sorted()
                          .collect(Collectors.toList());
  
                      final List<ProcessId> clients = Stream.of(clientList.get(i), clientList.get(j))
                          .sorted().collect(
                              Collectors.toList());
  
                      final Map<TaskId, ProcessId> taskClientMap = new HashMap<>();
                      final Map<ProcessId, Integer> originalAssignedTaskNumber = new HashMap<>();
                      final Graph<Integer> graph = graphConstructor.constructTaskGraph(
                          clients,
                          taskIdList,
                          clientStates,
                          taskClientMap,
                          originalAssignedTaskNumber,
                          ClientState::hasStandbyTask,
                          this::getCost,
                          trafficCost,
                          nonOverlapCost,
                          true,
                          true
                      );
                      graph.solveMinCostFlow();
  
                      taskMoved |= graphConstructor.assignTaskFromMinCostFlow(graph, clients, taskIdList, clientStates,
                          originalAssignedTaskNumber,
                          taskClientMap, ClientState::assignStandby, ClientState::unassignStandby,
                          ClientState::hasStandbyTask);
                  }
              }
          }
          final long cost = standByTasksCost(standbyTasks, clientStates, trafficCost, nonOverlapCost);
  
          final long duration = time.milliseconds() - startTime;
          log.info("Assignment after {} rounds and {} milliseconds for standby task optimization is {}\n with cost {}", round, duration, clientStates, cost);
          return cost;
      }
  }
  