MINOR: Skip further reads from log for share fetches if max records reached

core/src/main/java/kafka/server/share/DelayedShareFetch.java

@@ -26,6 +26,8 @@
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.apache.kafka.common.message.ShareFetchResponseData;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.internal.RecordBatch;
+import org.apache.kafka.common.record.internal.Records;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.raft.errors.NotLeaderException;
 import org.apache.kafka.server.metrics.KafkaMetricsGroup;
@@ -840,6 +842,14 @@ CompletableFuture<LinkedHashMap<TopicIdPartition, LogReadResult>> combineLogRead
             return CompletableFuture.completedFuture(existingFetchedData);
         }
 
+        // The already fetched records are handed out first, so only the remainder of the record budget
+        // is available to the missing partitions. If nothing remains then nothing from them could be
+        // acquired, so skip the read altogether; the partitions are simply absent from the response and
+        // their locks are released along with the rest of the acquired partitions.
+        if (remainingRecordBudget(shareFetch.maxFetchRecords(), topicPartitionData, existingFetchedData) <= 0) {
+            return CompletableFuture.completedFuture(existingFetchedData);
+        }
+
         return readFromLog(
             missingLogReadTopicPartitions,
             partitionMaxBytesStrategy.maxBytes(shareFetch.fetchParams().maxBytes,
@@ -850,6 +860,61 @@ CompletableFuture<LinkedHashMap<TopicIdPartition, LogReadResult>> combineLogRead
         });
     }
 
