KAFKA-10694: Implement zero copy for FetchSnapshot

clients/src/main/java/org/apache/kafka/common/protocol/SendBuilder.java

@@ -21,6 +21,7 @@
 import org.apache.kafka.common.record.BaseRecords;
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.record.MultiRecordsSend;
+import org.apache.kafka.common.record.UnalignedMemoryRecords;
 import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.requests.ResponseHeader;
 import org.apache.kafka.common.utils.ByteUtils;
@@ -137,6 +138,9 @@ public void writeRecords(BaseRecords records) {
         if (records instanceof MemoryRecords) {
             flushPendingBuffer();
             addBuffer(((MemoryRecords) records).buffer());
+        } else if (records instanceof UnalignedMemoryRecords) {
+            flushPendingBuffer();
+            addBuffer(((UnalignedMemoryRecords) records).buffer());
         } else {
             flushPendingSend();
             addSend(records.toSend());

clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java

@@ -62,8 +62,8 @@ public Iterable<Record> records() {
     }
 
     @Override
-    public DefaultRecordsSend toSend() {
-        return new DefaultRecordsSend(this);
+    public DefaultRecordsSend<Records> toSend() {
+        return new DefaultRecordsSend<>(this);
     }
 
     private Iterator<Record> recordsIterator() {

clients/src/main/java/org/apache/kafka/common/record/DefaultRecordsSend.java

@@ -20,12 +20,12 @@
 
 import java.io.IOException;
 
-public class DefaultRecordsSend extends RecordsSend<Records> {
-    public DefaultRecordsSend(Records records) {
+public class DefaultRecordsSend<T extends TransferableRecords> extends RecordsSend<T> {
+    public DefaultRecordsSend(T records) {
         this(records, records.sizeInBytes());
     }
 
-    public DefaultRecordsSend(Records records, int maxBytesToWrite) {
+    public DefaultRecordsSend(T records, int maxBytesToWrite) {
         super(records, maxBytesToWrite);
     }
 

clients/src/main/java/org/apache/kafka/common/record/FileRecords.java

@@ -134,6 +134,28 @@ public void readInto(ByteBuffer buffer, int position) throws IOException {
      * @return A sliced wrapper on this message set limited based on the given position and size
      */
     public FileRecords slice(int position, int size) throws IOException {
+        int availableBytes = availableBytes(position, size);
+        int startPosition = this.start + position;
+        return new FileRecords(file, channel, startPosition, startPosition + availableBytes, true);
+    }
+
+    /**
+     * Return a slice of records from this instance, the difference with {@link FileRecords#slice(int, int)} is
+     * that the position is not necessarily on an offset boundary.
+     *
+     * This method is reserved for cases where offset alignment is not necessary, such as in the replication of raft
+     * snapshots.
+     *
+     * @param position The start position to begin the read from
+     * @param size The number of bytes after the start position to include
+     * @return A unaligned slice of records on this message set limited based on the given position and size
+     */
+    public UnalignedFileRecords sliceUnaligned(int position, int size) {
+        int availableBytes = availableBytes(position, size);
+        return new UnalignedFileRecords(channel, this.start + position, availableBytes);
+    }
+
+    private int availableBytes(int position, int size) {
         // Cache current size in case concurrent write changes it
         int currentSizeInBytes = sizeInBytes();
 
@@ -145,10 +167,10 @@ public FileRecords slice(int position, int size) throws IOException {
             throw new IllegalArgumentException("Invalid size: " + size + " in read from " + this);
 
         int end = this.start + position + size;
-        // handle integer overflow or if end is beyond the end of the file
+        // Handle integer overflow or if end is beyond the end of the file
         if (end < 0 || end > start + currentSizeInBytes)
-            end = start + currentSizeInBytes;
-        return new FileRecords(file, channel, this.start + position, end, true);
+            end = this.start + currentSizeInBytes;
+        return end - (this.start + position);
     }
 
     /**

clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecordsSend.java

@@ -91,7 +91,7 @@ public long writeTo(TransferableChannel channel, long previouslyWritten, int rem
                 convertedRecords = buildOverflowBatch(remaining);
             }
 
-            convertedRecordsWriter = new DefaultRecordsSend(convertedRecords, Math.min(convertedRecords.sizeInBytes(), remaining));
+            convertedRecordsWriter = new DefaultRecordsSend<>(convertedRecords, Math.min(convertedRecords.sizeInBytes(), remaining));
         }
         return convertedRecordsWriter.writeTo(channel);
     }

clients/src/main/java/org/apache/kafka/common/record/MemoryRecords.java

@@ -26,6 +26,7 @@
 import org.apache.kafka.common.utils.CloseableIterator;
 import org.apache.kafka.common.utils.Time;
 
+import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -71,11 +72,7 @@ public long writeTo(TransferableChannel channel, long position, int length) thro
             throw new IllegalArgumentException("position+length should not be greater than buffer.limit(), position: "
                     + position + ", length: " + length + ", buffer.limit(): " + buffer.limit());
 
-        int pos = (int) position;
-        ByteBuffer dup = buffer.duplicate();
-        dup.position(pos);
-        dup.limit(pos + length);
-        return channel.write(dup);
+        return Utils.tryWriteTo(channel, (int) position, length, buffer);
     }
 
     /**

clients/src/main/java/org/apache/kafka/common/record/Records.java

@@ -16,11 +16,9 @@
  */
 package org.apache.kafka.common.record;
 
-import org.apache.kafka.common.network.TransferableChannel;
 import org.apache.kafka.common.utils.AbstractIterator;
 import org.apache.kafka.common.utils.Time;
 
-import java.io.IOException;
 import java.util.Iterator;
 
 
@@ -43,7 +41,7 @@
  *
  * See {@link MemoryRecords} for the in-memory representation and {@link FileRecords} for the on-disk representation.
  */
-public interface Records extends BaseRecords {
+public interface Records extends TransferableRecords {
     int OFFSET_OFFSET = 0;
     int OFFSET_LENGTH = 8;
     int SIZE_OFFSET = OFFSET_OFFSET + OFFSET_LENGTH;
@@ -56,16 +54,6 @@ public interface Records extends BaseRecords {
     int MAGIC_LENGTH = 1;
     int HEADER_SIZE_UP_TO_MAGIC = MAGIC_OFFSET + MAGIC_LENGTH;
 
-    /**
-     * Attempts to write the contents of this buffer to a channel.
-     * @param channel The channel to write to
-     * @param position The position in the buffer to write from
-     * @param length The number of bytes to write
-     * @return The number of bytes actually written
-     * @throws IOException For any IO errors
-     */
-    long writeTo(TransferableChannel channel, long position, int length) throws IOException;
-
     /**
      * Get the record batches. Note that the signature allows subclasses
      * to return a more specific batch type. This enables optimizations such as in-place offset

clients/src/main/java/org/apache/kafka/common/record/TransferableRecords.java

@@ -0,0 +1,39 @@
+/*
+ * 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.common.record;
+
+import org.apache.kafka.common.network.TransferableChannel;
+
+import java.io.IOException;
+
+/**
+ * Represents a record set which can be transferred to a channel
+ * @see Records
+ * @see UnalignedRecords
+ */
+public interface TransferableRecords extends BaseRecords {
+
+    /**
+     * Attempts to write the contents of this buffer to a channel.
+     * @param channel The channel to write to
+     * @param position The position in the buffer to write from
+     * @param length The number of bytes to write
+     * @return The number of bytes actually written
+     * @throws IOException For any IO errors
+     */
+    long writeTo(TransferableChannel channel, long position, int length) throws IOException;
+}

clients/src/main/java/org/apache/kafka/common/record/UnalignedFileRecords.java

@@ -0,0 +1,50 @@
+/*
+ * 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.common.record;
+
+import org.apache.kafka.common.network.TransferableChannel;
+
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+
+/**
+ * Represents a file record set which is not necessarily offset-aligned
+ */
+public class UnalignedFileRecords implements UnalignedRecords {
+
+    private final FileChannel channel;
+    private final long position;
+    private final int size;
+
+    public UnalignedFileRecords(FileChannel channel, long position, int size) {
+        this.channel = channel;
+        this.position = position;
+        this.size = size;
+    }
+
+    @Override
+    public int sizeInBytes() {
+        return size;
+    }
+
+    @Override
+    public long writeTo(TransferableChannel destChannel, long previouslyWritten, int remaining) throws IOException {
+        long position = this.position + previouslyWritten;
+        long count = Math.min(remaining, sizeInBytes() - previouslyWritten);
+        return destChannel.transferFrom(channel, position, count);
+    }
+}

clients/src/main/java/org/apache/kafka/common/record/UnalignedMemoryRecords.java

@@ -0,0 +1,56 @@
+/*
+ * 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.common.record;
+
+import org.apache.kafka.common.network.TransferableChannel;
+import org.apache.kafka.common.utils.Utils;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Objects;
+
+/**
+ * Represents a memory record set which is not necessarily offset-aligned
+ */
+public class UnalignedMemoryRecords implements UnalignedRecords {
+
+    private final ByteBuffer buffer;
+
+    public UnalignedMemoryRecords(ByteBuffer buffer) {
+        this.buffer = Objects.requireNonNull(buffer);
+    }
+
+    public ByteBuffer buffer() {
+        return buffer.duplicate();
+    }
+
+    @Override
+    public int sizeInBytes() {
+        return buffer.remaining();
+    }
+
+    @Override
+    public long writeTo(TransferableChannel channel, long position, int length) throws IOException {
+        if (position > Integer.MAX_VALUE)
+            throw new IllegalArgumentException("position should not be greater than Integer.MAX_VALUE: " + position);
+        if (position + length > buffer.limit())
+            throw new IllegalArgumentException("position+length should not be greater than buffer.limit(), position: "
+                    + position + ", length: " + length + ", buffer.limit(): " + buffer.limit());
+        return Utils.tryWriteTo(channel, (int) position, length, buffer);
+    }
+
+}

clients/src/main/java/org/apache/kafka/common/record/UnalignedRecords.java

@@ -0,0 +1,29 @@
+/*
+ * 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.common.record;
+
+/**
+ * Represents a record set which is not necessarily offset-aligned, and is
+ * only used when fetching raft snapshot
+ */
+public interface UnalignedRecords extends TransferableRecords {
+
+    @Override
+    default RecordsSend<? extends BaseRecords> toSend() {
+        return new DefaultRecordsSend<>(this, sizeInBytes());
+    }
+}

clients/src/main/java/org/apache/kafka/common/utils/Utils.java

@@ -23,6 +23,7 @@
 import java.util.TreeSet;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.network.TransferableChannel;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -1077,7 +1078,7 @@ public static void readFully(FileChannel channel, ByteBuffer destinationBuffer,
      *
      * @throws IOException If an I/O error occurs
      */
-    public static final void readFully(InputStream inputStream, ByteBuffer destinationBuffer) throws IOException {
+    public static void readFully(InputStream inputStream, ByteBuffer destinationBuffer) throws IOException {
         if (!destinationBuffer.hasArray())
             throw new IllegalArgumentException("destinationBuffer must be backed by an array");
         int initialOffset = destinationBuffer.arrayOffset() + destinationBuffer.position();
@@ -1098,6 +1099,29 @@ public static void writeFully(FileChannel channel, ByteBuffer sourceBuffer) thro
             channel.write(sourceBuffer);
     }
 
+    /**
+     * Trying to write data in source buffer to a {@link TransferableChannel}, we may need to call this method multiple
+     * times since this method doesn't ensure the data in the source buffer can be fully written to the destination channel.
+     *
+     * @param destChannel The destination channel
+     * @param position From which the source buffer will be written
+     * @param length The max size of bytes can be written
+     * @param sourceBuffer The source buffer
+     *
+     * @return The length of the actual written data
+     * @throws IOException If an I/O error occurs
+     */
+    public static long tryWriteTo(TransferableChannel destChannel,
+                                  int position,
+                                  int length,
+                                  ByteBuffer sourceBuffer) throws IOException {
+
+        ByteBuffer dup = sourceBuffer.duplicate();
+        dup.position(position);
+        dup.limit(position + length);
+        return destChannel.write(dup);
+    }
+
     /**
      * Write the contents of a buffer to an output stream. The bytes are copied from the current position
      * in the buffer.

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

@@ -51,8 +51,8 @@
           "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": "UnalignedRecords", "type": "records", "versions": "0+",
+          "about": "Snapshot data in records format which may not be aligned on an offset boundary" }
       ]}
     ]}
   ]