KAFKA-21234: FileRecords#readInto should honor slice boundaries

[kevin-wu24]

Issue Description

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 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 HW, 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 HW catches up.

The underlying cause is the same as
https://issues.apache.org/jira/browse/KAFKA-20841. Thanks to @asdf2014
for working on this issue as part of
#22963.

What Changed

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

• 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.
• 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.

Testing

Added unit tests in RecordsIteratorTest and FileRecordsTest

Alternatives

• Limiting the buffer in RecordsIterator prior to calling
  FileRecords#readInto. The motivation behind this is that other callers
  may want to use this functionality in the future, so the change belongs
  in FileRecords, not a caller of that class.

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

Reviewers: Jun Rao junrao@gmail.com, Chia-Ping Tsai chia7712@gmail.com, José Armando García Sancio jsancio@apache.org

[jsancio]

Thanks for the fix @kevin-wu24 . I just have one high-level comment.

[jsancio]

Instead of introducing a new method (readUntil). I would like to understand if we should fix readInto to honor the limit of the FileRecords slice. What do you think @junrao

[kevin-wu24]

I am not opposed to that approach. I mainly need to understand what each of the other users of FileRecords#readInto assumes/can handle. I thought it might be better to see if others already knew some of these things. I will look into them myself, but it just may take some time.

[jsancio]

Take a look at writeTo.

I would expect the read semantic of readInto to be similar to writeTo.

[junrao]

@jsancio : Yes, I agree. #22963 seems to fix the same issue with the approach you suggested.

[kevin-wu24]

@jsancio : Yes, I agree. #22963 seems to fix the same issue with the approach you suggested.

Sounds good. I will close this PR and the JIRA as a duplicate.

[junrao]

@kevin-wu24 : Thanks for the updated PR. A few more comments.

[junrao]

Claude found the following issue with RecordsIterator.createMemoryRecords(). The following comment is no longer true can throw a NullPointerException.

    // firstBatchSize() is always non-null because the minimum buffer is HEADER_SIZE_UP_TO_MAGIC.

Example: slice(0, B1 + 10) with batchSize = B1.
The first read gets B1, which is consumed.
compact() leaves an empty buffer.
The capped read returns only 10 bytes.
firstBatchSize() returns null because fewer than HEADER_SIZE_UP_TO_MAGIC (17) bytes remain, and unboxing it into int firstBatchSize throws.
The old code read past the slice instead, which was also wrong but didn't crash. The comment at :139 ("always non-null because the minimum buffer is HEADER_SIZE_UP_TO_MAGIC") is no longer true.
This can't happen in practice today. Every raft caller I checked ends its slice on a batch boundary: the high-watermark or log-end positions for fireHandleCommit (KafkaRaftClient:442) and maybeLoadLog (KRaftControlRecordStateMachine:228), and whole files for snapshots. The iterator should still handle a partial tail.
Fix: once bytesRead == records.sizeInBytes(), the slice is used up, so return memoryRecords as is. Its batch iterator yields nothing and iteration ends. This also skips the "grow buffer" path, which would otherwise allocate a buffer of firstBatchSize for a tail that can't be completed.


Integer firstBatchSize = memoryRecords.firstBatchSize();
if (firstBatchSize == null || firstBatchSize <= buffer.remaining() || bytesRead >= fileRecords.sizeInBytes())
    return memoryRecords;

[kevin-wu24]

I see, because before we would fill up the buffer or go to the file end. I have added Claude's suggestion.

[junrao]

Could we add below that this method now throws IllegalArgumentException ?

[junrao]

Was this comment addressed?

[kevin-wu24]

Apologies, fixed now.

[junrao]

This check is unnecessary since it's always true. Also, the error message in availableBytes() seems very specific to the slice() call. Claude suggested the following change.

if (position < 0 || position > size)
    throw new IllegalArgumentException("Read position " + position + " is outside " + this);
buffer.limit(buffer.position() + Math.min(buffer.remaining(), size - position));

[kevin-wu24]

This check is unnecessary since it's always true.

I am not sure if I agree with this statement. For example, on the first call for an 8MiB buffer at position 0 and trying to read a FileRecords slice of > 8MiB, available = availableBytes(position, buffer.remaining()) = buffer.remaining() == 8MiB, which makes this if evaluate to false. What do you think?

Also, the error message in availableBytes() seems very specific to the slice() call