+    /**
+     * Returns the record budget left for further reads once the already fetched data has been handed out,
+     * or 0 if the budget is exhausted. Records are counted by offset span, from the fetch offset to the last
+     * offset of the last batch, which is the same approximation {@link SharePartition#acquire} makes when it
+     * bounds the acquired records by max fetch records. The span starts at the fetch offset rather than at
+     * the base offset of the first batch because the fetched data starts at a batch boundary, and the records
+     * of that batch prior to the fetch offset are already in-flight or acknowledged, hence are not handed out
+     * again and must not consume the budget.
+     * <ul>
+     *     <li>{@link RecordBatch#lastOffset()} is called at most once per partition, on the last batch, rather
+     *     than on every batch. For file backed records the base offset comes from the bytes already read while
+     *     iterating, whereas the last offset may require loading the batch header from the file channel.</li>
+     *     <li>Iteration stops as soon as the base offsets alone prove the budget is spent. Base offsets are
+     *     strictly increasing, so the offsets between the start of the span and a batch's base offset are
+     *     already accounted for and the batch itself holds at least one more record.</li>
+     * </ul>
+     *
+     * @param maxFetchRecords The record budget of the share fetch request.
+     * @param fetchOffsets The fetch offsets of the acquired topic partitions.
+     * @param fetchedData The already fetched data, keyed by topic partition.
+     * @return The remaining record budget, or 0 if none remains.
+     */
+    private static int remainingRecordBudget(int maxFetchRecords,
+                                             Map<TopicIdPartition, Long> fetchOffsets,
+                                             Map<TopicIdPartition, LogReadResult> fetchedData) {
+        long remaining = maxFetchRecords;
+        for (Map.Entry<TopicIdPartition, LogReadResult> entry : fetchedData.entrySet()) {
+            FetchDataInfo fetchDataInfo = entry.getValue().info();
+            if (fetchDataInfo == null || fetchDataInfo.records == null) {
+                continue;
+            }
+
+            Records records = fetchDataInfo.records;
+            // The span starts at the fetch offset, or at the base offset of the first batch if that is
+            // higher or the fetch offset is unknown.
+            long startOffset = fetchOffsets.getOrDefault(entry.getKey(), -1L);
+            RecordBatch lastBatch = null;
+            for (RecordBatch batch : records.batches()) {
+                if (lastBatch == null) {
+                    startOffset = Math.max(startOffset, batch.baseOffset());
+                } else if (batch.baseOffset() - startOffset >= remaining) {
+                    return 0;
+                }
+                lastBatch = batch;
+            }
+            if (lastBatch != null) {
+                remaining -= Math.max(0, lastBatch.lastOffset() - startOffset + 1);
+                if (remaining <= 0) {
+                    return 0;
+                }
+            }
+        }
+        return (int) remaining;
+    }
+
     // Visible for testing.
     void releasePartitionLocks(Set<TopicIdPartition> topicIdPartitions) {
         topicIdPartitions.forEach(tp -> {

core/src/test/java/kafka/server/share/DelayedShareFetchTest.java

@@ -65,6 +65,7 @@
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -74,6 +75,7 @@
 import java.util.OptionalInt;
 import java.util.OptionalLong;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.concurrent.RejectedExecutionException;
@@ -93,6 +95,7 @@
 import static kafka.server.share.SharePartitionManagerTest.buildEmptyLogReadResult;
 import static kafka.server.share.SharePartitionManagerTest.buildLogReadResult;
 import static kafka.server.share.SharePartitionManagerTest.mockReplicaManagerDelayedShareFetch;
+import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.createFileRecords;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -869,6 +872,116 @@ public void testCombineLogReadResponse() {
         assertEquals(combinedLogReadResponse.get(tp1), logReadResponse.get(tp1));
     }
 
+    @Test
+    public void testCombineLogReadResponseSkipsReadWhenRecordBudgetExhausted() throws IOException {
+        String groupId = "grp";
+        ReplicaManager replicaManager = mock(ReplicaManager.class);
+        TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
+        TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
+        TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 2));
+
+        LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
+        sharePartitions.put(tp0, mock(SharePartition.class));
+        sharePartitions.put(tp1, mock(SharePartition.class));
+        sharePartitions.put(tp2, mock(SharePartition.class));
+
+        ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, Uuid.randomUuid().toString(),
+            new CompletableFuture<>(), List.of(tp0, tp1, tp2), BATCH_OPTIMIZED, BATCH_SIZE, MAX_FETCH_RECORDS,
+            BROKER_TOPIC_STATS);
+
+        DelayedShareFetch delayedShareFetch = DelayedShareFetchBuilder.builder()
+            .withShareFetchData(shareFetch)
+            .withReplicaManagerLogReader(replicaManager)
+            .withSharePartitions(sharePartitions)
+            .withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp2)))
+            .build();
+        doAnswer(invocation -> buildLogReadResult(List.of(tp2))).when(replicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
+
+        LinkedHashMap<TopicIdPartition, Long> topicPartitionData = new LinkedHashMap<>();
+        topicPartitionData.put(tp0, 0L);
+        topicPartitionData.put(tp1, 0L);
+        topicPartitionData.put(tp2, 0L);
+
+        // Case 1 - the already fetched data holds 99 records in 2 batches, fewer than max fetch records,
+        // hence the missing partition is read.
+        LinkedHashMap<TopicIdPartition, LogReadResult> logReadResponse = buildLogReadResultMap(List.of(tp0),
+            createFileRecords(new TreeMap<>(Map.of(5L, 50, 55L, 49))));
+        logReadResponse.putAll(buildLogReadResultMap(List.of(tp1), MemoryRecords.EMPTY));
+        LinkedHashMap<TopicIdPartition, LogReadResult> combinedLogReadResponse = delayedShareFetch.combineLogReadResponse(topicPartitionData, logReadResponse).join();
+        assertEquals(topicPartitionData.keySet(), combinedLogReadResponse.keySet());
+        Mockito.verify(replicaManager, times(1)).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
+
+        // Case 2 - a single already fetched partition holds exactly max fetch records in 2 batches, hence
+        // nothing could be acquired from the missing partition and the read is skipped.
+        Mockito.clearInvocations(replicaManager);
+        logReadResponse = buildLogReadResultMap(List.of(tp0), createFileRecords(new TreeMap<>(Map.of(0L, 50, 50L, 50))));
+        logReadResponse.putAll(buildLogReadResultMap(List.of(tp1), MemoryRecords.EMPTY));
+        combinedLogReadResponse = delayedShareFetch.combineLogReadResponse(topicPartitionData, logReadResponse).join();
+        assertEquals(Set.of(tp0, tp1), combinedLogReadResponse.keySet());
+        Mockito.verify(replicaManager, never()).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
+
+        // Case 3 - the already fetched partitions together hold max fetch records, hence the read is skipped.
+        logReadResponse = buildLogReadResultMap(List.of(tp0), createFileRecords(new TreeMap<>(Map.of(10L, 30, 40L, 30))));
+        logReadResponse.putAll(buildLogReadResultMap(List.of(tp1), createFileRecords(new TreeMap<>(Map.of(0L, 40)))));
+        combinedLogReadResponse = delayedShareFetch.combineLogReadResponse(topicPartitionData, logReadResponse).join();
+        assertEquals(Set.of(tp0, tp1), combinedLogReadResponse.keySet());
+        Mockito.verify(replicaManager, never()).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
+
+        // Case 4 - the already fetched data spans well past max fetch records in 4 batches. The base offsets
+        // alone prove the budget is spent, so the iteration stops before the last batch and the read is skipped.
+        logReadResponse = buildLogReadResultMap(List.of(tp0), createFileRecords(new TreeMap<>(Map.of(5L, 50, 55L, 50, 105L, 50, 155L, 50))));
+        combinedLogReadResponse = delayedShareFetch.combineLogReadResponse(topicPartitionData, logReadResponse).join();
+        assertEquals(Set.of(tp0), combinedLogReadResponse.keySet());
+        Mockito.verify(replicaManager, never()).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
+    }
+
+    @Test
+    public void testCombineLogReadResponseRecordBudgetCountsFromFetchOffset() throws IOException {
+        String groupId = "grp";
+        ReplicaManager replicaManager = mock(ReplicaManager.class);
+        TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
+        TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
+
+        LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
+        sharePartitions.put(tp0, mock(SharePartition.class));
+        sharePartitions.put(tp1, mock(SharePartition.class));
+
+        ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, Uuid.randomUuid().toString(),
+            new CompletableFuture<>(), List.of(tp0, tp1), BATCH_OPTIMIZED, BATCH_SIZE, MAX_FETCH_RECORDS,
+            BROKER_TOPIC_STATS);
+
+        DelayedShareFetch delayedShareFetch = DelayedShareFetchBuilder.builder()
+            .withShareFetchData(shareFetch)
+            .withReplicaManagerLogReader(replicaManager)
+            .withSharePartitions(sharePartitions)
+            .withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp1)))
+            .build();
+        doAnswer(invocation -> buildLogReadResult(List.of(tp1))).when(replicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
+
+        // The fetched data starts at a batch boundary, hence the first batch can hold records prior to the
+        // fetch offset which are already in-flight or acknowledged. Those are not handed out again, so they
+        // must not consume the record budget. The data holds offsets 5 to 154 in 3 batches.
+        LinkedHashMap<TopicIdPartition, LogReadResult> logReadResponse = buildLogReadResultMap(List.of(tp0),
+            createFileRecords(new TreeMap<>(Map.of(5L, 50, 55L, 50, 105L, 50))));
+
+        // Case 1 - fetch offset inside the first batch: only offsets 60 to 154 count, so budget remains.
+        LinkedHashMap<TopicIdPartition, Long> topicPartitionData = new LinkedHashMap<>();
+        topicPartitionData.put(tp0, 60L);
+        topicPartitionData.put(tp1, 0L);
+
+        LinkedHashMap<TopicIdPartition, LogReadResult> combinedLogReadResponse = delayedShareFetch.combineLogReadResponse(topicPartitionData, logReadResponse).join();
+        assertEquals(topicPartitionData.keySet(), combinedLogReadResponse.keySet());
+        Mockito.verify(replicaManager, times(1)).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
+
+        // Case 2 - fetch offset at the base of the first batch: all 150 count, so the budget is exhausted.
+        Mockito.clearInvocations(replicaManager);
+        topicPartitionData.put(tp0, 0L);
+
+        combinedLogReadResponse = delayedShareFetch.combineLogReadResponse(topicPartitionData, logReadResponse).join();
+        assertEquals(Set.of(tp0), combinedLogReadResponse.keySet());
+        Mockito.verify(replicaManager, never()).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
+    }
+
     @Test
     public void testExceptionInMinBytesCalculation() {
         String groupId = "grp";
@@ -3346,6 +3459,16 @@ private static LinkedHashMap<TopicIdPartition, LogReadResult> buildLogReadResult
             ));
     }
 
+    private static LinkedHashMap<TopicIdPartition, LogReadResult> buildLogReadResultMap(List<TopicIdPartition> topicIdPartitions, Records records) {
+        return CollectionConverters.asJava(buildLogReadResult(topicIdPartitions, records)).stream()
+            .collect(Collectors.toMap(
+                Tuple2::_1,
+                Tuple2::_2,
+                (oldValue, newValue) -> oldValue,
+                LinkedHashMap::new
+            ));
+    }
+
     @SuppressWarnings("unchecked")
     static class DelayedShareFetchBuilder {
         private ShareFetch shareFetch = mock(ShareFetch.class);