KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks

checkstyle/suppressions.xml

@@ -201,7 +201,7 @@
               files="Murmur3.java"/>
 
     <suppress checks="(NPathComplexity|CyclomaticComplexity)"
-              files="KStreamSlidingWindowAggregate.java"/>
+              files="(KStreamSlidingWindowAggregate|RackAwareTaskAssignor).java"/>
 
     <!-- suppress FinalLocalVariable outside of the streams package. -->
     <suppress checks="FinalLocalVariable"

streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java

@@ -268,24 +268,44 @@ public static class AssignmentConfigs {
         public final long probingRebalanceIntervalMs;
         public final List<String> rackAwareAssignmentTags;
 
+        // TODO: get from streamsConfig after we add the config
+        public final Integer trafficCost;
+        public final Integer nonOverlapCost;
+
         private AssignmentConfigs(final StreamsConfig configs) {
             acceptableRecoveryLag = configs.getLong(StreamsConfig.ACCEPTABLE_RECOVERY_LAG_CONFIG);
             maxWarmupReplicas = configs.getInt(StreamsConfig.MAX_WARMUP_REPLICAS_CONFIG);
             numStandbyReplicas = configs.getInt(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG);
             probingRebalanceIntervalMs = configs.getLong(StreamsConfig.PROBING_REBALANCE_INTERVAL_MS_CONFIG);
             rackAwareAssignmentTags = configs.getList(StreamsConfig.RACK_AWARE_ASSIGNMENT_TAGS_CONFIG);
+            // TODO: get from streamsConfig after we add the config
+            trafficCost = null;
+            nonOverlapCost = null;
         }
 
         AssignmentConfigs(final Long acceptableRecoveryLag,
                           final Integer maxWarmupReplicas,
                           final Integer numStandbyReplicas,
                           final Long probingRebalanceIntervalMs,
                           final List<String> rackAwareAssignmentTags) {
+            this(acceptableRecoveryLag, maxWarmupReplicas, numStandbyReplicas, probingRebalanceIntervalMs, rackAwareAssignmentTags, null, null);
+        }
+
+        AssignmentConfigs(final Long acceptableRecoveryLag,
+                          final Integer maxWarmupReplicas,
+                          final Integer numStandbyReplicas,
+                          final Long probingRebalanceIntervalMs,
+                          final List<String> rackAwareAssignmentTags,
+                          final Integer trafficCost,
+                          final Integer nonOverlapCost) {
             this.acceptableRecoveryLag = validated(StreamsConfig.ACCEPTABLE_RECOVERY_LAG_CONFIG, acceptableRecoveryLag);
             this.maxWarmupReplicas = validated(StreamsConfig.MAX_WARMUP_REPLICAS_CONFIG, maxWarmupReplicas);
             this.numStandbyReplicas = validated(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, numStandbyReplicas);
             this.probingRebalanceIntervalMs = validated(StreamsConfig.PROBING_REBALANCE_INTERVAL_MS_CONFIG, probingRebalanceIntervalMs);
             this.rackAwareAssignmentTags = validated(StreamsConfig.RACK_AWARE_ASSIGNMENT_TAGS_CONFIG, rackAwareAssignmentTags);
+            // TODO: validate after we add config
+            this.trafficCost = trafficCost;
+            this.nonOverlapCost = nonOverlapCost;
         }
 
         private static <T> T validated(final String configKey, final T value) {

streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/Graph.java

@@ -42,8 +42,12 @@ public Edge(final V destination, final int capacity, final int cost, final int r
             this(destination, capacity, cost, residualFlow, flow, true);
         }
 
-        public Edge(final V destination, final int capacity, final int cost, final int residualFlow, final int flow,
-            final boolean forwardEdge) {
+        public Edge(final V destination,
+                    final int capacity,
+                    final int cost,
+                    final int residualFlow,
+                    final int flow,
+                    final boolean forwardEdge) {
             Objects.requireNonNull(destination);
             if (capacity < 0) {
                 throw new IllegalArgumentException("Edge capacity cannot be negative");
@@ -72,8 +76,11 @@ public boolean equals(final Object other) {
 
             final Graph<?>.Edge otherEdge = (Graph<?>.Edge) other;
 
-            return destination.equals(otherEdge.destination) && capacity == otherEdge.capacity
-                && cost == otherEdge.cost && residualFlow == otherEdge.residualFlow && flow == otherEdge.flow
+            return destination.equals(otherEdge.destination)
+                && capacity == otherEdge.capacity
+                && cost == otherEdge.cost
+                && residualFlow == otherEdge.residualFlow
+                && flow == otherEdge.flow
                 && forwardEdge == otherEdge.forwardEdge;
         }
 
@@ -84,8 +91,15 @@ public int hashCode() {
 
         @Override
         public String toString() {
-            return "{destination= " + destination + ", capacity=" + capacity + ", cost=" + cost
-                + ", residualFlow=" + residualFlow + ", flow=" + flow + ", forwardEdge=" + forwardEdge;
+            return "Edge {"
+                + "destination= " + destination
+                + ", capacity=" + capacity
+                + ", cost=" + cost
+                + ", residualFlow=" + residualFlow
+                + ", flow=" + flow
+                + ", forwardEdge=" + forwardEdge
+                + "}";
+
         }
     }
 
@@ -106,12 +120,13 @@ public void addEdge(final V u, final V v, final int capacity, final int cost, fi
         addEdge(u, new Edge(v, capacity, cost, capacity - flow, flow));
     }
 
-    public Set<V> nodes() {
+    public SortedSet<V> nodes() {
         return nodes;
     }
 
-    public Map<V, Edge> edges(final V node) {
-        return adjList.get(node);
+    public SortedMap<V, Edge> edges(final V node) {
+        final SortedMap<V, Edge> edge = adjList.get(node);
+        return edge == null ? new TreeMap<>() : edge;
     }
 
     public boolean isResidualGraph() {
@@ -126,12 +141,12 @@ public void setSinkNode(final V node) {
         sinkNode = node;
     }
 
-    public int totalCost() {
-        int totalCost = 0;
+    public long totalCost() {
+        long totalCost = 0;
         for (final Map.Entry<V, SortedMap<V, Edge>> nodeEdges : adjList.entrySet()) {
             final SortedMap<V, Edge> edges = nodeEdges.getValue();
-            for (final Entry<V, Edge> nodeEdge : edges.entrySet()) {
-                totalCost += nodeEdge.getValue().cost * nodeEdge.getValue().flow;
+            for (final Edge nodeEdge : edges.values()) {
+                totalCost += nodeEdge.cost * nodeEdge.flow;
             }
         }
         return totalCost;