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.TopicPartition;
  import org.apache.kafka.streams.processor.TaskId;
  import org.apache.kafka.streams.processor.assignment.ProcessId;
  import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
  
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.ValueSource;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  import java.util.SortedMap;
  import java.util.SortedSet;
  import java.util.TreeMap;
  
  import static java.util.Collections.emptySet;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertBalancedTasks;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertValidAssignment;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomClientState;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMap;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTasksForTopicGroup;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  
  public class RackAwareGraphConstructorTest {
      private static final String MIN_COST = "min_cost";
      private static final String BALANCE_SUBTOPOLOGY = "balance_sub_topology";
  
      private static final int TP_SIZE = 40;
      private static final int PARTITION_SIZE = 3;
      private static final int TOPIC_GROUP_SIZE = TP_SIZE;
      private static final int CLIENT_SIZE = 20;
  
      private Graph<Integer> graph;
      private final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(
          TP_SIZE, PARTITION_SIZE, false);
      private final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
      private final List<TaskId> taskIdList = new ArrayList<>(taskIds);
      private final SortedMap<ProcessId, ClientState> clientStateMap = getRandomClientState(CLIENT_SIZE,
          TP_SIZE, PARTITION_SIZE, 1, false, taskIds);
      private final List<ProcessId> clientList = new ArrayList<>(clientStateMap.keySet());
      private final Map<TaskId, ProcessId> taskClientMap = new HashMap<>();
      private final Map<ProcessId, Integer> originalAssignedTaskNumber = new HashMap<>();
      private final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = getTasksForTopicGroup(TP_SIZE,
          PARTITION_SIZE);
      private RackAwareGraphConstructor<ClientState> constructor;
  
  
      public void setUp(final String constructorType) {
          randomAssignTasksToClient(taskIdList, clientStateMap);
  
          if (constructorType.equals(MIN_COST)) {
              constructor = new MinTrafficGraphConstructor<>();
          } else if (constructorType.equals(BALANCE_SUBTOPOLOGY)) {
              final List<Set<TaskId>> taskSetsPerTopicGroup = new ArrayList<>(new TreeMap<>(tasksForTopicGroup).values());
              constructor = new BalanceSubtopologyGraphConstructor<>(taskSetsPerTopicGroup);
          }
          graph = constructor.constructTaskGraph(
              clientList, taskIdList, clientStateMap, taskClientMap, originalAssignedTaskNumber, ClientState::hasAssignedTask, this::getCost, 10, 1, false, false);
      }
  
      private int getCost(final TaskId taskId, final ProcessId processId, final boolean inCurrentAssignment, final int trafficCost, final int nonOverlapCost, final boolean isStandby) {
          return 1;
      }
  
      @ParameterizedTest
      @ValueSource(strings = {MIN_COST, BALANCE_SUBTOPOLOGY})
      public void testSubtopologyShouldContainAllTasks(final String constructorType) {
          setUp(constructorType);
          if (constructorType.equals(MIN_COST)) {
              return;
          }
          taskIdList.add(new TaskId(41, 0)); // Extra task not in subtopology map
          assertThrows(IllegalStateException.class, () -> graph = constructor.constructTaskGraph(
              clientList, taskIdList, clientStateMap, taskClientMap, originalAssignedTaskNumber,
              ClientState::hasAssignedTask, this::getCost, 10, 1, false, false));
      }
  
      @ParameterizedTest
      @ValueSource(strings = {MIN_COST, BALANCE_SUBTOPOLOGY})
      public void testMinCostGraphConstructor(final String constructorType) {
          setUp(constructorType);
          if (constructorType.equals(BALANCE_SUBTOPOLOGY)) {
              return;
          }
          // Flow should equal to task size which means each task is assigned now
          assertEquals(taskIdList.size(), graph.flow());
  
          // Total node size should match. Extra 2 is for source and sink nodes
          assertEquals(taskIdList.size() + clientList.size() + 2, graph.nodes().size());
  
          // Check edges from source to task nodes
          SortedMap<Integer, Graph<Integer>.Edge> edges = graph.edges(RackAwareGraphConstructor.SOURCE_ID);
          for (final Graph<Integer>.Edge edge : edges.values()) {
              assertEquals(1, edge.flow);
              assertEquals(1, edge.capacity);
              assertEquals(0, edge.residualFlow);
              assertEquals(0, edge.cost);
              Assertions.assertTrue(edge.forwardEdge);
          }
  
          // Check edges from task nodes to clients nodes,
          for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
              edges = graph.edges(taskNodeId);
              assertEquals(clientList.size(), edges.size());
              int assignedClient = 0;
              for (final Graph<Integer>.Edge edge : edges.values()) {
                  final int flow = edge.flow;
                  if (flow == 1) {
                      assignedClient++;
                  }
                  assertEquals(1, edge.capacity);
                  assertEquals(flow == 1 ? 0 : 1, edge.residualFlow);
                  assertEquals(1, edge.cost);
                  Assertions.assertTrue(edge.forwardEdge);
              }
              // Task should be assigned to exact 1 client
              assertEquals(1, assignedClient);
              taskNodeId++;
          }
  
          // Check edges from second stage client node to sink
          final int sinkId = clientList.size() + taskIdList.size();
          int totalFlow = 0;
          for (int i = 0; i < clientList.size(); i++) {
              final ProcessId clientId = clientList.get(i);
              final int originalAssignedCount = originalAssignedTaskNumber.get(clientId);
  
              final int clientNodeId = i + taskIdList.size();
              edges = graph.edges(clientNodeId);
              assertEquals(1, edges.size());
              for (final Entry<Integer, Graph<Integer>.Edge> nodeEdge : edges.entrySet()) {
                  final Integer nodeId = nodeEdge.getKey();
                  assertEquals(sinkId, nodeId);
                  totalFlow += nodeEdge.getValue().flow;
                  assertEquals(originalAssignedCount, nodeEdge.getValue().capacity);
                  Assertions.assertTrue(nodeEdge.getValue().forwardEdge);
              }
          }
          assertEquals(taskIdList.size(), totalFlow);
      }
  
      @ParameterizedTest
      @ValueSource(strings = {MIN_COST, BALANCE_SUBTOPOLOGY})
      public void testBalanceSubtopologyGraphConstructor(final String constructorType) {
          setUp(constructorType);
          if (constructorType.equals(MIN_COST)) {
              return;
          }
          // Flow should equal to task size which means each task is assigned now
          assertEquals(taskIdList.size(), graph.flow());
  
          // Total node size should match. Extra 2 is for source and sink nodes
          assertEquals(taskIdList.size() + TOPIC_GROUP_SIZE * clientList.size() + clientList.size() + 2, graph.nodes().size());
  
          // Check edges from source to task nodes
          SortedMap<Integer, Graph<Integer>.Edge> edges = graph.edges(RackAwareGraphConstructor.SOURCE_ID);
          for (final Graph<Integer>.Edge edge : edges.values()) {
              assertEquals(1, edge.flow);
              assertEquals(1, edge.capacity);
              assertEquals(0, edge.residualFlow);
              assertEquals(0, edge.cost);
              Assertions.assertTrue(edge.forwardEdge);
          }
  
          // Check edges from task nodes to first stage clients node,
          int taskNodeId = 0;
          for (final Set<TaskId> unsortedTaskIds : tasksForTopicGroup.values()) {
              for (int i = 0; i < unsortedTaskIds.size(); i++) {
                  edges = graph.edges(taskNodeId);
                  assertEquals(clientList.size(), edges.size());
                  int assignedClient = 0;
                  for (final Graph<Integer>.Edge edge : edges.values()) {
                      final int flow = edge.flow;
                      if (flow == 1) {
                          assignedClient++;
                      }
                      assertEquals(1, edge.capacity);
                      assertEquals(flow == 1 ? 0 : 1, edge.residualFlow);
                      assertEquals(1, edge.cost);
                      Assertions.assertTrue(edge.forwardEdge);
                  }
                  // Task should be assigned to exact 1 client
                  assertEquals(1, assignedClient);
                  taskNodeId++;
              }
          }
  
          // Check edges from first stage client node to second stage client node
          int topicGroupIndex = 0;
          for (final Set<TaskId> tasks : tasksForTopicGroup.values()) {
              final int taskCount = tasks.size();
              for (int j = 0; j < clientList.size(); j++) {
                  final ProcessId clientId = clientList.get(j);
                  final int originalAssignedCount = originalAssignedTaskNumber.get(clientId);
                  final int expectedCapacity = (int) Math.ceil(originalAssignedCount * 1.0 / taskIdList.size() * taskCount);
                  final int clientNodeId = topicGroupIndex * clientList.size() + taskIdList.size() + j;
                  edges = graph.edges(clientNodeId);
                  assertEquals(1, edges.size());
                  for (final Entry<Integer, Graph<Integer>.Edge> nodeEdge : edges.entrySet()) {
                      final Integer nodeId = nodeEdge.getKey();
                      assertEquals(clientList.size() * tasksForTopicGroup.size() + taskIdList.size() + j, nodeId);
                      final Graph<Integer>.Edge edge = nodeEdge.getValue();
                      assertEquals(expectedCapacity, edge.capacity);
                      assertEquals(0, edge.cost);
                      Assertions.assertTrue(edge.forwardEdge);
                  }
              }
              topicGroupIndex++;
          }
  
          // Check edges from second stage client node to sink
          final int sinkId = clientList.size() + tasksForTopicGroup.size() * clientList.size() + taskIdList.size();
          int totalFlow = 0;
          for (int i = 0; i < clientList.size(); i++) {
              final ProcessId clientId = clientList.get(i);
              final int originalAssignedCount = originalAssignedTaskNumber.get(clientId);
  
              final int clientNodeId =
                  i + tasksForTopicGroup.size() * clientList.size() + taskIdList.size();
              edges = graph.edges(clientNodeId);
              assertEquals(1, edges.size());
              for (final Entry<Integer, Graph<Integer>.Edge> nodeEdge : edges.entrySet()) {
                  final Integer nodeId = nodeEdge.getKey();
                  assertEquals(sinkId, nodeId);
                  totalFlow += nodeEdge.getValue().flow;
                  assertEquals(originalAssignedCount, nodeEdge.getValue().capacity);
                  Assertions.assertTrue(nodeEdge.getValue().forwardEdge);
              }
          }
          assertEquals(taskIdList.size(), totalFlow);
      }
  
      @ParameterizedTest
      @ValueSource(strings = {MIN_COST, BALANCE_SUBTOPOLOGY})
      public void testAssignTaskFromMinCostFlow(final String constructorType) {
          setUp(constructorType);
          graph.solveMinCostFlow();
          constructor.assignTaskFromMinCostFlow(
              graph,
              clientList,
              taskIdList,
              clientStateMap,
              originalAssignedTaskNumber,
              taskClientMap,
              ClientState::assignActive,
              ClientState::unassignActive,
              ClientState::hasAssignedTask
          );
          assertValidAssignment(0, taskIds, emptySet(), clientStateMap, new StringBuilder());
          if (constructorType.equals(BALANCE_SUBTOPOLOGY)) {
              assertBalancedTasks(clientStateMap);
          }
      }
  
      private void randomAssignTasksToClient(final List<TaskId> taskIdList, final SortedMap<ProcessId, ClientState> clientStateMap) {
          int totalAssigned = 0;
          for (final ClientState clientState : clientStateMap.values()) {
              clientState.assignActive(taskIdList.get(totalAssigned++));
              clientState.assignActive(taskIdList.get(totalAssigned++));
          }
          while (totalAssigned < taskIdList.size()) {
              for (final ClientState clientState : clientStateMap.values()) {
                  if (AssignmentTestUtils.getRandom().nextInt(3) == 0) {
                      clientState.assignActive(taskIdList.get(totalAssigned));
                      totalAssigned++;
                      if (totalAssigned >= taskIdList.size()) {
                          break;
                      }
                  }
              }
          }
      }
  }
  