KAFKA-10694: Implement zero copy for FetchSnapshot

clients/src/main/resources/common/message/FetchSnapshotResponse.json

@@ -51,8 +51,7 @@
           "about": "The total size of the snapshot." },
         { "name": "Position", "type": "int64", "versions": "0+",
           "about": "The starting byte position within the snapshot included in the Bytes field." },
-        { "name": "Bytes", "type": "bytes", "versions": "0+", "zeroCopy": true,
-          "about": "Snapshot data." }
+        { "name": "Bytes", "type": "records", "versions": "0+", "about": "Snapshot data." }
       ]}
     ]}
   ]

clients/src/test/java/org/apache/kafka/test/TestUtils.java

@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.test;
 
-import java.io.FileWriter;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.common.Cluster;
@@ -26,6 +25,7 @@
 import org.apache.kafka.common.network.NetworkReceive;
 import org.apache.kafka.common.network.Send;
 import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.record.BaseRecords;
 import org.apache.kafka.common.requests.ByteBufferChannel;
 import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.utils.Exit;
@@ -34,6 +34,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.File;
+import java.io.FileWriter;
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.nio.ByteBuffer;
@@ -437,6 +438,10 @@ public static ByteBuffer toBuffer(Send send) {
         }
     }
 
+    public static ByteBuffer toBuffer(BaseRecords records) {
+        return toBuffer(records.toSend());
+    }
+
     public static Set<TopicPartition> generateRandomTopicPartitions(int numTopic, int numPartitionPerTopic) {
         Set<TopicPartition> tps = new HashSet<>();
         for (int i = 0; i < numTopic; i++) {

raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java

@@ -40,6 +40,7 @@
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.ApiMessage;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.BaseRecords;
 import org.apache.kafka.common.record.BufferSupplier;
 import org.apache.kafka.common.record.CompressionType;
 import org.apache.kafka.common.record.MemoryRecords;
@@ -74,7 +75,6 @@
 
 import java.io.IOException;
 import java.net.InetSocketAddress;
-import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
@@ -1191,10 +1191,8 @@ private FetchSnapshotResponseData handleFetchSnapshotRequest(
 
             return FetchSnapshotResponse.singleton(
                 unknownTopicPartition,
-                responsePartitionSnapshot -> {
-                    return responsePartitionSnapshot
-                        .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code());
-                }
+                responsePartitionSnapshot -> responsePartitionSnapshot
+                    .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code())
             );
         }
 
