KAFKA-21234: FileRecords#readInto should honor slice boundaries (#23737)

When the metadata layer (e.g. MetadataLoader) tries to read a slice of
records from KRaft which is greater than
KafkaRaftClient#MAX_BATCH_SIZE_BYTES, the RecordsIterator can return
records which are past the slice (HWM). As it iterates through this
slice of
records, FileRecords#readInto reads from disk the minimum of the bytes
until end of the file and the remaining space in the associated buffer.
In the case of the former, this can lead to a metadata error where the
node tries to snapshot with records past the HWM.

For the metadata error to occur, the following three pre-conditions must
be met:

1. One committed read handed to a listener is larger than 8 MiB.
2. The segment file has bytes after the HWM, meaning records this node
already appended but that aren't committed yet.
3. For the error to actually show up, the snapshot generator emits a
snapshot from the resulting image before the HWM catches up.

The underlying cause is the same as KAFKA-20841. Thanks to @asdf2014
for working on this issue as part of
https://github.com/apache/kafka/pull/22963.

Change `FileRecords#readInto` which reads from the underlying file to
the minimum of:

1. The remaining space in the associated buffer, which in the case of
RecordsIterator is the minimum of the FileRecords' record size and
the batch size defined by KRaft.
2. The end byte position of the FileRecords object, rather than the end
of the underlying file. This ensures when a caller constructs a
FileRecords slice of [start position, HWM), no bytes past the HWM are
read into memory.

Co-authored-by: Benedict Jin <asdf2014@apache.org>
Reviewers: José Armando García Sancio <jsancio@apache.org>, Jun Rao
<junrao@gmail.com>, Zhiyan Tang <ztang@confluent.io>, Chia-Ping Tsai
<chia7712@gmail.com>

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

@@ -152,15 +152,21 @@ public class FileRecords extends AbstractRecords implements Closeable {
     }
 
     /**
-     * Read log batches into the given buffer until there are no bytes remaining in the buffer or the end of the file
-     * is reached.
+     * Read log batches into the given buffer until there are no bytes remaining in the buffer or the end of
+     * this FileRecords object is reached.
      *
      * @param buffer The buffer to write the batches to
-     * @param position Position in the buffer to read from
+     * @param position Position to start reading from, relative to the start of this FileRecords object
      * @throws IOException If an I/O error occurs, see {@link FileChannel#read(ByteBuffer, long)} for details on the
-     * possible exceptions
+     * possible exceptions.
+     * @throws IllegalArgumentException If position is negative or greater than {@link #sizeInBytes()}
      */
     public void readInto(ByteBuffer buffer, int position) throws IOException {
+        int currentSize = sizeInBytes();
+        if (position < 0 || position > currentSize) {
+            throw new IllegalArgumentException("Read position " + position + " is outside " + this);
+        }
+        buffer.limit(buffer.position() + Math.min(buffer.remaining(), currentSize - position));
         Utils.readFully(channel, buffer, position + this.start);
         buffer.flip();
     }

clients/src/test/java/org/apache/kafka/common/record/internal/FileRecordsTest.java

@@ -46,6 +46,7 @@ import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static java.util.Arrays.asList;
@@ -245,6 +246,64 @@ public class FileRecordsTest {
         assertEquals(Collections.singletonList(second), batches(read), "Read a single message starting from the second message");
     }
 