Sure, but I think the suggested implementation was what I had previously, so I can revert it. I wonder if Claude is running in circles haha. I personally found this previous implementation a little easier to understand than the availableBytes one, but outside of the exception content I agree they are the same.

[junrao]

No need for the new helper. The caller can just do batches(MemoryRecords.readableRecords(buffer)).

[kevin-wu24]

Thanks for the review @junrao. I have pushed a commit to address. Please let me know if there is anything else.

[junrao]

@kevin-wu24 : Thanks for the updated PR. A few more comments.

[junrao]

Was this comment addressed?

[junrao]

Claude suggested the following improvement.

1. Grow the buffer only when the batch can actually be completed (RecordsIterator.java:140-151). Today the grow path runs whenever the first batch doesn't fit in the buffer and the slice isn't fully read. The buffer is then sized to firstBatchSize, even when that batch runs past the end of the slice. That has two costs:

It allocates for a batch that can never be completed.
A corrupt length field (firstBatchSize() parses with maxMessageSize = Integer.MAX_VALUE) can trigger a very large allocation, e.g. an OutOfMemoryError on the controller while it reads a snapshot. The capped read would then fill only a small part of that buffer.
The old code had the same problem, but it's easy to fix here. Compare against the bytes the slice can still supply; that also replaces the bytesRead >= size check:



Integer firstBatchSize = memoryRecords.firstBatchSize();
int bytesAvailable = buffer.remaining() + (fileRecords.sizeInBytes() - bytesRead);
if (firstBatchSize == null || firstBatchSize <= buffer.remaining() || firstBatchSize > bytesAvailable) {
    return memoryRecords;
}
...

[kevin-wu24]

I see. The first clause is for when the buffer holds fewer than 17 bytes, meaning the slice is fully read. The third clause is for when the buffer holds the start of a batch that goes past the slices end, so we should not return that. I have added this suggestion and a test in RecordsIteratorTest.

[junrao]

Perhaps rename size to sth like currentSize to avoid hiding the size field.

[junrao]

Claude suggested the following.

preallocate parameter: only testReadIntoStopsAtEndOfRecords needs it. The other three tests read only written data, so the second run adds nothing.

[junrao]

@kevin-wu24 : Thanks for the updated PR. One more comment.

[junrao]

This change has a subtle side effect. It means that bytesRead may never reach records.sizeInBytes(), but nextBatches() depends on that to end the iteration. Currently, it doesn't seem to cause an issue, but it makes the code harder to understand.

So, it seems that it's simpler to leave the code as

if (firstBatchSize == null || firstBatchSize <= buffer.remaining() ||
            bytesRead >= fileRecords.sizeInBytes()

This way, the current convention of expecting bytesRead to reach records.sizeInBytes() will hold. This could lead to an unnecessary buffer grow, but it's rare and probably not worth optimizing. What do you think?

[kevin-wu24]

It means that bytesRead may never reach records.sizeInBytes(), but nextBatches() depends on that to end the iteration.

I'm not sure if I agree with this statement. I think in RecordsIterator#nextBatch(), RecordsIterator#nextBatches = memoryRecords.batchIterator() with RecordBatchIterator#hasNext() == false in the firstBatchSize == null or firstBatchSize > bytesAvailable case, because there is not a full batch in the buffer. I looked at ByteBufferLogInputStream#nextBatch(), RecordBatchIterator#makeNext(), and AbstractIterator. Because nextBatches.hasNext() == false, RecordsIterator#nextBatch() returns an empty optional, so RecordsIterator#hasNext will return false eventually. Any code that calls RecordsIterator#hasNext again after it evaluates false for the first time is wrong, and I am pretty sure that does not happen in the raft module.

Currently, it doesn't seem to cause an issue, but it makes the code harder to understand

I believe the reason for this is that it is the caller's (raft code's) responsibility that this iterator is requested at batch aligned offsets. This should always be the case because the HWM/LEO are batch aligned, but I don't think the code enforces that in any way.

I do agree this is harder to understand than the bytesRead >= fileRecords.sizeInBytes() condition. I think case 1 and 3 (firstBatchSize > bytesAvailable) of this if clause should only ever be hit if the user of this iterator calls it on a records slice which is not batch aligned, which I would argue is unexpected. In that case, should the code do something stronger: (e.g. log an error, or maybe throw)? I am not sure about throwing because that can impact availability.

What do you think? cc @jsancio @junrao

[junrao]