@@ -1205,10 +1203,8 @@ private FetchSnapshotResponseData handleFetchSnapshotRequest(
         if (leaderValidation.isPresent()) {
             return FetchSnapshotResponse.singleton(
                 log.topicPartition(),
-                responsePartitionSnapshot -> {
-                    return addQuorumLeader(responsePartitionSnapshot)
-                        .setErrorCode(leaderValidation.get().code());
-                }
+                responsePartitionSnapshot -> addQuorumLeader(responsePartitionSnapshot)
+                    .setErrorCode(leaderValidation.get().code())
             );
         }
 
@@ -1220,21 +1216,17 @@ private FetchSnapshotResponseData handleFetchSnapshotRequest(
         if (!snapshotOpt.isPresent()) {
             return FetchSnapshotResponse.singleton(
                 log.topicPartition(),
-                responsePartitionSnapshot -> {
-                    return addQuorumLeader(responsePartitionSnapshot)
-                        .setErrorCode(Errors.SNAPSHOT_NOT_FOUND.code());
-                }
+                responsePartitionSnapshot -> addQuorumLeader(responsePartitionSnapshot)
+                    .setErrorCode(Errors.SNAPSHOT_NOT_FOUND.code())
             );
         }
 
         try (RawSnapshotReader snapshot = snapshotOpt.get()) {
             if (partitionSnapshot.position() < 0 || partitionSnapshot.position() >= snapshot.sizeInBytes()) {
                 return FetchSnapshotResponse.singleton(
                     log.topicPartition(),
-                    responsePartitionSnapshot -> {
-                        return addQuorumLeader(responsePartitionSnapshot)
-                            .setErrorCode(Errors.POSITION_OUT_OF_RANGE.code());
-                    }
+                    responsePartitionSnapshot -> addQuorumLeader(responsePartitionSnapshot)
+                        .setErrorCode(Errors.POSITION_OUT_OF_RANGE.code())
                 );
             }
 
@@ -1245,9 +1237,11 @@ private FetchSnapshotResponseData handleFetchSnapshotRequest(
                 maxSnapshotSize = Integer.MAX_VALUE;
             }
 
-            ByteBuffer buffer = ByteBuffer.allocate(Math.min(data.maxBytes(), maxSnapshotSize));
-            snapshot.read(buffer, partitionSnapshot.position());
-            buffer.flip();
+            if (partitionSnapshot.position() > Integer.MAX_VALUE) {
+                throw new IllegalStateException(String.format("Trying to fetch a snapshot with position: %d lager than Int.MaxValue", partitionSnapshot.position()));
+            }
+
+            BaseRecords records = snapshot.read(partitionSnapshot.position(), Math.min(data.maxBytes(), maxSnapshotSize));
 
             long snapshotSize = snapshot.sizeInBytes();
 
@@ -1262,7 +1256,7 @@ private FetchSnapshotResponseData handleFetchSnapshotRequest(
                     return responsePartitionSnapshot
                         .setSize(snapshotSize)
                         .setPosition(partitionSnapshot.position())
-                        .setBytes(buffer);
+                        .setBytes(records);
                 }
             );
         }
@@ -1340,7 +1334,7 @@ private boolean handleFetchSnapshotResponse(
             throw new IllegalStateException(String.format("Received fetch snapshot response with an invalid position. Expected %s; Received %s", snapshot.sizeInBytes(), partitionSnapshot.position()));
         }
 
-        snapshot.append(partitionSnapshot.bytes());
+        snapshot.append((MemoryRecords) partitionSnapshot.bytes());
 
         if (snapshot.sizeInBytes() == partitionSnapshot.size()) {
             // Finished fetching the snapshot.

raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotReader.java

@@ -16,15 +16,16 @@
  */
 package org.apache.kafka.snapshot;
 
-import java.io.IOException;
-import java.nio.ByteBuffer;
-import java.nio.file.Path;
-import java.util.Iterator;
+import org.apache.kafka.common.record.BaseRecords;
 import org.apache.kafka.common.record.FileRecords;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.raft.OffsetAndEpoch;
 
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.Iterator;
+
 public final class FileRawSnapshotReader implements RawSnapshotReader {
     private final FileRecords fileRecords;
     private final OffsetAndEpoch snapshotId;
@@ -49,9 +50,8 @@ public Iterator<RecordBatch> iterator() {
         return Utils.covariantCast(fileRecords.batchIterator());
     }
 
-    @Override
-    public int read(ByteBuffer buffer, long position) throws IOException {
-        return fileRecords.channel().read(buffer, position);
+    public BaseRecords read(long position, int size) throws IOException {
+        return fileRecords.slice((int) position, size);
     }
 
     @Override

raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotWriter.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.snapshot;
 
+import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.raft.OffsetAndEpoch;
 
@@ -53,13 +54,13 @@ public long sizeInBytes() throws IOException {
     }
 
     @Override
-    public void append(ByteBuffer buffer) throws IOException {
+    public void append(MemoryRecords records) throws IOException {
         if (frozen) {
             throw new IllegalStateException(
                 String.format("Append is not supported. Snapshot is already frozen: id = %s; temp path = %s", snapshotId, tempSnapshotPath)
             );
         }
-
+        ByteBuffer buffer = records.buffer();
         Utils.writeFully(channel, buffer);
     }
 

raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotReader.java

@@ -16,12 +16,13 @@
  */
 package org.apache.kafka.snapshot;
 
-import java.io.Closeable;
-import java.io.IOException;
-import java.nio.ByteBuffer;
+import org.apache.kafka.common.record.BaseRecords;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.raft.OffsetAndEpoch;
 
+import java.io.Closeable;
+import java.io.IOException;
+
 /**
  * Interface for reading snapshots as a sequence of records.
  */
@@ -43,10 +44,10 @@ public interface RawSnapshotReader extends Closeable, Iterable<RecordBatch> {
      *
      * It is not guarantee that the given buffer will be filled.
      *
-     * @param buffer byte buffer to put the read files
+     * @param size size to read from snapshot file
      * @param position the starting position in the snapshot to read
-     * @return the number of bytes read
+     * @return the region read from snapshot
      * @throws IOException for any IO error while reading the snapshot
      */
-    public int read(ByteBuffer buffer, long position) throws IOException;
+    BaseRecords read(long position, int size) throws IOException;
 }

raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotWriter.java

@@ -17,10 +17,11 @@
 
 package org.apache.kafka.snapshot;
 
+import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.raft.OffsetAndEpoch;
+
 import java.io.Closeable;
 import java.io.IOException;
-import java.nio.ByteBuffer;
-import org.apache.kafka.raft.OffsetAndEpoch;
 
 /**
  * Interface for writing snapshot as a sequence of records.
@@ -44,10 +45,10 @@ public interface RawSnapshotWriter extends Closeable {
      * If the method returns without an exception the given buffer was fully writing the
      * snapshot.
      *
-     * @param buffer the buffer to append
+     * @param records the region to append
      * @throws IOException for any IO error during append
      */
-    public void append(ByteBuffer buffer) throws IOException;
+    public void append(MemoryRecords records) throws IOException;
 
     /**
      * Returns true if the snapshot has been frozen, otherwise false is returned.

raft/src/main/java/org/apache/kafka/snapshot/SnapshotWriter.java

@@ -17,16 +17,19 @@
 
 package org.apache.kafka.snapshot;
 
-import java.io.Closeable;
-import java.io.IOException;
-import java.util.List;
 import org.apache.kafka.common.memory.MemoryPool;
 import org.apache.kafka.common.record.CompressionType;
+import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.raft.OffsetAndEpoch;
 import org.apache.kafka.raft.RecordSerde;
-import org.apache.kafka.raft.internals.BatchAccumulator.CompletedBatch;
 import org.apache.kafka.raft.internals.BatchAccumulator;
+import org.apache.kafka.raft.internals.BatchAccumulator.CompletedBatch;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
 
 /**
  * A type for writing a snapshot fora given end offset and epoch.
@@ -147,7 +150,8 @@ public void close() throws IOException {
     private void appendBatches(List<CompletedBatch<T>> batches) throws IOException {
         try {
             for (CompletedBatch batch : batches) {
-                snapshot.append(batch.data.buffer());
+                ByteBuffer buffer = batch.data.buffer();
+                snapshot.append(MemoryRecords.readableRecords(buffer));
             }
         } finally {
             batches.forEach(CompletedBatch::release);