+    @Test
+    public void testReadIntoStopsAtEndOfSlice() throws IOException {
+        List<RecordBatch> items = batches(fileRecords);
+        RecordBatch second = items.get(1);
+        // The slice holds only the second batch. The third batch should not be read into the buffer.
+        FileRecords slice = fileRecords.slice(items.get(0).sizeInBytes(), second.sizeInBytes());
+        ByteBuffer buffer = ByteBuffer.allocate(fileRecords.sizeInBytes());
+
+        slice.readInto(buffer, 0);
+
+        assertEquals(second.sizeInBytes(), buffer.remaining());
+        assertEquals(
+            List.of(second.baseOffset()),
+            batches(MemoryRecords.readableRecords(buffer)).stream()
+                .map(RecordBatch::baseOffset).collect(Collectors.toList())
+        );
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    public void testReadIntoStopsAtEndOfRecords(boolean preallocate) throws IOException {
+        try (FileRecords fileRecords = createFileRecords(values, preallocate)) {
+            // The buffer has room past the records, where a preallocated file has zeros instead of ending
+            ByteBuffer buffer = ByteBuffer.allocate(2 * fileRecords.sizeInBytes());
+
+            fileRecords.readInto(buffer, 0);
+
+            assertEquals(fileRecords.sizeInBytes(), buffer.remaining());
+            assertEquals(
+                List.of(0L, 1L, 2L),
+                batches(MemoryRecords.readableRecords(buffer)).stream()
+                    .map(RecordBatch::baseOffset).collect(Collectors.toList())
+            );
+        }
+    }
+
+    @Test
+    public void testReadIntoAtEndOfRecordsReadsNothing() throws IOException {
+        List<RecordBatch> items = batches(fileRecords);
+        FileRecords slice = fileRecords.slice(items.get(0).sizeInBytes(), items.get(1).sizeInBytes());
+        ByteBuffer buffer = ByteBuffer.allocate(fileRecords.sizeInBytes());
+
+        slice.readInto(buffer, slice.sizeInBytes());
+
+        assertEquals(0, buffer.remaining());
+    }
+
+    @Test
+    public void testReadIntoWithPositionOutsideRecords() {
+        List<RecordBatch> items = batches(fileRecords);
+        FileRecords slice = fileRecords.slice(items.get(0).sizeInBytes(), items.get(1).sizeInBytes());
+        ByteBuffer buffer = ByteBuffer.allocate(fileRecords.sizeInBytes());
+
+        // Both positions are inside the file but outside the slice
+        assertThrows(IllegalArgumentException.class, () -> slice.readInto(buffer, -1));
+        assertThrows(IllegalArgumentException.class, () -> slice.readInto(buffer, slice.sizeInBytes() + 1));
+    }
+
     /**
      * Test the MessageSet.searchFor API.
      */
@@ -955,7 +1014,11 @@ public class FileRecordsTest {
     }
 
     private FileRecords createFileRecords(byte[][] values) throws IOException {
-        FileRecords fileRecords = FileRecords.open(tempFile());
+        return createFileRecords(values, false);
+    }
+
+    private FileRecords createFileRecords(byte[][] values, boolean preallocate) throws IOException {
+        FileRecords fileRecords = FileRecords.open(tempFile(), false, 1024 * 1024, preallocate);
         append(fileRecords, values);
         return fileRecords;
     }

raft/src/main/java/org/apache/kafka/raft/internals/RecordsIterator.java

@@ -136,9 +136,14 @@ public final class RecordsIterator<T> implements Iterator<Batch<T>>, AutoCloseab
 
         MemoryRecords memoryRecords = readFileRecords(fileRecords, buffer);
 
-        // firstBatchSize() is always non-null because the minimum buffer is HEADER_SIZE_UP_TO_MAGIC.
-        int firstBatchSize = memoryRecords.firstBatchSize();
-        if (firstBatchSize <= buffer.remaining()) {
+        Integer firstBatchSize = memoryRecords.firstBatchSize();
+        // firstBatchSize is null when fewer than HEADER_SIZE_UP_TO_MAGIC bytes were read,
+        // which only happens at the end of the records. Once all records have been read,
+        // a trailing partial batch can't be completed, so return the records read so far
+        // instead of growing the buffer.
+        if (firstBatchSize == null || firstBatchSize <= buffer.remaining() ||
+            bytesRead >= fileRecords.sizeInBytes()
+        ) {
             return memoryRecords;
         } else {
             logger.info(
@@ -147,7 +152,7 @@ public final class RecordsIterator<T> implements Iterator<Batch<T>>, AutoCloseab
                 firstBatchSize
             );
             // Not enough bytes read; create a bigger buffer
-            ByteBuffer newBuffer = bufferSupplier.get(memoryRecords.firstBatchSize());
+            ByteBuffer newBuffer = bufferSupplier.get(firstBatchSize);
             allocatedBuffer = Optional.of(newBuffer);
 
             newBuffer.put(buffer);

raft/src/test/java/org/apache/kafka/raft/internals/RecordsIteratorTest.java

@@ -55,6 +55,7 @@ import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Mockito;
 
 import java.io.IOException;
@@ -104,8 +105,8 @@ public final class RecordsIteratorTest {
             for (int i = 0; i < 50; i++) {
                 long seed = new Random(System.nanoTime() + i).nextLong();
                 assertDoesNotThrow(
-                        () -> compressionTypeSeed.accept(ct, seed),
-                        () -> "Failed with compressionType=" + ct + ", seed=" + seed);
+                    () -> compressionTypeSeed.accept(ct, seed),
+                    () -> "Failed with compressionType=" + ct + ", seed=" + seed);
             }
         }
     }
@@ -131,6 +132,87 @@ public final class RecordsIteratorTest {
         });
     }
 
