KAFKA-14944: Reduce CompletedFetch#parseRecord() memory copy

clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java

@@ -82,7 +82,6 @@
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Timer;
-import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.helpers.MessageFormatter;
 
@@ -109,6 +108,8 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static org.apache.kafka.clients.consumer.ConsumerRecord.NULL_SIZE;
+
 /**
  * This class manages the fetching process with the brokers.
  * <p>
@@ -1413,22 +1414,19 @@ private ConsumerRecord<K, V> parseRecord(TopicPartition partition,
                                              RecordBatch batch,
                                              Record record) {
         try {
-            long offset = record.offset();
-            long timestamp = record.timestamp();
-            Optional<Integer> leaderEpoch = maybeLeaderEpoch(batch.partitionLeaderEpoch());
-            TimestampType timestampType = batch.timestampType();
-            Headers headers = new RecordHeaders(record.headers());
-            ByteBuffer keyBytes = record.key();
-            byte[] keyByteArray = keyBytes == null ? null : Utils.toArray(keyBytes);
-            K key = keyBytes == null ? null : this.keyDeserializer.deserialize(partition.topic(), headers, keyByteArray);
-            ByteBuffer valueBytes = record.value();
-            byte[] valueByteArray = valueBytes == null ? null : Utils.toArray(valueBytes);
-            V value = valueBytes == null ? null : this.valueDeserializer.deserialize(partition.topic(), headers, valueByteArray);
+            final long offset = record.offset();
+            final long timestamp = record.timestamp();
+            final Optional<Integer> leaderEpoch = maybeLeaderEpoch(batch.partitionLeaderEpoch());
+            final TimestampType timestampType = batch.timestampType();
+            final Headers headers = new RecordHeaders(record.headers());
+            final ByteBuffer keyBytes = record.key();
+            final int keySize = keyBytes == null ? NULL_SIZE : keyBytes.remaining();
+            final K key = keyBytes == null ? null : this.keyDeserializer.deserialize(partition.topic(), headers, keyBytes);
+            final ByteBuffer valueBytes = record.value();
+            final int valueSize = valueBytes == null ? NULL_SIZE : valueBytes.remaining();
+            final V value = valueBytes == null ? null : this.valueDeserializer.deserialize(partition.topic(), headers, valueBytes);
             return new ConsumerRecord<>(partition.topic(), partition.partition(), offset,
-                                        timestamp, timestampType,
-                                        keyByteArray == null ? ConsumerRecord.NULL_SIZE : keyByteArray.length,
-                                        valueByteArray == null ? ConsumerRecord.NULL_SIZE : valueByteArray.length,
-                                        key, value, headers, leaderEpoch);
+                    timestamp, timestampType, keySize, valueSize, key, value, headers, leaderEpoch);
         } catch (RuntimeException e) {
             throw new RecordDeserializationException(partition, record.offset(),
                 "Error deserializing key/value for partition " + partition +

clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferDeserializer.java

@@ -16,13 +16,22 @@
  */
 package org.apache.kafka.common.serialization;
 
+import org.apache.kafka.common.header.Headers;
+
 import java.nio.ByteBuffer;
 
 public class ByteBufferDeserializer implements Deserializer<ByteBuffer> {
+
+    @Override
     public ByteBuffer deserialize(String topic, byte[] data) {
-        if (data == null)
+        if (data == null) {
             return null;
-
+        }
         return ByteBuffer.wrap(data);
     }
+
+    @Override
+    public ByteBuffer deserialize(String topic, Headers headers, ByteBuffer data) {
+        return data;
+    }
 }

clients/src/main/java/org/apache/kafka/common/serialization/Deserializer.java

@@ -17,8 +17,10 @@
 package org.apache.kafka.common.serialization;
 
 import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.utils.Utils;
 
 import java.io.Closeable;