My concern is on the caller logic. The check if (bytesRead < recordSize) used to always reach false when the iteration finishes for the records. However, with this change, this is no longer the case. It doesn't cause an issue now since existing callers don't depend on that. However, this is a bit dangerous since future callers may not realize this. I think it's better to make the behavior more consistent. We can either change createMemoryRecords() to drain all bytes in the records or keep the optimization in createMemoryRecords(), but remove if (bytesRead < recordSize) in nextBatches().

    private Iterator<MutableRecordBatch> nextBatches() {
        int recordSize = records.sizeInBytes();
        if (bytesRead < recordSize) {
            final MemoryRecords memoryRecords;
            if (records instanceof MemoryRecords) {
                bytesRead = recordSize;
                memoryRecords = (MemoryRecords) records;
            } else if (records instanceof FileRecords) {
                memoryRecords = createMemoryRecords((FileRecords) records);
            } else {
                throw new IllegalStateException(String.format("Unexpected Records type %s", records.getClass()));
            }

            return memoryRecords.batchIterator();
        }

        return Collections.emptyIterator();
    }

[kevin-wu24]

Maybe I am missing something, but I believe that check still always reaches false when the last readFileRecords that finishes reading all the bytes in the FileRecords#slice executes. Each call to FileRecords#readInto reads to the minimum of the remaining buffer space and the remaining bytes in the slice, and then updates bytesRead. Eventually this means bytesRead == recordSize right? Do you have a counter example?

[junrao]

Here is an example. Suppose that the batch size is 1MB. The slice contains record batch of 8MB but has a length of only 3MB. createMemoryRecords() is called, readFileRecords() fills memoryRecords with 1MB from the slice. It then notices firstBatchSize (8MB) > bytesAvailable (3MB) and returns memoryRecords without growing the buffer. The caller will see no records in the returned memoryRecords but with bytesRead < recordSize. createMemoryRecords() is called again. Since memoryRecords is full, readFileRecords() reads nothing. firstBatchSize (8MB) > bytesAvailable (3MB) is still true and the same memoryRecords is returned. This will loop forever.

Note that the above won't happen without the optimization since it will grow the buffer and drain all bytes of the partial batch from the records.

[kevin-wu24]