+    @Test
+    public void testFileRecordsSliceLargerThanBatchSize() throws IOException {
+        // Models KafkaRaftClient reading committed records: the slice ends at the high-watermark
+        // and the file has uncommitted batches after the end of the slice.
+        List<TestBatch<String>> sliceBatches = List.of(
+            new TestBatch<>(0, 1, 0, List.of("a", "b", "c", "d", "e")),
+            new TestBatch<>(5, 1, 0, List.of("f", "g", "h", "i", "j"))
+        );
+        List<TestBatch<String>> batchesAfterSlice = List.of(
+            new TestBatch<>(10, 1, 0, List.of("k"))
+        );
+
+        try (FileRecords fileRecords = FileRecords.open(TestUtils.tempFile())) {
+            fileRecords.append(buildRecords(CompressionType.NONE, sliceBatches));
+            int sliceSize = fileRecords.sizeInBytes();
+            fileRecords.append(buildRecords(CompressionType.NONE, batchesAfterSlice));
+            FileRecords slice = fileRecords.slice(0, sliceSize);
+
+            // A batch size one byte smaller than the slice forces a second read, which has only
+            // one byte of the slice left to read, but the buffer can hold `batchesAfterSlice`
+            List<TestBatch<String>> actualBatches = new ArrayList<>();
+            try (RecordsIterator<String> iterator = createIterator(
+                slice,
+                BufferSupplier.NO_CACHING,
+                sliceSize - 1,
+                true
+            )) {
+                iterator.forEachRemaining(batch -> actualBatches.add(TestBatch.from(batch)));
+            }
+
+            assertEquals(sliceBatches, actualBatches);
+        }
+    }
+
+    @ParameterizedTest
+    @ValueSource(ints = {1, Records.HEADER_SIZE_UP_TO_MAGIC - 1, Records.HEADER_SIZE_UP_TO_MAGIC})
+    public void testRecordsEndingInPartialBatch(int partialBatchBytes) throws IOException {
+        TestBatch<String> firstBatch = new TestBatch<>(0, 1, 0, List.of("a", "b", "c", "d", "e"));
+        TestBatch<String> secondBatch = new TestBatch<>(5, 1, 0, Collections.nCopies(10, "f"));
+        MemoryRecords records = buildRecords(
+            CompressionType.NONE,
+            List.of(firstBatch, secondBatch)
+        );
+        int firstBatchSize = records.firstBatchSize();
+        // The records end partialBatchBytes into the second batch
+        int partialRecordsSize = firstBatchSize + partialBatchBytes;
+
+        try (FileRecords fileRecords = FileRecords.open(TestUtils.tempFile())) {
+            fileRecords.append(records);
+
+            testIterator(List.of(firstBatch), records.slice(0, partialRecordsSize), true);
+            testIterator(List.of(firstBatch), fileRecords.slice(0, partialRecordsSize), true);
+        }
+    }
+
+    @Test
+    public void testFileRecordsSliceEndingInPartialBatchAfterGrowingBuffer() throws IOException {
+        TestBatch<String> batch = new TestBatch<>(0, 1, 0, Collections.nCopies(10, "a"));
+
+        try (FileRecords fileRecords = FileRecords.open(TestUtils.tempFile())) {
+            fileRecords.append(buildRecords(CompressionType.NONE, List.of(batch)));
+            int batchSize = fileRecords.sizeInBytes();
+            // The slice ends one byte before the end of the batch
+            FileRecords slice = fileRecords.slice(0, batchSize - 1);
+
+            Set<ByteBuffer> allocatedBuffers = Collections.newSetFromMap(new IdentityHashMap<>());
+            BufferSupplier bufferSupplier = mockBufferSupplier(allocatedBuffers);
+            try (RecordsIterator<String> iterator = createIterator(
+                slice,
+                bufferSupplier,
+                true
+            )) {
+                assertFalse(iterator::hasNext);
+                assertThrows(NoSuchElementException.class, iterator::next);
+            }
+
+            Mockito.verify(bufferSupplier).get(batchSize);
+            assertEquals(Set.of(), allocatedBuffers);
+        }
+    }
+
     @Test
     public void testCrcValidation() {
         forCompressionTypeSeed((compressionType, seed) -> {
@@ -323,12 +405,21 @@ public final class RecordsIteratorTest {
         Records records,
         BufferSupplier bufferSupplier,
         boolean validateCrc
+    ) {
+        return createIterator(records, bufferSupplier, Records.HEADER_SIZE_UP_TO_MAGIC, validateCrc);
+    }
+
+    static RecordsIterator<String> createIterator(
+        Records records,
+        BufferSupplier bufferSupplier,
+        int batchSize,
+        boolean validateCrc
     ) {
         return new RecordsIterator<>(
             records,
             RecordsDecodingStrategy.dataAndControl(STRING_SERDE),
             bufferSupplier,
-            Records.HEADER_SIZE_UP_TO_MAGIC,
+            batchSize,
             validateCrc,
             new LogContext()
         );