KAFKA-17125: Streams Sticky Task Assignor

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/GroupSpec.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.coordinator.group.streams.assignor;
 
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -28,6 +29,11 @@ public interface GroupSpec {
      */
     Map<String, AssignmentMemberSpec> members();
 
+    /**
+     * @return The list of subtopologies.
+     */
+    List<String> subtopologies();
+
     /**
      * @return Any configurations passed to the assignor.
      */

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/GroupSpecImpl.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.coordinator.group.streams.assignor;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
@@ -26,10 +27,12 @@
  * @param assignmentConfigs Any configurations passed to the assignor.
  */
 public record GroupSpecImpl(Map<String, AssignmentMemberSpec> members,
+                            List<String> subtopologies,
                             Map<String, String> assignmentConfigs) implements GroupSpec {
 
     public GroupSpecImpl {
         Objects.requireNonNull(members);
+        Objects.requireNonNull(subtopologies);
         Objects.requireNonNull(assignmentConfigs);
     }
 

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/ProcessState.java

@@ -0,0 +1,144 @@
+/*
+ * 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.assignor;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static java.util.Collections.unmodifiableSet;
+import static org.apache.kafka.common.utils.Utils.union;
+
+public class ProcessState {
+    private final String processId;
+    // number of members
+    private int capacity;
+    private double load;
+    private final Map<String, Integer> memberToTaskCounts;
+    private final Map<String, Set<TaskId>> assignedActiveTasks;
+    private final Map<String, Set<TaskId>> assignedStandbyTasks;
+
+
+    ProcessState(final String processId) {
+        this.processId = processId;
+        this.capacity = 0;
+        this.load = Double.MAX_VALUE;
+        this.assignedActiveTasks = new HashMap<>();
+        this.assignedStandbyTasks = new HashMap<>();
+        this.memberToTaskCounts = new HashMap<>();
+    }
+
+
+    public String processId() {
+        return processId;
+    }
+
+    public int capacity() {
+        return capacity;
+    }
+
+    public int totalTaskCount() {
+        return assignedStandbyTasks().size() + assignedActiveTasks().size();
+    }
+
+    public double load() {
+        return load;
+    }
+
+    public Map<String, Integer> memberToTaskCounts() {
+        return memberToTaskCounts;
+    }
+
+    public Set<TaskId> assignedActiveTasks() {
+        return assignedActiveTasks.values().stream()
+            .flatMap(Set::stream)
+            .collect(Collectors.toSet());
+    }
+
+    public Map<String, Set<TaskId>> assignedActiveTasksByMember() {
+        return assignedActiveTasks;
+    }
+
+    public Set<TaskId> assignedStandbyTasks() {
+        return assignedStandbyTasks.values().stream()
+            .flatMap(Set::stream)
+            .collect(Collectors.toSet());
+    }
+
+    public Map<String, Set<TaskId>> assignedStandbyTasksByMember() {
+        return assignedStandbyTasks;
+    }
+
+    public void addTask(final String memberId, final TaskId taskId, final boolean isActive) {
+        if (isActive) {
+            assignedActiveTasks.putIfAbsent(memberId, new HashSet<>());
+            assignedActiveTasks.get(memberId).add(taskId);
+        } else {
+            assignedStandbyTasks.putIfAbsent(memberId, new HashSet<>());
+            assignedStandbyTasks.get(memberId).add(taskId);
+        }
+        memberToTaskCounts.put(memberId, memberToTaskCounts.get(memberId) + 1);
+        computeLoad();
+    }
+
+    private void incrementCapacity() {
+        capacity++;
+        computeLoad();
+    }
+    public void computeLoad() {
+        if (capacity <= 0) {
+            this.load = -1;
+        } else {
+            this.load = (double) totalTaskCount() / capacity;
+        }
+    }
+
+    public void addMember(final String member) {
+        this.memberToTaskCounts.put(member, 0);
+        incrementCapacity();
+    }
+
+    public boolean hasCapacity() {
+        return totalTaskCount() < capacity;
+    }
+
+    public int compareTo(final ProcessState other) {
+        int loadCompare = Double.compare(this.load, other.load());
+        if (loadCompare == 0) {
+            return Integer.compare(other.capacity, this.capacity);
+        }
+        return loadCompare;
+    }
+
+    public boolean hasTask(final TaskId taskId) {
+        return assignedActiveTasks().contains(taskId) || assignedStandbyTasks().contains(taskId);    }
+
+
+    Set<TaskId> assignedTasks() {
+        final Set<TaskId> assignedActiveTaskIds = assignedActiveTasks();
+        final Set<TaskId> assignedStandbyTaskIds = assignedStandbyTasks();
+        return unmodifiableSet(
+            union(
+                () -> new HashSet<>(assignedActiveTaskIds.size() + assignedStandbyTaskIds.size()),
+                assignedActiveTaskIds,
+                assignedStandbyTaskIds
+            )
+        );
+    }
+}