+import java.nio.ByteBuffer;
 import java.util.Map;
 
 /**
@@ -60,6 +62,17 @@ default T deserialize(String topic, Headers headers, byte[] data) {
         return deserialize(topic, data);
     }
 
+    /**
+     * Deserialize a record value from a ByteBuffer into a value or object.
+     * @param topic topic associated with the data
+     * @param headers headers associated with the record; may be empty.
+     * @param data serialized ByteBuffer; may be null; implementations are recommended to handle null by returning a value or null rather than throwing an exception.
+     * @return deserialized typed data; may be null
+     */
+    default T deserialize(String topic, Headers headers, ByteBuffer data) {
+        return deserialize(topic, headers, Utils.toNullableArray(data));
+    }
+
     /**
      * Close this deserializer.
      * <p>

clients/src/main/java/org/apache/kafka/common/serialization/StringDeserializer.java

@@ -17,8 +17,11 @@
 package org.apache.kafka.common.serialization;
 
 import org.apache.kafka.common.errors.SerializationException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.utils.Utils;
 
 import java.io.UnsupportedEncodingException;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.Map;
 
@@ -27,27 +30,44 @@
  *  value.deserializer.encoding or deserializer.encoding. The first two take precedence over the last.
  */
 public class StringDeserializer implements Deserializer<String> {
+
     private String encoding = StandardCharsets.UTF_8.name();
 
     @Override
     public void configure(Map<String, ?> configs, boolean isKey) {
-        String propertyName = isKey ? "key.deserializer.encoding" : "value.deserializer.encoding";
+        final String propertyName = isKey ? "key.deserializer.encoding" : "value.deserializer.encoding";
         Object encodingValue = configs.get(propertyName);
-        if (encodingValue == null)
+        if (encodingValue == null) {
             encodingValue = configs.get("deserializer.encoding");
-        if (encodingValue instanceof String)
+        }
+
+        if (encodingValue instanceof String) {
             encoding = (String) encodingValue;
+        }
     }
 
     @Override
     public String deserialize(String topic, byte[] data) {
         try {
-            if (data == null)
-                return null;
-            else
-                return new String(data, encoding);
+            return data == null ? null : new String(data, encoding);
         } catch (UnsupportedEncodingException e) {
             throw new SerializationException("Error when deserializing byte[] to string due to unsupported encoding " + encoding);
         }
     }
+
+    @Override
+    public String deserialize(String topic, Headers headers, ByteBuffer data) {
+        if (data == null) {
+            return null;
+        }
+
+        try {
+            if (data.hasArray()) {
+                return new String(data.array(), data.position() + data.arrayOffset(), data.remaining(), encoding);
+            }
+            return new String(Utils.toArray(data), encoding);
+        } catch (UnsupportedEncodingException e) {
+            throw new SerializationException("Error when deserializing ByteBuffer to string due to unsupported encoding " + encoding);
+        }
+    }
 }

clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java

@@ -50,6 +50,7 @@
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.message.FetchResponseData;
 import org.apache.kafka.common.message.HeartbeatResponseData;
@@ -263,6 +264,16 @@ public String deserialize(String topic, byte[] data) {
                     return super.deserialize(topic, data);
                 }
             }
+
+            @Override
+            public String deserialize(String topic, Headers headers, ByteBuffer data) {
+                if (i == recordIndex) {
+                    throw new SerializationException();
+                } else {
+                    i++;
+                    return super.deserialize(topic, headers, data);
+                }
+            }
         };
     }
 

clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java

@@ -70,6 +70,7 @@
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.FetchRequest.PartitionData;
+import org.apache.kafka.common.serialization.ByteBufferDeserializer;
 import org.apache.kafka.common.utils.BufferSupplier;
 import org.apache.kafka.common.record.CompressionType;
 import org.apache.kafka.common.record.ControlRecordType;
@@ -4819,6 +4820,34 @@ public void testEndOffsetsEmpty() {
         assertEquals(emptyMap(), fetcher.endOffsets(emptyList(), time.timer(5000L)));
     }
 
