KAFKA-18827: Initialize share group state group coordinator impl. [3/N] (#19026)

* This PR adds impl for the initialize share groups call from the Group
Coordinator perspective.
* The initialize call on persister instance will be invoked by the
`GroupCoordinatorService`, based on the response of the
`GroupCoordinatorShard.shareGroupHeartbeat`. If there is new topic
subscription or member assignment change (topic paritions incremented),
the delta share partitions corresponding to the share group in question
are returned as an optional initialize request.
* The request is then sent to the share coordinator as an encapsulated
timer task because we want the heartbeat response to go asynchronously.
* Tests have been added for `GroupCoordinatorService` and
`GroupMetadataManager`. Existing tests have also been updated.
* A new formatter `ShareGroupStatePartitionMetadataFormatter` has been
added for debugging.

Reviewers: Andrew Schofield <[email protected]>

Author: smjn

clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ShareConsumerTest.java

@@ -2726,8 +2726,15 @@ class PlaintextAdminIntegrationTest extends BaseAdminIntegrationTest {
           client.listGroups(options).all.get.stream().filter(_.groupId == testGroupId).count() == 0
         }, s"Expected to find zero groups")
 
-        val describeWithFakeGroupResult = client.describeShareGroups(util.Arrays.asList(testGroupId, fakeGroupId),
-          new DescribeShareGroupsOptions().includeAuthorizedOperations(true))
+        var describeWithFakeGroupResult: DescribeShareGroupsResult = null
+
+        TestUtils.waitUntilTrue(() => {
+          describeWithFakeGroupResult = client.describeShareGroups(util.Arrays.asList(testGroupId, fakeGroupId),
+            new DescribeShareGroupsOptions().includeAuthorizedOperations(true))
+          val members = describeWithFakeGroupResult.describedGroups().get(testGroupId).get().members()
+          members.asScala.flatMap(_.assignment().topicPartitions().asScala).groupBy(_.topic()).nonEmpty
+        }, s"Could not get partitions assigned. Last response $describeWithFakeGroupResult.")
+
         assertEquals(2, describeWithFakeGroupResult.describedGroups().size())
 
         // Test that we can get information about the test share group.

core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala

@@ -125,7 +125,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
       }, msg = s"Could not get partitions assigned. Last response $shareGroupHeartbeatResponse.")
 
       // Verify the response.
-      assertEquals(2, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
       assertEquals(expectedAssignment, shareGroupHeartbeatResponse.data.assignment)
 
       // Leave the group.
@@ -232,14 +232,26 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
 
       // Heartbeats until the partitions are assigned for member 1.
       shareGroupHeartbeatResponse = null
+
       TestUtils.waitUntilTrue(() => {
         shareGroupHeartbeatResponse = connectAndReceive(shareGroupHeartbeatRequest)
-        shareGroupHeartbeatResponse.data.errorCode == Errors.NONE.code && shareGroupHeartbeatResponse.data.assignment != null
+        if (shareGroupHeartbeatResponse.data.errorCode == Errors.NONE.code && shareGroupHeartbeatResponse.data().assignment() != null) {
+          true
+        } else {
+          shareGroupHeartbeatRequest = new ShareGroupHeartbeatRequest.Builder(
+            new ShareGroupHeartbeatRequestData()
+              .setGroupId("grp")
+              .setMemberId(memberId1)
+              .setMemberEpoch(shareGroupHeartbeatResponse.data.memberEpoch()),
+            true
+          ).build()
+          false
+        }
       }, msg = s"Could not get partitions assigned. Last response $shareGroupHeartbeatResponse.")
 
       val topicPartitionsAssignedToMember1 = shareGroupHeartbeatResponse.data.assignment.topicPartitions()
       // Verify the response.
-      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(4, shareGroupHeartbeatResponse.data.memberEpoch)
 
       // Prepare the next heartbeat for member 2.
       shareGroupHeartbeatRequest = new ShareGroupHeartbeatRequest.Builder(
@@ -259,7 +271,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
 
       val topicPartitionsAssignedToMember2 = shareGroupHeartbeatResponse.data.assignment.topicPartitions()
       // Verify the response.
-      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(4, shareGroupHeartbeatResponse.data.memberEpoch)
 
       val partitionsAssigned: util.Set[Integer] = new util.HashSet[Integer]()
       topicPartitionsAssignedToMember1.forEach(topicPartition => {
@@ -290,7 +302,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
       }, msg = s"Could not get partitions assigned. Last response $shareGroupHeartbeatResponse.")
 
       // Verify the response.
-      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(4, shareGroupHeartbeatResponse.data.memberEpoch)
     } finally {
       admin.close()
     }
@@ -369,7 +381,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
       }, msg = s"Could not get partitions assigned. Last response $shareGroupHeartbeatResponse.")
 
       // Verify the response.
