KAFKA-18409: ShareGroupStateMessageFormatter should use CoordinatorRecordMessageFormatter (apache#18510)

ShareGroupStateMessageFormatter should extend
CoordinatorRecordMessageFormatter in order to have a consistent handling
of records of coordinators.

Reviewers: Ken Huang <[email protected]>, David Jacot <[email protected]>

Author: brandboat

tools/src/main/java/org/apache/kafka/tools/consumer/group/share/ShareGroupStateMessageFormatter.java

@@ -16,162 +16,34 @@
  */
 package org.apache.kafka.tools.consumer.group.share;
 
-import org.apache.kafka.clients.consumer.ConsumerRecord;
-import org.apache.kafka.common.MessageFormatter;
-import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.protocol.ApiMessage;
-import org.apache.kafka.common.protocol.ByteBufferAccessor;
-import org.apache.kafka.coordinator.share.generated.CoordinatorRecordType;
-import org.apache.kafka.coordinator.share.generated.ShareSnapshotKey;
-import org.apache.kafka.coordinator.share.generated.ShareSnapshotKeyJsonConverter;
-import org.apache.kafka.coordinator.share.generated.ShareSnapshotValue;
-import org.apache.kafka.coordinator.share.generated.ShareSnapshotValueJsonConverter;
-import org.apache.kafka.coordinator.share.generated.ShareUpdateKey;
-import org.apache.kafka.coordinator.share.generated.ShareUpdateKeyJsonConverter;
-import org.apache.kafka.coordinator.share.generated.ShareUpdateValue;
-import org.apache.kafka.coordinator.share.generated.ShareUpdateValueJsonConverter;
+import org.apache.kafka.coordinator.share.ShareCoordinatorRecordSerde;
+import org.apache.kafka.coordinator.share.generated.CoordinatorRecordJsonConverters;
 import org.apache.kafka.tools.consumer.CoordinatorRecordMessageFormatter;
 
 import com.fasterxml.jackson.databind.JsonNode;
-import com.fasterxml.jackson.databind.node.JsonNodeFactory;
-import com.fasterxml.jackson.databind.node.NullNode;
-import com.fasterxml.jackson.databind.node.ObjectNode;
-import com.fasterxml.jackson.databind.node.TextNode;
-
-import java.io.IOException;
-import java.io.PrintStream;
-import java.nio.ByteBuffer;
-import java.util.Objects;
-import java.util.Optional;
-
-import static java.nio.charset.StandardCharsets.UTF_8;
 
 /**
  * Formatter for records of in __share_group_state topic.
  */
-public class ShareGroupStateMessageFormatter implements MessageFormatter {
-
-    private static final String VERSION = "version";
-    private static final String DATA = "data";
-    private static final String KEY = "key";
-    private static final String VALUE = "value";
-    private static final String UNKNOWN = "unknown";
-
-    @Override
-    public void writeTo(ConsumerRecord<byte[], byte[]> consumerRecord, PrintStream output) {
-        ObjectNode json = new ObjectNode(JsonNodeFactory.instance);
-
-        byte[] key = consumerRecord.key();
-        short keyVersion = -1;
-        if (Objects.nonNull(key)) {
-            keyVersion = ByteBuffer.wrap(key).getShort();
-            JsonNode dataNode = readToKeyJson(ByteBuffer.wrap(key), keyVersion);
-
-            if (dataNode instanceof NullNode) {
-                return;
-            }
-            json.putObject(KEY)
-                .put(VERSION, keyVersion)
-                .set(DATA, dataNode);
-        } else {
-            json.set(KEY, NullNode.getInstance());
-        }
-
-        byte[] value = consumerRecord.value();
-        if (Objects.nonNull(value)) {
-            short valueVersion = ByteBuffer.wrap(value).getShort();
-            JsonNode dataNode = readToValueJson(ByteBuffer.wrap(value), keyVersion, valueVersion);
-
-            json.putObject(VALUE)
-                .put(VERSION, valueVersion)
-                .set(DATA, dataNode);
-        } else {
-            json.set(VALUE, NullNode.getInstance());
-        }
+public class ShareGroupStateMessageFormatter extends CoordinatorRecordMessageFormatter {
 
-        try {
-            output.write(json.toString().getBytes(UTF_8));
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        }
+    public ShareGroupStateMessageFormatter() {
+        super(new ShareCoordinatorRecordSerde());
     }
 
-    private JsonNode readToKeyJson(ByteBuffer byteBuffer, short version) {
-        return readToSnapshotMessageKey(byteBuffer)
-            .map(logKey -> transferKeyMessageToJsonNode(logKey, version))
-            .orElseGet(() -> new TextNode(UNKNOWN));
-    }
-
-    private Optional<ApiMessage> readToSnapshotMessageKey(ByteBuffer byteBuffer) {
-        short version = byteBuffer.getShort();
-        try {
-            switch (CoordinatorRecordType.fromId(version)) {
-                case SHARE_SNAPSHOT:
-                    return Optional.of(new ShareSnapshotKey(new ByteBufferAccessor(byteBuffer), version));
-                case SHARE_UPDATE:
-                    return Optional.of(new ShareUpdateKey(new ByteBufferAccessor(byteBuffer), version));
-                default:
-                    return Optional.empty();
-            }
-        } catch (UnsupportedVersionException ex) {
-            return Optional.empty();
-        }
-    }
-
-    private JsonNode transferKeyMessageToJsonNode(ApiMessage logKey, short keyVersion) {
-        if (logKey instanceof ShareSnapshotKey) {
-            return ShareSnapshotKeyJsonConverter.write((ShareSnapshotKey) logKey, keyVersion);
-        } else if (logKey instanceof ShareUpdateKey) {
-            return ShareUpdateKeyJsonConverter.write((ShareUpdateKey) logKey, keyVersion);
-        }
-        return new TextNode(UNKNOWN);
-    }
-
-    /**
-     * Here the valueVersion is not enough to identity the deserializer for the ByteBuffer.
-     * This is because both {@link ShareSnapshotValue} and {@link ShareUpdateValue} have version 0
-     * as per RPC spec.
-     * To differentiate, we need to use the corresponding key versions. This is acceptable as
-     * the records will always appear in pairs (key, value). However, this means that we cannot
-     * extend {@link CoordinatorRecordMessageFormatter} as it requires overriding
-     * readToValueJson whose signature does not allow for passing keyversion.
-     *
-     * @param byteBuffer - Represents the raw data read from the topic
-     * @param keyVersion - Version of the actual key component of the data read from topic
-     * @param valueVersion - Version of the actual value component of the data read from topic
-     * @return JsonNode corresponding to the raw data value component
-     */
-    protected JsonNode readToValueJson(ByteBuffer byteBuffer, short keyVersion, short valueVersion) {
-        return readToSnapshotMessageValue(byteBuffer, keyVersion)
-            .map(logValue -> transferValueMessageToJsonNode(logValue, valueVersion))
-            .orElseGet(() -> new TextNode(UNKNOWN));
+    @Override
+    protected boolean isRecordTypeAllowed(short recordType) {
+        return true;
     }
 
-    private JsonNode transferValueMessageToJsonNode(ApiMessage logValue, short version) {
-        if (logValue instanceof ShareSnapshotValue) {
-            return ShareSnapshotValueJsonConverter.write((ShareSnapshotValue) logValue, version);
-        } else if (logValue instanceof ShareUpdateValue) {
-            return ShareUpdateValueJsonConverter.write((ShareUpdateValue) logValue, version);
-        }
-        return new TextNode(UNKNOWN);
+    @Override
+    protected JsonNode keyAsJson(ApiMessage message) {
+        return CoordinatorRecordJsonConverters.writeRecordKeyAsJson(message);
     }
 
-    private Optional<ApiMessage> readToSnapshotMessageValue(ByteBuffer byteBuffer, short keyVersion) {
-        short version = byteBuffer.getShort();
-        // Check the key version here as that will determine which type
-        // of value record to fetch. Both share update and share snapshot
-        // value records can have the same version.
-        try {
-            switch (CoordinatorRecordType.fromId(keyVersion)) {
-                case SHARE_SNAPSHOT:
-                    return Optional.of(new ShareSnapshotValue(new ByteBufferAccessor(byteBuffer), version));
-                case SHARE_UPDATE:
-                    return Optional.of(new ShareUpdateValue(new ByteBufferAccessor(byteBuffer), version));
-                default:
-                    return Optional.empty();
-            }
-        } catch (UnsupportedVersionException ex) {
-            return Optional.empty();
-        }
+    @Override
+    protected JsonNode valueAsJson(ApiMessage message, short version) {
+        return CoordinatorRecordJsonConverters.writeRecordValueAsJson(message, version);
     }
 }

tools/src/test/java/org/apache/kafka/tools/consumer/group/share/ShareGroupStateMessageFormatterTest.java

@@ -30,6 +30,8 @@
 import org.apache.kafka.coordinator.share.generated.ShareUpdateValue;
 import org.apache.kafka.server.share.SharePartitionKey;
 import org.apache.kafka.server.share.persister.PersisterStateBatch;
+import org.apache.kafka.tools.consumer.CoordinatorRecordMessageFormatter;
+import org.apache.kafka.tools.consumer.CoordinatorRecordMessageFormatterTest;
 
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
@@ -46,7 +48,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
-public class ShareGroupStateMessageFormatterTest {
+public class ShareGroupStateMessageFormatterTest extends CoordinatorRecordMessageFormatterTest {
     private static final SharePartitionKey KEY_1 = SharePartitionKey.getInstance("gs1", Uuid.fromString("gtb2stGYRk-vWZ2zAozmoA"), 0);
     private static final ShareGroupOffset SHARE_GROUP_OFFSET_1 = new ShareGroupOffset.Builder()
         .setSnapshotEpoch(0)
@@ -132,28 +134,51 @@ public class ShareGroupStateMessageFormatterTest {
                 .collect(Collectors.toList())
         );
 
-    private static Stream<Arguments> parameters() {
+    @Override
+    protected CoordinatorRecordMessageFormatter formatter() {
+        return new ShareGroupStateMessageFormatter();
+    }
+
+    @Override
+    protected Stream<Arguments> parameters() {
         return Stream.of(
             Arguments.of(
                 MessageUtil.toVersionPrefixedByteBuffer((short) 0, SHARE_SNAPSHOT_KEY).array(),
                 MessageUtil.toVersionPrefixedByteBuffer((short) 0, SHARE_SNAPSHOT_VALUE).array(),
-                "{\"key\":{\"version\":0,\"data\":{\"groupId\":\"gs1\",\"topicId\":\"gtb2stGYRk-vWZ2zAozmoA\",\"partition\":0}},\"value\":{\"version\":0,\"data\":{\"snapshotEpoch\":0,\"stateEpoch\":1,\"leaderEpoch\":20,\"startOffset\":50,\"stateBatches\":[{\"firstOffset\":100,\"lastOffset\":200,\"deliveryState\":1,\"deliveryCount\":10},{\"firstOffset\":201,\"lastOffset\":210,\"deliveryState\":2,\"deliveryCount\":10}]}}}"
+                """
+                    {"key":{"type":0,"data":{"groupId":"gs1","topicId":"gtb2stGYRk-vWZ2zAozmoA","partition":0}},
+                     "value":{"version":0,
+                              "data":{"snapshotEpoch":0,
+                                      "stateEpoch":1,
+                                      "leaderEpoch":20,
+                                      "startOffset":50,
+                                      "stateBatches":[{"firstOffset":100,"lastOffset":200,"deliveryState":1,"deliveryCount":10},
+                                                      {"firstOffset":201,"lastOffset":210,"deliveryState":2,"deliveryCount":10}]}}}
+                """
             ),
             Arguments.of(
                 MessageUtil.toVersionPrefixedByteBuffer((short) 1, SHARE_UPDATE_KEY).array(),
                 MessageUtil.toVersionPrefixedByteBuffer((short) 0, SHARE_UPDATE_VALUE).array(),
-                "{\"key\":{\"version\":1,\"data\":{\"groupId\":\"gs2\",\"topicId\":\"r9Nq4xGAQf28jvu36t7gQQ\",\"partition\":0}},\"value\":{\"version\":0,\"data\":{\"snapshotEpoch\":1,\"leaderEpoch\":25,\"startOffset\":55,\"stateBatches\":[{\"firstOffset\":100,\"lastOffset\":150,\"deliveryState\":1,\"deliveryCount\":12},{\"firstOffset\":151,\"lastOffset\":200,\"deliveryState\":2,\"deliveryCount\":15}]}}}"
+                """
+                    {"key":{"type":1,"data":{"groupId":"gs2","topicId":"r9Nq4xGAQf28jvu36t7gQQ","partition":0}},
+                     "value":{"version":0,
+                              "data":{"snapshotEpoch":1,
+                                      "leaderEpoch":25,
+                                      "startOffset":55,
+                                      "stateBatches":[{"firstOffset":100,"lastOffset":150,"deliveryState":1,"deliveryCount":12},
+                                                      {"firstOffset":151,"lastOffset":200,"deliveryState":2,"deliveryCount":15}]}}}
+                """
             ),
             // wrong versions
             Arguments.of(
                 MessageUtil.toVersionPrefixedByteBuffer((short) 10, SHARE_SNAPSHOT_KEY).array(),
                 MessageUtil.toVersionPrefixedByteBuffer((short) 0, SHARE_SNAPSHOT_VALUE).array(),
-                "{\"key\":{\"version\":10,\"data\":\"unknown\"},\"value\":{\"version\":0,\"data\":\"unknown\"}}"
+                ""
             ),
             Arguments.of(
                 MessageUtil.toVersionPrefixedByteBuffer((short) 15, SHARE_UPDATE_KEY).array(),
                 MessageUtil.toVersionPrefixedByteBuffer((short) 0, SHARE_UPDATE_VALUE).array(),
-                "{\"key\":{\"version\":15,\"data\":\"unknown\"},\"value\":{\"version\":0,\"data\":\"unknown\"}}"
+                ""
             )
         );
     }
@@ -164,37 +189,22 @@ private static Stream<Arguments> exceptions() {
             Arguments.of(
                 MessageUtil.toVersionPrefixedByteBuffer((short) 0, SHARE_SNAPSHOT_KEY).array(),
                 MessageUtil.toVersionPrefixedByteBuffer((short) 0, SHARE_UPDATE_VALUE).array(),
-                new RuntimeException("non-nullable field stateBatches was serialized as null")
+                new RuntimeException("""
+                    Could not read record at offset 0 due to: \
+                    Could not read record with version 0 from value's buffer due to: \
+                    non-nullable field stateBatches was serialized as null.""")
             ),
             Arguments.of(
                 MessageUtil.toVersionPrefixedByteBuffer((short) 1, SHARE_UPDATE_KEY).array(),
                 MessageUtil.toVersionPrefixedByteBuffer((short) 0, SHARE_SNAPSHOT_VALUE).array(),
-                new RuntimeException("non-nullable field stateBatches was serialized as null")
+                new RuntimeException("""
+                    Could not read record at offset 0 due to: \
+                    Could not read record with version 0 from value's buffer due to: \
+                    non-nullable field stateBatches was serialized as null.""")
             )
         );
     }
 
-    @ParameterizedTest
-    @MethodSource("parameters")
-    public void testShareGroupStateMessageFormatter(
-        byte[] keyBuffer,
-        byte[] valueBuffer,
-        String expectedOutput
-    ) {
-        ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>(
-            Topic.SHARE_GROUP_STATE_TOPIC_NAME, 0, 0,
-            0L, TimestampType.CREATE_TIME, 0,
-            0, keyBuffer, valueBuffer,
-            new RecordHeaders(), Optional.empty());
-
-        try (MessageFormatter formatter = new ShareGroupStateMessageFormatter()) {
-            formatter.configure(emptyMap());
-            ByteArrayOutputStream out = new ByteArrayOutputStream();
-            formatter.writeTo(record, new PrintStream(out));
-            assertEquals(expectedOutput, out.toString());
-        }
-    }
-
     @ParameterizedTest
     @MethodSource("exceptions")
     public void testShareGroupStateMessageFormatterException(