KAFKA-20116: Forwad task-(end)-offset to task assignor (5/5)

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

@@ -160,7 +160,6 @@
 import static org.apache.kafka.coordinator.group.Utils.throwIfEmptyString;
 import static org.apache.kafka.coordinator.group.Utils.throwIfNotEmptyCollection;
 import static org.apache.kafka.coordinator.group.Utils.throwIfNotNull;
-import static org.apache.kafka.coordinator.group.Utils.throwIfNotNullOrEmpty;
 import static org.apache.kafka.coordinator.group.Utils.throwIfNull;
 
 /**
@@ -617,7 +616,6 @@ private static void throwIfStreamsGroupHeartbeatRequestIsInvalid(
     private static void throwIfStreamsGroupHeartbeatRequestIsUsingUnsupportedFeatures(
         StreamsGroupHeartbeatRequestData request
     ) throws InvalidRequestException {
-        throwIfNotNullOrEmpty(request.warmupTasks(), "WarmupTasks are not supported yet.");
         if (request.topology() != null) {
             for (StreamsGroupHeartbeatRequestData.Subtopology subtopology : request.topology().subtopologies()) {
                 throwIfNotEmptyCollection(subtopology.sourceTopicRegex(), "Regular expressions for source topics are not supported yet.");

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

@@ -65,6 +65,7 @@
 import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.Endpoint;
 import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.KeyValue;
 import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.TaskIds;
+import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.TaskOffset;
 import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.Topology;
 import org.apache.kafka.common.message.StreamsGroupHeartbeatResponseData;
 import org.apache.kafka.common.message.StreamsGroupHeartbeatResponseData.Status;
@@ -2078,6 +2079,8 @@ private CoordinatorResult<StreamsGroupHeartbeatResult, CoordinatorRecord> stream
         List<TaskIds> ownedActiveTasks,
         List<TaskIds> ownedStandbyTasks,
         List<TaskIds> ownedWarmupTasks,
+        List<TaskOffset> taskOffsets,
+        List<TaskOffset> taskEndOffsets,
         String processId,
         Endpoint userEndpoint,
         List<KeyValue> clientTags,
@@ -2129,6 +2132,14 @@ private CoordinatorResult<StreamsGroupHeartbeatResult, CoordinatorRecord> stream
             );
         }
 
+        // Store the latest task changelog offsets/end-offsets reported by the member. These are transient telemetry
+        // (used by the assignor to estimate task lag) and are not persisted. Task offsets and end-offsets are reported
+        // independently: a null list means "unchanged since the last heartbeat", so we retain the previously reported
+        // value for whichever of the two is null and only update when at least one is reported.
+        if (taskOffsets != null || taskEndOffsets != null) {
+            group.updateTaskOffsets(memberId, group.taskOffsets(memberId).update(taskOffsets, taskEndOffsets));
+        }
+
         // 1. Create or update the member.
         StreamsGroupMember.Builder updatedMemberBuilder = new StreamsGroupMember.Builder(member)
             .maybeUpdateInstanceId(Optional.ofNullable(instanceId))
@@ -4354,7 +4365,8 @@ private UpdateTargetAssignmentResult<TasksTuple> maybeUpdateStreamsTargetAssignm
                 .withMembers(updatedMembersAndTargetAssignment.members())
                 .withTopology(configuredTopology)
                 .withMetadataImage(metadataImage)
-                .withTargetAssignment(updatedMembersAndTargetAssignment.targetAssignment());
+                .withTargetAssignment(updatedMembersAndTargetAssignment.targetAssignment())
+                .withTaskOffsets(group.taskOffsets());
 
             long startTimeMs = time.milliseconds();
             org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult assignmentResult =
@@ -5405,6 +5417,8 @@ public CoordinatorResult<StreamsGroupHeartbeatResult, CoordinatorRecord> streams
                 request.activeTasks(),
                 request.standbyTasks(),
                 request.warmupTasks(),
+                request.taskOffsets(),
+                request.taskEndOffsets(),
                 request.processId(),
                 request.userEndpoint(),
                 request.clientTags(),

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

@@ -0,0 +1,65 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData;
+import org.apache.kafka.coordinator.group.streams.assignor.TaskId;
+
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * The latest per-task cumulative changelog offsets and end-offsets that a member reported in its heartbeat.
+ * <p>
+ * These values are transient telemetry that the assignor uses to estimate task lag (for warm-up promotion).
+ * Per KIP-1071 they are <b>not</b> persisted to the {@code __consumer_offsets} topic ("not persisted, as they are
+ * constantly changing"); they are held in memory on the group coordinator and re-reported by the member on the
+ * task-offset interval.
+ *
+ * @param taskOffsets    Cumulative changelog offsets per task.
+ * @param taskEndOffsets Cumulative changelog end-offsets per task.
+ */
+public record MemberTaskOffsets(Map<TaskId, Long> taskOffsets, Map<TaskId, Long> taskEndOffsets) {
+
+    public static final MemberTaskOffsets EMPTY = new MemberTaskOffsets(Map.of(), Map.of());
+
+    /**
+     * Returns a copy of these offsets updated with the values reported in a heartbeat. Task offsets and task
+     * end-offsets are reported independently: a {@code null} list means "unchanged since the last heartbeat", so the
+     * corresponding map is retained from this instance even when the other one is updated.
+     *
+     * @param reportedTaskOffsets    The reported task offsets, or {@code null} if unchanged.
+     * @param reportedTaskEndOffsets The reported task end-offsets, or {@code null} if unchanged.
+     */
+    public MemberTaskOffsets update(
+        final List<StreamsGroupHeartbeatRequestData.TaskOffset> reportedTaskOffsets,
+        final List<StreamsGroupHeartbeatRequestData.TaskOffset> reportedTaskEndOffsets
+    ) {
+        return new MemberTaskOffsets(
+            reportedTaskOffsets == null ? taskOffsets : toTaskIdMap(reportedTaskOffsets),
+            reportedTaskEndOffsets == null ? taskEndOffsets : toTaskIdMap(reportedTaskEndOffsets)
+        );
+    }
+
+    private static Map<TaskId, Long> toTaskIdMap(final List<StreamsGroupHeartbeatRequestData.TaskOffset> taskOffsets) {
+        return taskOffsets.stream().collect(Collectors.toMap(
+            taskOffset -> new TaskId(taskOffset.subtopologyId(), taskOffset.partition()),
+            StreamsGroupHeartbeatRequestData.TaskOffset::offset
+        ));
+    }
+}

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