-      assertEquals(2, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
 
       // Member leaves the group.
       shareGroupHeartbeatRequest = new ShareGroupHeartbeatRequest.Builder(
@@ -402,7 +414,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
       shareGroupHeartbeatResponse = connectAndReceive(shareGroupHeartbeatRequest)
 
       // Verify the response for member 1.
-      assertEquals(4, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(5, shareGroupHeartbeatResponse.data.memberEpoch)
       assertEquals(memberId, shareGroupHeartbeatResponse.data.memberId)
       // Partition assignment remains intact on rejoining.
       assertEquals(expectedAssignment, shareGroupHeartbeatResponse.data.assignment)
@@ -491,7 +503,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
           shareGroupHeartbeatResponse.data.assignment.topicPartitions.containsAll(expectedAssignment.topicPartitions)
       }, msg = s"Could not get partitions for topic foo and bar assigned. Last response $shareGroupHeartbeatResponse.")
       // Verify the response.
-      assertEquals(2, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
       // Create the topic baz.
       val bazTopicId = TestUtils.createTopicWithAdminRaw(
         admin = admin,
@@ -515,7 +527,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
         new ShareGroupHeartbeatRequestData()
           .setGroupId("grp")
           .setMemberId(memberId)
-          .setMemberEpoch(2),
+          .setMemberEpoch(3),
         true
       ).build()
 
@@ -527,7 +539,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
           shareGroupHeartbeatResponse.data.assignment.topicPartitions.containsAll(expectedAssignment.topicPartitions)
       }, msg = s"Could not get partitions for topic baz assigned. Last response $shareGroupHeartbeatResponse.")
       // Verify the response.
-      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(5, shareGroupHeartbeatResponse.data.memberEpoch)
       // Increasing the partitions of topic bar which is already being consumed in the share group.
       increasePartitions(admin, "bar", 6, Seq.empty)
 
@@ -547,7 +559,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
         new ShareGroupHeartbeatRequestData()
           .setGroupId("grp")
           .setMemberId(memberId)
-          .setMemberEpoch(3),
+          .setMemberEpoch(5),
         true
       ).build()
 
@@ -559,7 +571,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
           shareGroupHeartbeatResponse.data.assignment.topicPartitions.containsAll(expectedAssignment.topicPartitions)
       }, msg = s"Could not update partitions assignment for topic bar. Last response $shareGroupHeartbeatResponse.")
       // Verify the response.
-      assertEquals(4, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(7, shareGroupHeartbeatResponse.data.memberEpoch)
       // Delete the topic foo.
       TestUtils.deleteTopicWithAdmin(
         admin = admin,
@@ -581,7 +593,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
         new ShareGroupHeartbeatRequestData()
           .setGroupId("grp")
           .setMemberId(memberId)
-          .setMemberEpoch(4),
+          .setMemberEpoch(7),
         true
       ).build()
 
@@ -593,7 +605,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
           shareGroupHeartbeatResponse.data.assignment.topicPartitions.containsAll(expectedAssignment.topicPartitions)
       }, msg = s"Could not update partitions assignment for topic foo. Last response $shareGroupHeartbeatResponse.")
       // Verify the response.
-      assertEquals(5, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(8, shareGroupHeartbeatResponse.data.memberEpoch)
     } finally {
       admin.close()
     }
@@ -704,12 +716,24 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
             .setTopicId(barId)
             .setPartitions(List[Integer](0).asJava)).asJava)
 
+      shareGroupHeartbeatResponse = null
+
       TestUtils.waitUntilTrue(() => {
         shareGroupHeartbeatResponse = connectAndReceive(shareGroupHeartbeatRequest)
-        shareGroupHeartbeatResponse.data.errorCode == Errors.NONE.code &&
-          shareGroupHeartbeatResponse.data.assignment != null &&
+        if (shareGroupHeartbeatResponse.data.assignment != null &&
           expectedAssignment.topicPartitions.containsAll(shareGroupHeartbeatResponse.data.assignment.topicPartitions) &&
-          shareGroupHeartbeatResponse.data.assignment.topicPartitions.containsAll(expectedAssignment.topicPartitions)
+          shareGroupHeartbeatResponse.data.assignment.topicPartitions.containsAll(expectedAssignment.topicPartitions)) {
+          true
+        } else {
+          shareGroupHeartbeatRequest = new ShareGroupHeartbeatRequest.Builder(
+            new ShareGroupHeartbeatRequestData()
+              .setGroupId("grp")
+              .setMemberId(memberId)
+              .setMemberEpoch(shareGroupHeartbeatResponse.data.memberEpoch),
+            true
+          ).build()
+          false
+        }
       }, msg = s"Could not get bar partitions assigned. Last response $shareGroupHeartbeatResponse.")
 
       // Verify the response, the epoch should have been bumped.