+    @Test
+    public void testFetchNormalByteBuffer() {
+        buildFetcher(new ByteBufferDeserializer(), new ByteBufferDeserializer());
+
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        // normal fetch
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        client.prepareResponse(fullFetchResponse(tidp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+
+        Map<TopicPartition, List<ConsumerRecord<ByteBuffer, ByteBuffer>>> partitionRecords = fetchedRecords();
+        assertTrue(partitionRecords.containsKey(tp0));
+
+        List<ConsumerRecord<ByteBuffer, ByteBuffer>> records = partitionRecords.get(tp0);
+        assertEquals(3, records.size());
+        assertEquals(4L, subscriptions.position(tp0).offset); // this is the next fetching position
+        long offset = 1;
+        for (ConsumerRecord<ByteBuffer, ByteBuffer> record : records) {
+            assertEquals(offset, record.offset());
+            offset += 1;
+        }
+    }
+
     private MockClient.RequestMatcher offsetsForLeaderEpochRequestMatcher(
         TopicPartition topicPartition,
         int currentLeaderEpoch,

clients/src/test/java/org/apache/kafka/common/serialization/SerializationTest.java

@@ -41,16 +41,20 @@ public class SerializationTest {
     final private String topic = "testTopic";
     final private Map<Class<?>, List<Object>> testData = new HashMap<Class<?>, List<Object>>() {
         {
-            put(String.class, Arrays.asList("my string"));
-            put(Short.class, Arrays.asList((short) 32767, (short) -32768));
-            put(Integer.class, Arrays.asList(423412424, -41243432));
-            put(Long.class, Arrays.asList(922337203685477580L, -922337203685477581L));
-            put(Float.class, Arrays.asList(5678567.12312f, -5678567.12341f));
-            put(Double.class, Arrays.asList(5678567.12312d, -5678567.12341d));
-            put(byte[].class, Arrays.asList("my string".getBytes()));
-            put(ByteBuffer.class, Arrays.asList(ByteBuffer.allocate(10).put("my string".getBytes())));
-            put(Bytes.class, Arrays.asList(new Bytes("my string".getBytes())));
-            put(UUID.class, Arrays.asList(UUID.randomUUID()));
+            put(String.class, Arrays.asList(null, "my string"));
+            put(Short.class, Arrays.asList(null, (short) 32767, (short) -32768));
+            put(Integer.class, Arrays.asList(null, 423412424, -41243432));
+            put(Long.class, Arrays.asList(null, 922337203685477580L, -922337203685477581L));
+            put(Float.class, Arrays.asList(null, 5678567.12312f, -5678567.12341f));
+            put(Double.class, Arrays.asList(null, 5678567.12312d, -5678567.12341d));
+            put(byte[].class, Arrays.asList(null, "my string".getBytes()));
+            put(ByteBuffer.class, Arrays.asList(
+                    null,
+                    ByteBuffer.wrap("my string".getBytes()),
+                    ByteBuffer.allocate(10).put("my string".getBytes()),
+                    ByteBuffer.allocateDirect(10).put("my string".getBytes())));
+            put(Bytes.class, Arrays.asList(null, new Bytes("my string".getBytes())));
+            put(UUID.class, Arrays.asList(null, UUID.randomUUID()));
         }
     };
 
@@ -78,6 +82,8 @@ public void allSerdesShouldSupportNull() {
                     "Should support null in " + cls.getSimpleName() + " serialization");
                 assertNull(serde.deserializer().deserialize(topic, null),
                     "Should support null in " + cls.getSimpleName() + " deserialization");
+                assertNull(serde.deserializer().deserialize(topic, null, (ByteBuffer) null),
+                        "Should support null in " + cls.getSimpleName() + " deserialization");
             }
         }
     }
@@ -355,6 +361,23 @@ public void voidDeserializerShouldThrowOnNotNullValues() {
         }
     }
 
+    @Test
+    public void stringDeserializerSupportByteBuffer() {
+        final String data = "Hello, ByteBuffer!";
+        try (Serde<String> serde = Serdes.String()) {
+            final Serializer<String> serializer = serde.serializer();
+            final Deserializer<String> deserializer = serde.deserializer();
+            final byte[] serializedBytes = serializer.serialize(topic, data);
+            final ByteBuffer heapBuff = ByteBuffer.allocate(serializedBytes.length << 1).put(serializedBytes);
+            heapBuff.flip();
+            assertEquals(data, deserializer.deserialize(topic, null, heapBuff));
+
+            final ByteBuffer directBuff = ByteBuffer.allocateDirect(serializedBytes.length << 2).put(serializedBytes);
+            directBuff.flip();
+            assertEquals(data, deserializer.deserialize(topic, null, directBuff));
+        }
+    }
+
     private Serde<String> getStringSerde(String encoder) {
         Map<String, Object> serializerConfigs = new HashMap<String, Object>();
         serializerConfigs.put("key.serializer.encoding", encoder);

streams/src/test/java/org/apache/kafka/streams/state/internals/SerdeThatDoesntHandleNull.java

@@ -16,12 +16,15 @@
  */
 package org.apache.kafka.streams.state.internals;
 
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.Serde;
 import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.serialization.StringSerializer;
 
+import java.nio.ByteBuffer;
+
 class SerdeThatDoesntHandleNull implements Serde<String> {
     @Override
     public Serializer<String> serializer() {
@@ -38,6 +41,14 @@ public String deserialize(final String topic, final byte[] data) {
                 }
                 return super.deserialize(topic, data);
             }
+
+            @Override
+            public String deserialize(final String topic, final Headers headers, final ByteBuffer data) {
+                if (data == null) {
+                    throw new NullPointerException();
+                }
+                return super.deserialize(topic, headers, data);
+            }
         };
     }
 }