I see. Basically, repeated createMemoryRecords() invocations no longer guarantee that bytesRead >= recordSize when: buffer capacity < (FileRecords#sizeInBytes − position in buffer) < firstBatchSize. Thanks for the explanation.

I am going to revert to the bytesRead >= fileRecords.sizeInBytes() condition since it is easier to understand. If the extra allocation case ever comes up, we can fix it then. My motivation for this over the other approach is that raft callers should not be requesting non-batch-aligned offsets, which is when the extra allocation would occur.

[junrao]

@kevin-wu24 : Thanks for the updated PR. Just a couple of minor comments.

[junrao]

We can use firstBatchSize in line 151 now.

[junrao]

testReadIntoStopsAtEndOfSlice, testReadIntoAtEndOfRecordsReadsNothing and testReadIntoWithPositionOutsideRecords can just use the fileRecords created in setup().

[0xffff-zhiyan]

@kevin-wu24 Thanks for the PR. It's really a great find! The fix LGTM.

I think this bug also reveals we're actually missing a check that MetadataLoader only gets committed records
Could we add a cheap check in MetadataLoader#handleCommit using the existing highWaterMarkAccessor,
e.g:

OptionalLong hwm = highWaterMarkAccessor.get();
if (hwm.isPresent() && batch.lastOffset() >= hwm.getAsLong()) {
    throw new IllegalStateException(...);
}

With this fix in theory it'd never fire anyway, just a tripwire for future regressions. Or do you think it's overkill?

[kevin-wu24]

@kevin-wu24 Thanks for the PR. It's really a great find! The fix LGTM.

I think this bug also reveals we're actually missing a check that MetadataLoader only gets committed records Could we add a cheap check in MetadataLoader#handleCommit using the existing highWaterMarkAccessor, e.g:

OptionalLong hwm = highWaterMarkAccessor.get();
if (hwm.isPresent() && batch.lastOffset() >= hwm.getAsLong()) {
    throw new IllegalStateException(...);
}

With this fix in theory it'd never fire anyway, just a tripwire for future regressions. Or do you think it's overkill?

Hi @0xffff-zhiyan,

Specifically for the issue we observed, the SnapshotEmitter is the one who logged a metadata error that it was trying to write a snapshot past the HWM, so we did become aware of it via an exception being thrown. For brokers, metadata errors like that are non-fatal, and do not crash the process.

[0xffff-zhiyan]

@kevin-wu24 Thanks for the PR. It's really a great find! The fix LGTM.
I think this bug also reveals we're actually missing a check that MetadataLoader only gets committed records Could we add a cheap check in MetadataLoader#handleCommit using the existing highWaterMarkAccessor, e.g:

OptionalLong hwm = highWaterMarkAccessor.get();
if (hwm.isPresent() && batch.lastOffset() >= hwm.getAsLong()) {
    throw new IllegalStateException(...);
}

With this fix in theory it'd never fire anyway, just a tripwire for future regressions. Or do you think it's overkill?

Hi @0xffff-zhiyan,

Specifically for the issue we observed, the SnapshotEmitter is the one who logged a metadata error that it was trying to write a snapshot past the HWM, so we did become aware of it via an exception being thrown. For brokers, metadata errors like that are non-fatal, and do not crash the process.

@kevin-wu24 Yeah, SnapshotEmitter is what made us aware. But it only reported "snapshot past the HWM," which could come from several things (offset miscalculation, HWM regression, …), so it was hard to tell where the problem actually was.

So I'm wondering is there a check we could add that points directly at the root cause e.g. confirm the loader actually received records past the HWM

[kevin-wu24]

So I'm wondering is there a check we could add that points directly at the root cause e.g. confirm the loader actually received records past the HWM

I think it is probably okay to make this change, but I would like to do it in a follow-up PR. Mainly because it looks like the MetadataLoader is also used in SnapshotFileReader by the kafka-metadata-shell.

[junrao]

@kevin-wu24 : Thanks for the updated PR. LGTM. I will let @jsancio take another look.

[chia7712]

Should we add a log when the slice ends in a partial batch? It should not happen in raft, so WARN level seems suitable.

[kevin-wu24]

I agree. I think we should log at ERROR because this indicates an issue with raft/the log, or that the records have been corrupted.

[chia7712]

{@link #sizeInBytes()}} -> {@link #sizeInBytes()}

[junrao]

@kevin-wu24 : Thanks for the updated PR. One more comment.

[junrao]

If this is not expected to happen, show we throw am exception instead of logging?

[kevin-wu24]

I believe it is better to throw an exception. The callers of this should all be at batch aligned offsets, since the HWM, LEO, and any snapshot offset should be batch aligned.

From my reading of the code, the HWM, LEO, and snapshot offsets being batch aligned are invariants. Therefore, I think it is appropriate to throw an unhandled IllegalStateException. What do you think?

cc @jsancio

[junrao]

@kevin-wu24 : Thanks for the updated PR. A couple of more comments.

[junrao]

Claude suggested the following.

2. testFileRecordsSliceEndingInPartialBatch doesn't check what it sets up. It builds actualBatches and allocatedBuffers but asserts on neither. After the try block, add:



assertEquals(List.of(firstBatch), actualBatches);   // the complete batch is still returned before the throw
assertEquals(Set.of(), allocatedBuffers);           // buffers are released on the exception path

[junrao]

Claude suggests the following.

1. An incomplete batch on the grow path still ends iteration without an error (RecordsIterator.java:~143-161). The new branch only throws once everything has been read (bytesRead >= size). If the slice still has unread bytes but the declared batch is bigger than what's left, the code grows the buffer, reads the rest, and returns records with no complete batch. hasNext() is then false, with no exception. The reviewer's example:

batchSize = 17 (the default in createIterator), and a slice holding the first 30 bytes of a 100-byte batch.

Read 1 gets 17 bytes, firstBatchSize = 100, and bytesRead = 17 < 30.
The buffer grows to 100 bytes, read 2 gets 13 bytes, and iteration ends with nothing returned and no error.

3. Say what went wrong in the exception message. If this fires while loading a KRaft snapshot or a committed slice, "Unexpected partial batch of records" says nothing about where. Including details costs nothing:

throw new IllegalStateException(String.format(
    "Records %s end in a partial batch: %d bytes buffered at position %d, batch size %s, %d of %d bytes read",
    fileRecords, buffer.remaining(), bytesRead - buffer.remaining(),
    firstBatchSize == null ? "unknown" : firstBatchSize, bytesRead, fileRecords.sizeInBytes()));
bytesRead

RecordsIterator supports both MemoryRecords and FileRecords. MemoryRecords input still ignores a trailing incomplete batch through batchIterator(), while FileRecords now throws. Could we make the behavior consistent and also document this in the javadoc?

[kevin-wu24]

RecordsIterator supports both MemoryRecords and FileRecords. MemoryRecords input still ignores a trailing incomplete batch through batchIterator(), while FileRecords now throws. Could we make the behavior consistent and also document this in the javadoc?

The MemoryRecords branch in nextBatches is only hit during unit tests because those use MockLog. But yeah, I can make the change.

Also moved the validation on batch alignment in createMemoryRecords into readFileRecords.

[jsancio]

I am not sure that we should throw an exception if FileRecords doesn't have enough bytes to complete a batch. This can happen when concurrently reading uncommitted data to the LEO and another thread appends to the LEO. Or when the slice doesn't end in a batch boundary.

I think we need to distinguish the two cases:

1. The buffer is not big enough to read the entire batch - create a new batch that is big enough to read the batch from the file records.
2. The file record (slice or not) doesn't contain an entire batch - return to the user that there are no more batches.

I guess the issue is that if there are concurrent read and write from the LEO RecordsIterator::hasNext may return false in one call to later return true. Is this what we are trying to fix? I took a look at the Java documentation for Iterator and the documentation does forbid this. I would assume that concurrent iterator need to be able to return false in one call and true in a future call.

What do we think?

[kevin-wu24]

This can happen when concurrently reading uncommitted data to the LEO and another thread appends to the LEO.

This case should never happen in KRaft right? KRaftControlRecordStateMachine#maybeLoadLog is the only uncommitted reader, and it runs on the raft poll/driver thread, which is also responsible for appending to the LEO right?

Or when the slice doesn't end in a batch boundary

I agree that this case leads to a FileRecords slice which does not have enough bytes to complete a batch. However, the only "correct/valid" FileRecords slices in the KRaft/metdata code should be ending at the HWM and the LEO, which are batch aligned. This is the motivation for throwing, since it is not something that is expected.

[kevin-wu24]

Discussed this issue offline with @jsancio. The main takeaways were:

• Enforcing batch alignment of records at the RecordsIterator level via throwing an exception is mainly motivated by assumptions around the current caller behaviors, which are raft/metadata module codes. If we want to enforce this without assuming things about the callers of this iterator, batch alignment should instead be enforced at the underlying records level. Meaning, if the underlyingFileRecords type does not currently enforce batch alignment of the end offset in either the slice or the mutable case, why should RecordsIterator? If the raft module wants to enforce this, perhaps it can use some type like BatchAlignedBoundedRecords, which itself enforces the records are an immutable slice, and are batch aligned. This way, the RecordsIterator type can remain generic and reusable. The case specifically mentioned was that the one caller of RecordsIterator on uncommitted data in KRaftControlRecordStateMachine could have an implementation that holds the same RecordsIterator object and waits for new bytes (which may not be batch aligned) to come in.
• Allowing for a partial batch at the end of the requested slice is okay because memoryRecords.batchIterator()#hasNext() == false, and so will the actual RecordsIterator#hasNext() == false. This means the overall Iterator contract is not violated. Meaning, a partial batch at the end of a FileRecords slice do not make it impossible for hasNext() to return false, so this is not an invalid state.

I am going to revert the exception code and tests because of the above. What do you think @junrao?

[jsancio]

Thanks for the fix @kevin-wu24

[jsancio]

Can we enclose this statement using { }:

        if (position < 0 || position > currentSize) {
            throw new IllegalArgumentException("Read position " + position + " is outside " + this);
        }

[jsancio]

I am not sure that we should throw an exception if FileRecords doesn't have enough bytes to complete a batch. This can happen when concurrently reading uncommitted data to the LEO and another thread appends to the LEO. Or when the slice doesn't end in a batch boundary.

I think we need to distinguish the two cases:

1. The buffer is not big enough to read the entire batch - create a new batch that is big enough to read the batch from the file records.
2. The file record (slice or not) doesn't contain an entire batch - return to the user that there are no more batches.

I guess the issue is that if there are concurrent read and write from the LEO RecordsIterator::hasNext may return false in one call to later return true. Is this what we are trying to fix? I took a look at the Java documentation for Iterator and the documentation does forbid this. I would assume that concurrent iterator need to be able to return false in one call and true in a future call.

What do we think?