@@ -840,7 +864,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
           shareGroupHeartbeatResponse.data.assignment == expectedAssignment
       }, msg = s"Could not get partitions assigned. Last response $shareGroupHeartbeatResponse.")
       // Verify the response.
-      assertEquals(2, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
 
       // Restart the only running broker.
       val broker = cluster.brokers().values().iterator().next()
@@ -864,7 +888,7 @@ class ShareGroupHeartbeatRequestTest(cluster: ClusterInstance) {
 
       // Verify the response. Epoch should not have changed and null assignments determines that no
       // change in old assignment.
-      assertEquals(2, shareGroupHeartbeatResponse.data.memberEpoch)
+      assertEquals(3, shareGroupHeartbeatResponse.data.memberEpoch)
       assertNull(shareGroupHeartbeatResponse.data.assignment)
     } finally {
       admin.close()

core/src/test/scala/unit/kafka/server/ShareGroupHeartbeatRequestTest.scala

@@ -30,7 +30,7 @@ public interface SubscribedTopicDescriber {
      *
      * @param topicId   Uuid corresponding to the topic.
      * @return The number of partitions corresponding to the given topic Id,
-     *         or -1 if the topic Id does not exist.
+     *         or -1 if the topic id does not exist.
      */
     int numPartitions(Uuid topicId);
 
@@ -43,4 +43,14 @@ public interface SubscribedTopicDescriber {
      *         If the topic Id does not exist, an empty set is returned.
      */
     Set<String> racksForPartition(Uuid topicId, int partition);
+
+    /**
+     * Returns a set of partitions corresponding to a topic id which
+     * are allowlisted based on some criteria. For example, for share groups
+     * only partitions which have been initialized are returned.
+     *
+     * @param topicId   The uuid of the topic
+     * @return The set of integers representing assignable partitions. Could be empty, or contain all partitions.
+     */
+    Set<Integer> assignablePartitions(Uuid topicId);
 }

group-coordinator/group-coordinator-api/src/main/java/org/apache/kafka/coordinator/group/api/assignor/SubscribedTopicDescriber.java

@@ -47,6 +47,8 @@
 import org.apache.kafka.coordinator.group.generated.ShareGroupMetadataValue;
 import org.apache.kafka.coordinator.group.generated.ShareGroupPartitionMetadataKey;
 import org.apache.kafka.coordinator.group.generated.ShareGroupPartitionMetadataValue;
+import org.apache.kafka.coordinator.group.generated.ShareGroupStatePartitionMetadataKey;
+import org.apache.kafka.coordinator.group.generated.ShareGroupStatePartitionMetadataValue;
 import org.apache.kafka.coordinator.group.generated.ShareGroupTargetAssignmentMemberKey;
 import org.apache.kafka.coordinator.group.generated.ShareGroupTargetAssignmentMemberValue;
 import org.apache.kafka.coordinator.group.generated.ShareGroupTargetAssignmentMetadataKey;
@@ -812,6 +814,44 @@ public static CoordinatorRecord newShareGroupCurrentAssignmentTombstoneRecord(
         );
     }
 
+    /**
+     * Creates a ShareGroupStatePartitionMetadata record.
+     *
+     * @param groupId   The share group id.
+     * @param initializedTopics  Topics which have been initialized.
+     * @param deletingTopics  Topics which are being deleted.
+     * @return The record.
+     */
+    public static CoordinatorRecord newShareGroupStatePartitionMetadataRecord(
+        String groupId,
+        Map<Uuid, Map.Entry<String, Set<Integer>>> initializedTopics,
+        Map<Uuid, String> deletingTopics
+    ) {
+        List<ShareGroupStatePartitionMetadataValue.TopicPartitionsInfo> initializedTopicPartitionInfo = initializedTopics.entrySet().stream()
+            .map(entry -> new ShareGroupStatePartitionMetadataValue.TopicPartitionsInfo()
+                .setTopicId(entry.getKey())
+                .setTopicName(entry.getValue().getKey())
+                .setPartitions(entry.getValue().getValue().stream().toList()))
+            .toList();
+
+        List<ShareGroupStatePartitionMetadataValue.TopicInfo> deletingTopicsInfo = deletingTopics.entrySet().stream()
+            .map(entry -> new ShareGroupStatePartitionMetadataValue.TopicInfo()
+                .setTopicId(entry.getKey())
+                .setTopicName(entry.getValue()))
+            .toList();
+
+        return CoordinatorRecord.record(
+            new ShareGroupStatePartitionMetadataKey()
+                .setGroupId(groupId),
+            new ApiMessageAndVersion(
+                new ShareGroupStatePartitionMetadataValue()
+                    .setInitializedTopics(initializedTopicPartitionInfo)
+                    .setDeletingTopics(deletingTopicsInfo),
+                (short) 0
+            )
+        );
+    }
+
     private static List<ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions> toTopicPartitions(
         Map<Uuid, Set<Integer>> topicPartitions
     ) {