KAFKA-18613: Add StreamsGroupHeartbeat handler in the group coordinator (#19114)

Basic streams group heartbeat handling. The main part of are the unit
tests that make sure that we behave, for the most part, like a consumer
group.

- No support for static membership
- No support for configurations (using constants instead)
- No support for regular expressions

Reviewers: Bill Bejeck <[email protected]>, Bruno Cadonna
<[email protected]>

Author: lucasbru

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java

@@ -400,7 +400,7 @@ public CoordinatorResult<ConsumerGroupHeartbeatResponseData, CoordinatorRecord>
      * @param context The request context.
      * @param request The actual StreamsGroupHeartbeat request.
      *
-     * @return A Result containing the StreamsGroupHeartbeat response, a list of internal topics to be created and
+     * @return A result containing the StreamsGroupHeartbeat response, a list of internal topics to be created and
      *         a list of records to update the state machine.
      */
     public CoordinatorResult<StreamsGroupHeartbeatResult, CoordinatorRecord> streamsGroupHeartbeat(

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java

@@ -146,7 +146,7 @@ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(BiFunction<Strin
      * @param ownedAssignment A collection of active, standby and warm-up tasks
      * @return This object.
      */
-    protected CurrentAssignmentBuilder withOwnedAssignment(TasksTuple ownedAssignment) {
+    public CurrentAssignmentBuilder withOwnedAssignment(TasksTuple ownedAssignment) {
         this.ownedTasks = Optional.ofNullable(ownedAssignment);
         return this;
     }

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

@@ -330,29 +330,61 @@ public String staticMemberId(String groupInstanceId) {
     }
 
     /**
-     * Gets or creates a new member but without adding it to the group. Adding a member is done via the
+     * Gets a new member or throws an exception, if the member does not exist.
+     *
+     * @param memberId The member ID.
+     * @throws UnknownMemberIdException If the member is not found.
+     * @return A StreamsGroupMember.
+     */
+    public StreamsGroupMember getMemberOrThrow(
+        String memberId
+    ) throws UnknownMemberIdException {
+        StreamsGroupMember member = members.get(memberId);
+        if (member != null) {
+            return member;
+        }
+
+        throw new UnknownMemberIdException(
+            String.format("Member %s is not a member of group %s.", memberId, groupId)
+        );
+    }
+
+    /**
+     * Gets or creates a new member, but keeping its fields uninitialized. This is used on the replay-path.
+     * The member is not added to the group, adding a member is done via the
      * {@link StreamsGroup#updateMember(StreamsGroupMember)} method.
      *
      * @param memberId          The member ID.
-     * @param createIfNotExists Booleans indicating whether the member must be created if it does not exist.
      * @return A StreamsGroupMember.
      */
-    public StreamsGroupMember getOrMaybeCreateMember(
-        String memberId,
-        boolean createIfNotExists
+    public StreamsGroupMember getOrCreateUninitializedMember(
+        String memberId
     ) throws UnknownMemberIdException {
         StreamsGroupMember member = members.get(memberId);
         if (member != null) {
             return member;
         }
 
-        if (!createIfNotExists) {
-            throw new UnknownMemberIdException(
-                String.format("Member %s is not a member of group %s.", memberId, groupId)
-            );
+        return new StreamsGroupMember.Builder(memberId).build();
+    }
+
+    /**
+     * Gets or creates a new member, setting default values on the fields. This is used on the replay-path.
+     * The member is not added to the group, adding a member is done via the
+     * {@link StreamsGroup#updateMember(StreamsGroupMember)} method.
+     *
+     * @param memberId          The member ID.
+     * @return A StreamsGroupMember.
+     */
+    public StreamsGroupMember getOrCreateDefaultMember(
+        String memberId
+    ) throws UnknownMemberIdException {
+        StreamsGroupMember member = members.get(memberId);
+        if (member != null) {
+            return member;
         }
 
-        return new StreamsGroupMember.Builder(memberId).build();
+        return StreamsGroupMember.Builder.withDefaults(memberId).build();
     }
 
     /**
@@ -363,7 +395,7 @@ public StreamsGroupMember getOrMaybeCreateMember(
      */
     public StreamsGroupMember staticMember(String instanceId) {
         String existingMemberId = staticMemberId(instanceId);
-        return existingMemberId == null ? null : getOrMaybeCreateMember(existingMemberId, false);
+        return existingMemberId == null ? null : getMemberOrThrow(existingMemberId);
     }
 
     /**
@@ -656,7 +688,7 @@ public void validateOffsetCommit(
             memberId.equals(JoinGroupRequest.UNKNOWN_MEMBER_ID) && groupInstanceId == null)
             return;
 
-        final StreamsGroupMember member = getOrMaybeCreateMember(memberId, false);
+        final StreamsGroupMember member = getMemberOrThrow(memberId);
 
         // If the commit is not transactional and the member uses the new streams protocol (KIP-1071),
         // the member should be using the OffsetCommit API version >= 9.

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

@@ -278,6 +278,21 @@ private static Map<String, Set<Integer>> assignmentFromTaskIds(
                 taskIds -> Set.copyOf(taskIds.partitions())));
         }
 
+        public static Builder withDefaults(String memberId) {
+            return new Builder(memberId)
+                .setRebalanceTimeoutMs(-1)
+                .setTopologyEpoch(-1)
+                .setInstanceId(null)
+                .setRackId(null)
+                .setProcessId("")
+                .setClientTags(Collections.emptyMap())
+                .setState(MemberState.STABLE)
+                .setMemberEpoch(0)
+                .setAssignedTasks(TasksTuple.EMPTY)
+                .setTasksPendingRevocation(TasksTuple.EMPTY)
+                .setUserEndpoint(null);
+        }
+
         public StreamsGroupMember build() {
             return new StreamsGroupMember(
                 memberId,

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

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.coordinator.group.streams;
 
+import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData;
 import org.apache.kafka.coordinator.group.generated.StreamsGroupTopologyValue;
 import org.apache.kafka.coordinator.group.generated.StreamsGroupTopologyValue.Subtopology;
 import org.apache.kafka.coordinator.group.generated.StreamsGroupTopologyValue.TopicInfo;
@@ -81,4 +82,17 @@ public static StreamsTopology fromRecord(StreamsGroupTopologyValue record) {
             record.subtopologies().stream().collect(Collectors.toMap(Subtopology::subtopologyId, x -> x))
         );
     }
+
+    /**
+     * Creates an instance of StreamsTopology from a StreamsGroupHeartbeatRequestData request.
+     *
+     * @param topology The topology supplied in the request.
+     * @return The instance of StreamsTopology created from the request.
+     */
+    public static StreamsTopology fromHeartbeatRequest(StreamsGroupHeartbeatRequestData.Topology topology) {
+        StreamsGroupTopologyValue recordValue = StreamsCoordinatorRecordHelpers.convertToStreamsGroupTopologyRecord(topology);
+        final Map<String, StreamsGroupTopologyValue.Subtopology> subtopologyMap = recordValue.subtopologies().stream()
+            .collect(Collectors.toMap(StreamsGroupTopologyValue.Subtopology::subtopologyId, x -> x));
+        return new StreamsTopology(topology.epoch(), subtopologyMap);
+    }
 }

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

@@ -16,12 +16,16 @@
  */
 package org.apache.kafka.coordinator.group.streams;
 
+import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData;
 import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
 
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -30,11 +34,11 @@
  * An immutable tuple containing active, standby and warm-up tasks.
  *
  * @param activeTasks           Active tasks.
- *                              The key of the map is the subtopology ID and the value is the set of partition IDs.
+ *                              The key of the map is the subtopology ID, and the value is the set of partition IDs.
  * @param standbyTasks          Standby tasks.
- *                              The key of the map is the subtopology ID and the value is the set of partition IDs.
+ *                              The key of the map is the subtopology ID, and the value is the set of partition IDs.
  * @param warmupTasks           Warm-up tasks.
- *                              The key of the map is the subtopology ID and the value is the set of partition IDs.
+ *                              The key of the map is the subtopology ID, and the value is the set of partition IDs.
  */
 public record TasksTuple(Map<String, Set<Integer>> activeTasks,
                          Map<String, Set<Integer>> standbyTasks,
@@ -88,7 +92,7 @@ private static Map<String, Set<Integer>> merge(final Map<String, Set<Integer>> t
     /**
      * Checks if this task tuple contains any of the tasks in another task tuple.
      *
-     * @param other The other task tuple.
+     * @param other Another task tuple.
      * @return true if there is at least one active, standby or warm-up task that is present in both tuples.
      */
     public boolean containsAny(TasksTuple other) {
@@ -130,4 +134,63 @@ public static TasksTuple fromTargetAssignmentRecord(StreamsGroupTargetAssignment
                 )
         );
     }
+
+    public String toString() {
+        return "(active=" + taskAssignmentToString(activeTasks) +
+            ", standby=" + taskAssignmentToString(standbyTasks) +
+            ", warmup=" + taskAssignmentToString(warmupTasks) +
+            ')';
+    }
+
+    public static TasksTuple fromHeartbeatRequest(final List<StreamsGroupHeartbeatRequestData.TaskIds> ownedActiveTasks,
+                                                  final List<StreamsGroupHeartbeatRequestData.TaskIds> ownedStandbyTasks,
+                                                  final List<StreamsGroupHeartbeatRequestData.TaskIds> ownedWarmupTasks) {
+        return new TasksTuple(
+            ownedActiveTasks.stream()
+                .collect(Collectors.toMap(
+                    StreamsGroupHeartbeatRequestData.TaskIds::subtopologyId,
+                        taskId -> new HashSet<>(taskId.partitions())
+                    )
+                ),
+            ownedStandbyTasks.stream()
+                .collect(Collectors.toMap(
+                    StreamsGroupHeartbeatRequestData.TaskIds::subtopologyId,
+                        taskId -> new HashSet<>(taskId.partitions())
+                    )
+                ),
+            ownedWarmupTasks.stream()
+                .collect(Collectors.toMap(
+                    StreamsGroupHeartbeatRequestData.TaskIds::subtopologyId,
+                        taskId -> new HashSet<>(taskId.partitions())
+                    )
+                )
+        );
+    }
+
+    /**
+     * @return The provided assignment as a String.
+     *
+     * Example:
+     * [subtopologyID1-0, subtopologyID1-1, subtopologyID2-0, subtopologyID2-1]
+     */
+    private static String taskAssignmentToString(
+        Map<String, Set<Integer>> assignment
+    ) {
+        StringBuilder builder = new StringBuilder("[");
+        Iterator<Entry<String, Set<Integer>>> subtopologyIterator = assignment.entrySet().iterator();
+        while (subtopologyIterator.hasNext()) {
+            Map.Entry<String, Set<Integer>> entry = subtopologyIterator.next();
+            Iterator<Integer> partitionsIterator = entry.getValue().iterator();
+            while (partitionsIterator.hasNext()) {
+                builder.append(entry.getKey());
+                builder.append("-");
+                builder.append(partitionsIterator.next());
+                if (partitionsIterator.hasNext() || subtopologyIterator.hasNext()) {
+                    builder.append(", ");
+                }
+            }
+        }
+        builder.append("]");
+        return builder.toString();
+    }
 }

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

@@ -23,7 +23,9 @@
 /**
  * The task assignment for a Streams group member.
  *
- * @param activeTasks The target tasks assigned to this member keyed by subtopologyId.
+ * @param activeTasks  The active tasks assigned to this member keyed by subtopologyId.
+ * @param standbyTasks The standby tasks assigned to this member keyed by subtopologyId.
+ * @param warmupTasks  The warm-up tasks assigned to this member keyed by subtopologyId.
  */
 public record MemberAssignment(Map<String, Set<Integer>> activeTasks,
                                Map<String, Set<Integer>> standbyTasks,