@@ -194,6 +194,14 @@ public static class DeadlineAndEpoch {
     private final TimelineHashMap<String, TimelineHashMap<Integer, Set<String>>> currentStandbyTaskToProcessIds;
     private final TimelineHashMap<String, TimelineHashMap<Integer, Set<String>>> currentWarmupTaskToProcessIds;
 
+    /**
+     * The latest per-task changelog offsets and end-offsets reported by each member, keyed by member ID.
+     * This is transient telemetry that the assignor uses to estimate task lag for warm-up promotion. Per KIP-1071
+     * it is not persisted to the {@code __consumer_offsets} topic; it is held in memory and re-reported by members
+     * on the task-offset interval (and is therefore lost on coordinator failover until re-reported).
+     */
+    private final Map<String, MemberTaskOffsets> taskOffsets = new HashMap<>();
+
     /**
      * The Streams topology.
      */
@@ -522,6 +530,32 @@ public void removeMember(String memberId) {
         removeStaticMember(oldMember);
         maybeUpdateGroupState();
         endpointToPartitionsCache.remove(memberId);
+        taskOffsets.remove(memberId);
+    }
+
+    /**
+     * Updates the latest per-task changelog offsets reported by a member. These are transient and not persisted.
+     *
+     * @param memberId        The member ID.
+     * @param memberOffsets   The reported task offsets and end-offsets.
+     */
+    public void updateTaskOffsets(String memberId, MemberTaskOffsets memberOffsets) {
+        taskOffsets.put(memberId, memberOffsets);
+    }
+
+    /**
+     * @return The latest per-task changelog offsets reported by the given member, or
+     *         {@link MemberTaskOffsets#EMPTY} if the member has not reported any.
+     */
+    public MemberTaskOffsets taskOffsets(String memberId) {
+        return taskOffsets.getOrDefault(memberId, MemberTaskOffsets.EMPTY);
+    }
+
+    /**
+     * @return An immutable map of the latest per-task changelog offsets reported by each member, keyed by member ID.
+     */
+    public Map<String, MemberTaskOffsets> taskOffsets() {
+        return Collections.unmodifiableMap(taskOffsets);
     }
 
     /**

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

@@ -91,6 +91,12 @@ public class TargetAssignmentBuilder {
      */
     private ConfiguredTopology topology;
 
+    /**
+     * The latest per-task changelog offsets reported by each member, keyed by member ID. Transient (not persisted);
+     * fed to the assignor so it can estimate task lag.
+     */
+    private Map<String, MemberTaskOffsets> taskOffsets = Map.of();
+
     /**
      * Constructs the object.
      *
@@ -112,7 +118,8 @@ public TargetAssignmentBuilder(
 
     static AssignmentMemberSpec createAssignmentMemberSpec(
         StreamsGroupMember member,
-        TasksTuple targetAssignment
+        TasksTuple targetAssignment,
+        MemberTaskOffsets taskOffsets
     ) {
         return new AssignmentMemberSpec(
             member.instanceId(),
@@ -122,8 +129,8 @@ static AssignmentMemberSpec createAssignmentMemberSpec(
             targetAssignment.warmupTasks(),
             member.processId(),
             member.clientTags(),
-            Map.of(),
-            Map.of()
+            taskOffsets.taskOffsets(),
+            taskOffsets.taskEndOffsets()
         );
     }
 
@@ -151,6 +158,19 @@ public TargetAssignmentBuilder withMembers(
         return this;
     }
 
+    /**
+     * Adds the latest per-task changelog offsets reported by each member.
+     *
+     * @param taskOffsets The reported task offsets/end-offsets keyed by member ID.
+     * @return This object.
+     */
+    public TargetAssignmentBuilder withTaskOffsets(
+        Map<String, MemberTaskOffsets> taskOffsets
+    ) {
+        this.taskOffsets = taskOffsets;
+        return this;
+    }
+
     /**
      * Adds the metadata image to use.
      *
@@ -202,7 +222,8 @@ public TargetAssignmentResult build() throws TaskAssignorException {
         // Prepare the member spec for all members.
         members.forEach((memberId, member) -> memberSpecs.put(memberId, createAssignmentMemberSpec(
             member,
-            targetAssignment.getOrDefault(memberId, org.apache.kafka.coordinator.group.streams.TasksTuple.EMPTY)
+            targetAssignment.getOrDefault(memberId, org.apache.kafka.coordinator.group.streams.TasksTuple.EMPTY),
+            taskOffsets.getOrDefault(memberId, MemberTaskOffsets.EMPTY)
         )));
 
         // Compute the assignment.

group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java

@@ -601,7 +601,7 @@ public void testStreamsGroupHeartbeatFailsForUnsupportedFeatures() throws Except
             new StreamsGroupHeartbeatResult(
                 new StreamsGroupHeartbeatResponseData()
                     .setErrorCode(Errors.INVALID_REQUEST.code())
-                    .setErrorMessage("WarmupTasks are not supported yet."),
+                    .setErrorMessage("Regular expressions for source topics are not supported yet."),
                 Map.of(),
                 -1,
                 -1,
@@ -610,29 +610,55 @@ public void testStreamsGroupHeartbeatFailsForUnsupportedFeatures() throws Except
             service.streamsGroupHeartbeat(
                 context,
                 new StreamsGroupHeartbeatRequestData()
-                    .setWarmupTasks(List.of(new StreamsGroupHeartbeatRequestData.TaskIds()))
+                    .setTopology(new StreamsGroupHeartbeatRequestData.Topology()
+                        .setSubtopologies(List.of(new StreamsGroupHeartbeatRequestData.Subtopology()
+                            .setSourceTopicRegex(List.of("foo.*"))
+                        ))
+                    )
             ).get(5, TimeUnit.SECONDS)
         );
+    }
 
-        assertEquals(
+    @Test
+    public void testStreamsGroupHeartbeatAcceptsTaskOffsetsAndWarmupTasks() throws ExecutionException, InterruptedException, TimeoutException {
+        CoordinatorRuntime<GroupCoordinatorShard, CoordinatorRecord> runtime = mockRuntime();
+        GroupCoordinatorService service = new GroupCoordinatorServiceBuilder()
+            .setRuntime(runtime)
+            .setConfig(createConfig())
+            .build(true);
+
+        StreamsGroupHeartbeatRequestData request = new StreamsGroupHeartbeatRequestData()
+            .setGroupId("foo")
+            .setMemberId(Uuid.randomUuid().toString())
+            .setMemberEpoch(1)
+            .setActiveTasks(List.of())
+            .setStandbyTasks(List.of())
+            .setWarmupTasks(List.of(new StreamsGroupHeartbeatRequestData.TaskIds()))
+            .setTaskOffsets(List.of(new StreamsGroupHeartbeatRequestData.TaskOffset()))
+            .setTaskEndOffsets(List.of(new StreamsGroupHeartbeatRequestData.TaskOffset()));
+
+        when(runtime.scheduleWriteOperation(
+            ArgumentMatchers.eq("streams-group-heartbeat"),
+            ArgumentMatchers.eq(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0)),
+            ArgumentMatchers.any()
+        )).thenReturn(CompletableFuture.completedFuture(
             new StreamsGroupHeartbeatResult(
-                new StreamsGroupHeartbeatResponseData()
-                    .setErrorCode(Errors.INVALID_REQUEST.code())
-                    .setErrorMessage("Regular expressions for source topics are not supported yet."),
+                new StreamsGroupHeartbeatResponseData(),
                 Map.of(),
                 -1,
                 -1,
                 -1
-            ),
-            service.streamsGroupHeartbeat(
-                context,
-                new StreamsGroupHeartbeatRequestData()
-                    .setTopology(new StreamsGroupHeartbeatRequestData.Topology()
-                        .setSubtopologies(List.of(new StreamsGroupHeartbeatRequestData.Subtopology()
-                            .setSourceTopicRegex(List.of("foo.*"))
-                        ))
-                    )
-            ).get(5, TimeUnit.SECONDS)
+            )
+        ));
+
+        CompletableFuture<StreamsGroupHeartbeatResult> future = service.streamsGroupHeartbeat(
+            requestContext(ApiKeys.STREAMS_GROUP_HEARTBEAT),
+            request
+        );
+
+        assertEquals(
+            new StreamsGroupHeartbeatResult(new StreamsGroupHeartbeatResponseData(), Map.of(), -1, -1, -1),
+            future.get(5, TimeUnit.SECONDS)
         );
     }