diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/SendBuilder.java b/clients/src/main/java/org/apache/kafka/common/protocol/SendBuilder.java index b76e0d8dea833..46afacdc9224f 100644 --- a/clients/src/main/java/org/apache/kafka/common/protocol/SendBuilder.java +++ b/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()); diff --git a/clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java b/clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java index 0f55a508fca9f..3756e11f2556a 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java +++ b/clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java @@ -62,8 +62,8 @@ public Iterable records() { } @Override - public DefaultRecordsSend toSend() { - return new DefaultRecordsSend(this); + public DefaultRecordsSend toSend() { + return new DefaultRecordsSend<>(this); } private Iterator recordsIterator() { diff --git a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordsSend.java b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordsSend.java index 13b9428045bd4..bbb17d4b460c0 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordsSend.java +++ b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordsSend.java @@ -20,12 +20,12 @@ import java.io.IOException; -public class DefaultRecordsSend extends RecordsSend { - public DefaultRecordsSend(Records records) { +public class DefaultRecordsSend extends RecordsSend { + public DefaultRecordsSend(T records) { this(records, records.sizeInBytes()); } - public DefaultRecordsSend(Records records, int maxBytesToWrite) { + public DefaultRecordsSend(T records, int maxBytesToWrite) { super(records, maxBytesToWrite); } diff --git a/clients/src/main/java/org/apache/kafka/common/record/FileRecords.java b/clients/src/main/java/org/apache/kafka/common/record/FileRecords.java index 7da47c797f86d..fdb0d299231ca 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/FileRecords.java +++ b/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); } /** diff --git a/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecordsSend.java b/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecordsSend.java index a070b6d0d0ebc..17addef74de4e 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecordsSend.java +++ b/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); } diff --git a/clients/src/main/java/org/apache/kafka/common/record/MemoryRecords.java b/clients/src/main/java/org/apache/kafka/common/record/MemoryRecords.java index ebc5dc4233a14..82a54afe6f7f8 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/MemoryRecords.java +++ b/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); } /** diff --git a/clients/src/main/java/org/apache/kafka/common/record/Records.java b/clients/src/main/java/org/apache/kafka/common/record/Records.java index ae39dd5749b77..5278de4183761 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/Records.java +++ b/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 diff --git a/clients/src/main/java/org/apache/kafka/common/record/TransferableRecords.java b/clients/src/main/java/org/apache/kafka/common/record/TransferableRecords.java new file mode 100644 index 0000000000000..09c0304a0c285 --- /dev/null +++ b/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; +} diff --git a/clients/src/main/java/org/apache/kafka/common/record/UnalignedFileRecords.java b/clients/src/main/java/org/apache/kafka/common/record/UnalignedFileRecords.java new file mode 100644 index 0000000000000..96970f992bb75 --- /dev/null +++ b/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); + } +} diff --git a/clients/src/main/java/org/apache/kafka/common/record/UnalignedMemoryRecords.java b/clients/src/main/java/org/apache/kafka/common/record/UnalignedMemoryRecords.java new file mode 100644 index 0000000000000..23795e306488c --- /dev/null +++ b/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); + } + +} diff --git a/clients/src/main/java/org/apache/kafka/common/record/UnalignedRecords.java b/clients/src/main/java/org/apache/kafka/common/record/UnalignedRecords.java new file mode 100644 index 0000000000000..561d1afe8479e --- /dev/null +++ b/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 toSend() { + return new DefaultRecordsSend<>(this, sizeInBytes()); + } +} diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Utils.java b/clients/src/main/java/org/apache/kafka/common/utils/Utils.java index a1400e736beb2..2ad0957491850 100755 --- a/clients/src/main/java/org/apache/kafka/common/utils/Utils.java +++ b/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. diff --git a/clients/src/main/resources/common/message/FetchSnapshotResponse.json b/clients/src/main/resources/common/message/FetchSnapshotResponse.json index 711b5361c9546..4e6ff85b103e2 100644 --- a/clients/src/main/resources/common/message/FetchSnapshotResponse.json +++ b/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" } ]} ]} ] diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/SendBuilderTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/SendBuilderTest.java index bec74bb2fafc9..6d36395db4e9b 100644 --- a/clients/src/test/java/org/apache/kafka/common/protocol/SendBuilderTest.java +++ b/clients/src/test/java/org/apache/kafka/common/protocol/SendBuilderTest.java @@ -22,6 +22,7 @@ import org.apache.kafka.common.record.MemoryRecordsBuilder; import org.apache.kafka.common.record.SimpleRecord; import org.apache.kafka.common.record.TimestampType; +import org.apache.kafka.common.record.UnalignedMemoryRecords; import org.apache.kafka.common.utils.Utils; import org.apache.kafka.test.TestUtils; import org.junit.jupiter.api.Test; @@ -100,6 +101,38 @@ public void testZeroCopyRecords() { assertEquals(15, readBuffer.getInt()); } + @Test + public void testZeroCopyUnalignedRecords() { + ByteBuffer buffer = ByteBuffer.allocate(128); + MemoryRecords records = createRecords(buffer, "foo"); + + ByteBuffer buffer1 = records.buffer().duplicate(); + buffer1.limit(buffer1.limit() / 2); + + ByteBuffer buffer2 = records.buffer().duplicate(); + buffer2.position(buffer2.limit() / 2); + + UnalignedMemoryRecords records1 = new UnalignedMemoryRecords(buffer1); + UnalignedMemoryRecords records2 = new UnalignedMemoryRecords(buffer2); + + SendBuilder builder = new SendBuilder(8); + builder.writeInt(5); + builder.writeRecords(records1); + builder.writeRecords(records2); + builder.writeInt(15); + Send send = builder.build(); + + // Overwrite the original buffer in order to prove the data was not copied + buffer.rewind(); + MemoryRecords overwrittenRecords = createRecords(buffer, "bar"); + + ByteBuffer readBuffer = TestUtils.toBuffer(send); + assertEquals(5, readBuffer.getInt()); + assertEquals(overwrittenRecords, getRecords(readBuffer, records.sizeInBytes())); + assertEquals(15, readBuffer.getInt()); + } + + private String getString(ByteBuffer buffer, int size) { byte[] readData = new byte[size]; buffer.get(readData); diff --git a/clients/src/test/java/org/apache/kafka/common/record/FileRecordsTest.java b/clients/src/test/java/org/apache/kafka/common/record/FileRecordsTest.java index 34fe34e1bdbad..390ddeaeeb20a 100644 --- a/clients/src/test/java/org/apache/kafka/common/record/FileRecordsTest.java +++ b/clients/src/test/java/org/apache/kafka/common/record/FileRecordsTest.java @@ -23,6 +23,7 @@ import org.apache.kafka.common.utils.MockTime; import org.apache.kafka.common.utils.Time; import org.apache.kafka.test.TestUtils; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -74,6 +75,11 @@ public void setup() throws IOException { this.time = new MockTime(); } + @AfterEach + public void cleanup() throws IOException { + this.fileRecords.close(); + } + @Test public void testAppendProtectsFromOverflow() throws Exception { File fileMock = mock(File.class); diff --git a/clients/src/test/java/org/apache/kafka/common/record/UnalignedFileRecordsTest.java b/clients/src/test/java/org/apache/kafka/common/record/UnalignedFileRecordsTest.java new file mode 100644 index 0000000000000..9a05a22ca5dcc --- /dev/null +++ b/clients/src/test/java/org/apache/kafka/common/record/UnalignedFileRecordsTest.java @@ -0,0 +1,79 @@ +/* + * 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.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Iterator; + +import static org.apache.kafka.test.TestUtils.tempFile; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class UnalignedFileRecordsTest { + + private byte[][] values = new byte[][] { + "foo".getBytes(), + "bar".getBytes() + }; + private FileRecords fileRecords; + + @BeforeEach + public void setup() throws IOException { + this.fileRecords = createFileRecords(values); + } + + @AfterEach + public void cleanup() throws IOException { + this.fileRecords.close(); + } + + @Test + public void testWriteTo() throws IOException { + + org.apache.kafka.common.requests.ByteBufferChannel channel = new org.apache.kafka.common.requests.ByteBufferChannel(fileRecords.sizeInBytes()); + int size = fileRecords.sizeInBytes(); + + UnalignedFileRecords records1 = fileRecords.sliceUnaligned(0, size / 2); + UnalignedFileRecords records2 = fileRecords.sliceUnaligned(size / 2, size - size / 2); + + records1.writeTo(channel, 0, records1.sizeInBytes()); + records2.writeTo(channel, 0, records2.sizeInBytes()); + + channel.close(); + Iterator records = MemoryRecords.readableRecords(channel.buffer()).records().iterator(); + for (byte[] value : values) { + assertTrue(records.hasNext()); + assertEquals(records.next().value(), ByteBuffer.wrap(value)); + } + } + + private FileRecords createFileRecords(byte[][] values) throws IOException { + FileRecords fileRecords = FileRecords.open(tempFile()); + + for (byte[] value : values) { + fileRecords.append(MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord(value))); + } + + return fileRecords; + } +} diff --git a/clients/src/test/java/org/apache/kafka/test/TestUtils.java b/clients/src/test/java/org/apache/kafka/test/TestUtils.java index 5bf8a674e852a..24c170d846787 100644 --- a/clients/src/test/java/org/apache/kafka/test/TestUtils.java +++ b/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.UnalignedRecords; 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(UnalignedRecords records) { + return toBuffer(records.toSend()); + } + public static Set generateRandomTopicPartitions(int numTopic, int numPartitionPerTopic) { Set tps = new HashSet<>(); for (int i = 0; i < numTopic; i++) { diff --git a/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java b/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java index e4a0b3fd3124e..37e139f214644 100644 --- a/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java +++ b/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java @@ -44,6 +44,8 @@ import org.apache.kafka.common.record.CompressionType; import org.apache.kafka.common.record.MemoryRecords; import org.apache.kafka.common.record.Records; +import org.apache.kafka.common.record.UnalignedMemoryRecords; +import org.apache.kafka.common.record.UnalignedRecords; import org.apache.kafka.common.requests.BeginQuorumEpochRequest; import org.apache.kafka.common.requests.BeginQuorumEpochResponse; import org.apache.kafka.common.requests.DescribeQuorumRequest; @@ -74,7 +76,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 +1192,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 +1204,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,10 +1217,8 @@ 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()) ); } @@ -1231,10 +1226,8 @@ private FetchSnapshotResponseData handleFetchSnapshotRequest( 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 +1238,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())); + } + + UnalignedRecords records = snapshot.read(partitionSnapshot.position(), Math.min(data.maxBytes(), maxSnapshotSize)); long snapshotSize = snapshot.sizeInBytes(); @@ -1262,7 +1257,7 @@ private FetchSnapshotResponseData handleFetchSnapshotRequest( return responsePartitionSnapshot .setSize(snapshotSize) .setPosition(partitionSnapshot.position()) - .setBytes(buffer); + .setUnalignedRecords(records); } ); } @@ -1340,7 +1335,10 @@ 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()); + if (!(partitionSnapshot.unalignedRecords() instanceof MemoryRecords)) { + throw new IllegalStateException(String.format("Received unexpected fetch snapshot response: %s", partitionSnapshot)); + } + snapshot.append(new UnalignedMemoryRecords(((MemoryRecords) partitionSnapshot.unalignedRecords()).buffer())); if (snapshot.sizeInBytes() == partitionSnapshot.size()) { // Finished fetching the snapshot. diff --git a/raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotReader.java b/raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotReader.java index d517006d4d68f..d0218c79cc427 100644 --- a/raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotReader.java +++ b/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.FileRecords; import org.apache.kafka.common.record.RecordBatch; +import org.apache.kafka.common.record.UnalignedRecords; 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 iterator() { return Utils.covariantCast(fileRecords.batchIterator()); } - @Override - public int read(ByteBuffer buffer, long position) throws IOException { - return fileRecords.channel().read(buffer, position); + public UnalignedRecords read(long position, int size) { + return fileRecords.sliceUnaligned(Math.toIntExact(position), size); } @Override diff --git a/raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotWriter.java b/raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotWriter.java index 934a12481ccf1..9fe0c8253e1f1 100644 --- a/raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotWriter.java +++ b/raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotWriter.java @@ -16,11 +16,12 @@ */ package org.apache.kafka.snapshot; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.UnalignedMemoryRecords; import org.apache.kafka.common.utils.Utils; import org.apache.kafka.raft.OffsetAndEpoch; import java.io.IOException; -import java.nio.ByteBuffer; import java.nio.channels.FileChannel; import java.nio.file.Files; import java.nio.file.Path; @@ -53,14 +54,23 @@ public long sizeInBytes() throws IOException { } @Override - public void append(ByteBuffer buffer) throws IOException { + public void append(UnalignedMemoryRecords 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) ); } + Utils.writeFully(channel, records.buffer()); + } - Utils.writeFully(channel, buffer); + @Override + 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) + ); + } + Utils.writeFully(channel, records.buffer()); } @Override diff --git a/raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotReader.java b/raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotReader.java index 6b10fa89bba59..6d1ff28950b33 100644 --- a/raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotReader.java +++ b/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.RecordBatch; +import org.apache.kafka.common.record.UnalignedRecords; import org.apache.kafka.raft.OffsetAndEpoch; +import java.io.Closeable; +import java.io.IOException; + /** * Interface for reading snapshots as a sequence of records. */ @@ -29,24 +30,24 @@ public interface RawSnapshotReader extends Closeable, Iterable { /** * Returns the end offset and epoch for the snapshot. */ - public OffsetAndEpoch snapshotId(); + OffsetAndEpoch snapshotId(); /** * Returns the number of bytes for the snapshot. * * @throws IOException for any IO error while reading the size */ - public long sizeInBytes() throws IOException; + long sizeInBytes() throws IOException; /** * Reads bytes from position into the given buffer. * * 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; + UnalignedRecords read(long position, int size) throws IOException; } diff --git a/raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotWriter.java b/raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotWriter.java index 001d2ef37c78e..8bf1bc3703e08 100644 --- a/raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotWriter.java +++ b/raft/src/main/java/org/apache/kafka/snapshot/RawSnapshotWriter.java @@ -17,10 +17,12 @@ package org.apache.kafka.snapshot; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.UnalignedMemoryRecords; +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. @@ -29,39 +31,51 @@ public interface RawSnapshotWriter extends Closeable { /** * Returns the end offset and epoch for the snapshot. */ - public OffsetAndEpoch snapshotId(); + OffsetAndEpoch snapshotId(); /** * Returns the number of bytes for the snapshot. * * @throws IOException for any IO error while reading the size */ - public long sizeInBytes() throws IOException; + long sizeInBytes() throws IOException; + + /** + * Fully appends the memory record set to the snapshot. + * + * If the method returns without an exception the given record set was fully writing the + * snapshot. + * + * @param records the region to append + * @throws IOException for any IO error during append + */ + void append(MemoryRecords records) throws IOException; /** - * Fully appends the buffer to the snapshot. + * Fully appends the memory record set to the snapshot, the difference with {@link RawSnapshotWriter#append(MemoryRecords)} + * is that the record set are fetched from leader by FetchSnapshotRequest, so the records are unaligned. * - * If the method returns without an exception the given buffer was fully writing the + * If the method returns without an exception the given records 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; + void append(UnalignedMemoryRecords records) throws IOException; /** * Returns true if the snapshot has been frozen, otherwise false is returned. * * Modification to the snapshot are not allowed once it is frozen. */ - public boolean isFrozen(); + boolean isFrozen(); /** * Freezes the snapshot and marking it as immutable. * * @throws IOException for any IO error during freezing */ - public void freeze() throws IOException; + void freeze() throws IOException; /** * Closes the snapshot writer. @@ -70,5 +84,5 @@ public interface RawSnapshotWriter extends Closeable { * * @throws IOException for any IO error during close */ - public void close() throws IOException; + void close() throws IOException; } diff --git a/raft/src/main/java/org/apache/kafka/snapshot/SnapshotWriter.java b/raft/src/main/java/org/apache/kafka/snapshot/SnapshotWriter.java index a2a5fa5c6d409..542e54608bc2c 100644 --- a/raft/src/main/java/org/apache/kafka/snapshot/SnapshotWriter.java +++ b/raft/src/main/java/org/apache/kafka/snapshot/SnapshotWriter.java @@ -17,16 +17,17 @@ 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.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.util.List; /** * A type for writing a snapshot fora given end offset and epoch. @@ -146,8 +147,8 @@ public void close() throws IOException { private void appendBatches(List> batches) throws IOException { try { - for (CompletedBatch batch : batches) { - snapshot.append(batch.data.buffer()); + for (CompletedBatch batch : batches) { + snapshot.append(batch.data); } } finally { batches.forEach(CompletedBatch::release); diff --git a/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientSnapshotTest.java b/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientSnapshotTest.java index a26b0536e496b..1b7d862dfc946 100644 --- a/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientSnapshotTest.java +++ b/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientSnapshotTest.java @@ -16,12 +16,6 @@ */ package org.apache.kafka.raft; -import java.io.IOException; -import java.nio.ByteBuffer; -import java.util.Arrays; -import java.util.List; -import java.util.Optional; -import java.util.Set; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.memory.MemoryPool; import org.apache.kafka.common.message.FetchResponseData; @@ -29,6 +23,8 @@ import org.apache.kafka.common.message.FetchSnapshotResponseData; import org.apache.kafka.common.protocol.Errors; import org.apache.kafka.common.record.CompressionType; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.UnalignedMemoryRecords; import org.apache.kafka.common.requests.FetchSnapshotRequest; import org.apache.kafka.common.requests.FetchSnapshotResponse; import org.apache.kafka.common.utils.Utils; @@ -38,6 +34,14 @@ import org.apache.kafka.snapshot.SnapshotWriter; import org.apache.kafka.snapshot.SnapshotWriterTest; import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.List; +import java.util.Optional; +import java.util.Set; + import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -126,13 +130,11 @@ public void testFetchSnapshotRequestAsLeader() throws Exception { assertEquals(Errors.NONE, Errors.forCode(response.errorCode())); assertEquals(snapshot.sizeInBytes(), response.size()); assertEquals(0, response.position()); - assertEquals(snapshot.sizeInBytes(), response.bytes().remaining()); + assertEquals(snapshot.sizeInBytes(), response.unalignedRecords().sizeInBytes()); - ByteBuffer buffer = ByteBuffer.allocate(Math.toIntExact(snapshot.sizeInBytes())); - snapshot.read(buffer, 0); - buffer.flip(); + UnalignedMemoryRecords memoryRecords = (UnalignedMemoryRecords) snapshot.read(0, Math.toIntExact(snapshot.sizeInBytes())); - assertEquals(buffer.slice(), response.bytes()); + assertEquals(memoryRecords.buffer(), ((UnalignedMemoryRecords) response.unalignedRecords()).buffer()); } } @@ -172,14 +174,13 @@ public void testPartialFetchSnapshotRequestAsLeader() throws Exception { assertEquals(Errors.NONE, Errors.forCode(response.errorCode())); assertEquals(snapshot.sizeInBytes(), response.size()); assertEquals(0, response.position()); - assertEquals(snapshot.sizeInBytes() / 2, response.bytes().remaining()); + assertEquals(snapshot.sizeInBytes() / 2, response.unalignedRecords().sizeInBytes()); - ByteBuffer snapshotBuffer = ByteBuffer.allocate(Math.toIntExact(snapshot.sizeInBytes())); - snapshot.read(snapshotBuffer, 0); - snapshotBuffer.flip(); + UnalignedMemoryRecords memoryRecords = (UnalignedMemoryRecords) snapshot.read(0, Math.toIntExact(snapshot.sizeInBytes())); + ByteBuffer snapshotBuffer = memoryRecords.buffer(); ByteBuffer responseBuffer = ByteBuffer.allocate(Math.toIntExact(snapshot.sizeInBytes())); - responseBuffer.put(response.bytes()); + responseBuffer.put(((UnalignedMemoryRecords) response.unalignedRecords()).buffer()); ByteBuffer expectedBytes = snapshotBuffer.duplicate(); expectedBytes.limit(Math.toIntExact(snapshot.sizeInBytes() / 2)); @@ -203,9 +204,9 @@ public void testPartialFetchSnapshotRequestAsLeader() throws Exception { assertEquals(Errors.NONE, Errors.forCode(response.errorCode())); assertEquals(snapshot.sizeInBytes(), response.size()); assertEquals(responseBuffer.position(), response.position()); - assertEquals(snapshot.sizeInBytes() - (snapshot.sizeInBytes() / 2), response.bytes().remaining()); + assertEquals(snapshot.sizeInBytes() - (snapshot.sizeInBytes() / 2), response.unalignedRecords().sizeInBytes()); - responseBuffer.put(response.bytes()); + responseBuffer.put(((UnalignedMemoryRecords) response.unalignedRecords()).buffer()); assertEquals(snapshotBuffer, responseBuffer.flip()); } } @@ -1031,7 +1032,7 @@ private static FetchSnapshotResponseData fetchSnapshotResponse( return partitionSnapshot .setSize(size) .setPosition(position) - .setBytes(buffer); + .setUnalignedRecords(MemoryRecords.readableRecords(buffer.slice())); } ); } @@ -1111,11 +1112,22 @@ public long sizeInBytes() { } @Override - public void append(ByteBuffer buffer) { + public void append(UnalignedMemoryRecords records) { + if (frozen) { + throw new RuntimeException("Snapshot is already frozen " + snapshotId); + } + append(records.buffer()); + } + + @Override + public void append(MemoryRecords records) { if (frozen) { throw new RuntimeException("Snapshot is already frozen " + snapshotId); } + append(records.buffer()); + } + private void append(ByteBuffer buffer) { if (!(data.remaining() >= buffer.remaining())) { ByteBuffer old = data; old.flip(); diff --git a/raft/src/test/java/org/apache/kafka/raft/MockLog.java b/raft/src/test/java/org/apache/kafka/raft/MockLog.java index f338ed5897751..ea46229048171 100644 --- a/raft/src/test/java/org/apache/kafka/raft/MockLog.java +++ b/raft/src/test/java/org/apache/kafka/raft/MockLog.java @@ -26,6 +26,8 @@ import org.apache.kafka.common.record.Records; import org.apache.kafka.common.record.SimpleRecord; import org.apache.kafka.common.record.TimestampType; +import org.apache.kafka.common.record.UnalignedMemoryRecords; +import org.apache.kafka.common.record.UnalignedRecords; import org.apache.kafka.common.utils.ByteBufferOutputStream; import org.apache.kafka.common.utils.Utils; import org.apache.kafka.snapshot.RawSnapshotReader; @@ -517,17 +519,23 @@ public long sizeInBytes() { if (frozen) { throw new RuntimeException("Snapshot is already frozen " + snapshotId); } - return data.position(); } @Override - public void append(ByteBuffer buffer) { + public void append(UnalignedMemoryRecords records) { if (frozen) { throw new RuntimeException("Snapshot is already frozen " + snapshotId); } + data.write(records.buffer()); + } - data.write(buffer); + @Override + public void append(MemoryRecords records) { + if (frozen) { + throw new RuntimeException("Snapshot is already frozen " + snapshotId); + } + data.write(records.buffer()); } @Override @@ -577,14 +585,11 @@ public Iterator iterator() { } @Override - public int read(ByteBuffer buffer, long position) { - ByteBuffer copy = data.buffer(); - copy.position((int) position); - copy.limit((int) position + Math.min(copy.remaining(), buffer.remaining())); - - buffer.put(copy); - - return copy.remaining(); + public UnalignedRecords read(long position, int size) { + ByteBuffer buffer = data.buffer(); + buffer.position(Math.toIntExact(position)); + buffer.limit(Math.min(buffer.limit(), Math.toIntExact(position + size))); + return new UnalignedMemoryRecords(buffer.slice()); } @Override diff --git a/raft/src/test/java/org/apache/kafka/snapshot/FileRawSnapshotTest.java b/raft/src/test/java/org/apache/kafka/snapshot/FileRawSnapshotTest.java index e52943558f2cc..b7e6eda12a1f4 100644 --- a/raft/src/test/java/org/apache/kafka/snapshot/FileRawSnapshotTest.java +++ b/raft/src/test/java/org/apache/kafka/snapshot/FileRawSnapshotTest.java @@ -16,22 +16,27 @@ */ package org.apache.kafka.snapshot; -import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.Arrays; -import java.util.Iterator; -import java.util.stream.IntStream; import org.apache.kafka.common.record.BufferSupplier.GrowableBufferSupplier; import org.apache.kafka.common.record.CompressionType; import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.Record; import org.apache.kafka.common.record.RecordBatch; import org.apache.kafka.common.record.SimpleRecord; -import org.apache.kafka.common.record.Record; +import org.apache.kafka.common.record.UnalignedFileRecords; +import org.apache.kafka.common.record.UnalignedMemoryRecords; +import org.apache.kafka.common.utils.Utils; import org.apache.kafka.raft.OffsetAndEpoch; import org.apache.kafka.test.TestUtils; import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Iterator; +import java.util.stream.IntStream; + import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -49,9 +54,9 @@ public void testWritingSnapshot() throws IOException { try (FileRawSnapshotWriter snapshot = FileRawSnapshotWriter.create(tempDir, offsetAndEpoch)) { assertEquals(0, snapshot.sizeInBytes()); - MemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); + UnalignedMemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); for (int i = 0; i < batches; i++) { - snapshot.append(records.buffer()); + snapshot.append(records); expectedSize += records.sizeInBytes(); } @@ -75,9 +80,9 @@ public void testWriteReadSnapshot() throws IOException { ByteBuffer expectedBuffer = ByteBuffer.wrap(randomBytes(bufferSize)); try (FileRawSnapshotWriter snapshot = FileRawSnapshotWriter.create(tempDir, offsetAndEpoch)) { - MemoryRecords records = buildRecords(expectedBuffer); + UnalignedMemoryRecords records = buildRecords(expectedBuffer); for (int i = 0; i < batches; i++) { - snapshot.append(records.buffer()); + snapshot.append(records); } snapshot.freeze(); @@ -107,6 +112,46 @@ public void testWriteReadSnapshot() throws IOException { } } + @Test + public void testPartialWriteReadSnapshot() throws IOException { + Path tempDir = TestUtils.tempDirectory().toPath(); + OffsetAndEpoch offsetAndEpoch = new OffsetAndEpoch(10L, 3); + + ByteBuffer records = buildRecords(ByteBuffer.wrap(Utils.utf8("foo"))).buffer(); + + ByteBuffer expectedBuffer = ByteBuffer.wrap(records.array()); + + ByteBuffer buffer1 = expectedBuffer.duplicate(); + buffer1.position(0); + buffer1.limit(expectedBuffer.limit() / 2); + ByteBuffer buffer2 = expectedBuffer.duplicate(); + buffer2.position(expectedBuffer.limit() / 2); + buffer2.limit(expectedBuffer.limit()); + + try (FileRawSnapshotWriter snapshot = FileRawSnapshotWriter.create(tempDir, offsetAndEpoch)) { + snapshot.append(new UnalignedMemoryRecords(buffer1)); + snapshot.append(new UnalignedMemoryRecords(buffer2)); + snapshot.freeze(); + } + + try (FileRawSnapshotReader snapshot = FileRawSnapshotReader.open(tempDir, offsetAndEpoch)) { + int totalSize = Math.toIntExact(snapshot.sizeInBytes()); + assertEquals(expectedBuffer.remaining(), totalSize); + + UnalignedFileRecords record1 = (UnalignedFileRecords) snapshot.read(0, totalSize / 2); + UnalignedFileRecords record2 = (UnalignedFileRecords) snapshot.read(totalSize / 2, totalSize - totalSize / 2); + + assertEquals(buffer1, TestUtils.toBuffer(record1)); + assertEquals(buffer2, TestUtils.toBuffer(record2)); + + ByteBuffer readBuffer = ByteBuffer.allocate(record1.sizeInBytes() + record2.sizeInBytes()); + readBuffer.put(TestUtils.toBuffer(record1)); + readBuffer.put(TestUtils.toBuffer(record2)); + readBuffer.flip(); + assertEquals(expectedBuffer, readBuffer); + } + } + @Test public void testBatchWriteReadSnapshot() throws IOException { Path tempDir = TestUtils.tempDirectory().toPath(); @@ -121,7 +166,7 @@ public void testBatchWriteReadSnapshot() throws IOException { .range(0, batchSize) .mapToObj(ignore -> ByteBuffer.wrap(randomBytes(bufferSize))).toArray(ByteBuffer[]::new); - snapshot.append(buildRecords(buffers).buffer()); + snapshot.append(buildRecords(buffers)); } snapshot.freeze(); @@ -165,8 +210,8 @@ public void testBufferWriteReadSnapshot() throws IOException { .range(0, batchSize) .mapToObj(ignore -> ByteBuffer.wrap(randomBytes(bufferSize))).toArray(ByteBuffer[]::new); - MemoryRecords records = buildRecords(buffers); - snapshot.append(records.buffer()); + UnalignedMemoryRecords records = buildRecords(buffers); + snapshot.append(records); expectedSize += records.sizeInBytes(); } @@ -211,9 +256,9 @@ public void testAbortedSnapshot() throws IOException { int batches = 10; try (FileRawSnapshotWriter snapshot = FileRawSnapshotWriter.create(tempDir, offsetAndEpoch)) { - MemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); + UnalignedMemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); for (int i = 0; i < batches; i++) { - snapshot.append(records.buffer()); + snapshot.append(records); } } @@ -230,14 +275,14 @@ public void testAppendToFrozenSnapshot() throws IOException { int batches = 10; try (FileRawSnapshotWriter snapshot = FileRawSnapshotWriter.create(tempDir, offsetAndEpoch)) { - MemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); + UnalignedMemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); for (int i = 0; i < batches; i++) { - snapshot.append(records.buffer()); + snapshot.append(records); } snapshot.freeze(); - assertThrows(RuntimeException.class, () -> snapshot.append(records.buffer())); + assertThrows(RuntimeException.class, () -> snapshot.append(records)); } // File should exist and the size should be greater than the sum of all the buffers @@ -253,9 +298,9 @@ public void testCreateSnapshotWithSameId() throws IOException { int batches = 1; try (FileRawSnapshotWriter snapshot = FileRawSnapshotWriter.create(tempDir, offsetAndEpoch)) { - MemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); + UnalignedMemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); for (int i = 0; i < batches; i++) { - snapshot.append(records.buffer()); + snapshot.append(records); } snapshot.freeze(); @@ -263,9 +308,9 @@ public void testCreateSnapshotWithSameId() throws IOException { // Create another snapshot with the same id try (FileRawSnapshotWriter snapshot = FileRawSnapshotWriter.create(tempDir, offsetAndEpoch)) { - MemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); + UnalignedMemoryRecords records = buildRecords(ByteBuffer.wrap(randomBytes(bufferSize))); for (int i = 0; i < batches; i++) { - snapshot.append(records.buffer()); + snapshot.append(records); } snapshot.freeze(); @@ -280,10 +325,11 @@ private static byte[] randomBytes(int size) { return array; } - private static MemoryRecords buildRecords(ByteBuffer... buffers) { - return MemoryRecords.withRecords( + private static UnalignedMemoryRecords buildRecords(ByteBuffer... buffers) { + MemoryRecords records = MemoryRecords.withRecords( CompressionType.NONE, - Arrays.stream(buffers).map(buffer -> new SimpleRecord(buffer)).toArray(SimpleRecord[]::new) + Arrays.stream(buffers).map(SimpleRecord::new).toArray(SimpleRecord[]::new) ); + return new UnalignedMemoryRecords(records.buffer()); } }