removed partition_max_bytes from delayed share fetch

Author: adixitconfluent

clients/src/main/java/org/apache/kafka/common/requests/FetchRequest.java

@@ -59,9 +59,18 @@ public static final class PartitionData {
         public final Uuid topicId;
         public final long fetchOffset;
         public final long logStartOffset;
-        public final int maxBytes;
         public final Optional<Integer> currentLeaderEpoch;
         public final Optional<Integer> lastFetchedEpoch;
+        public int maxBytes;
+
+        public PartitionData(
+            Uuid topicId,
+            long fetchOffset,
+            long logStartOffset,
+            Optional<Integer> currentLeaderEpoch
+        ) {
+            this(topicId, fetchOffset, logStartOffset, 0, currentLeaderEpoch, Optional.empty());
+        }
 
         public PartitionData(
             Uuid topicId,
@@ -89,6 +98,10 @@ public PartitionData(
             this.lastFetchedEpoch = lastFetchedEpoch;
         }
 
+        public void updateMaxBytes(int maxBytes) {
+            this.maxBytes = maxBytes;
+        }
+
         @Override
         public boolean equals(Object o) {
             if (this == o) return true;

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

@@ -41,6 +41,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.Lock;
 import java.util.stream.Collectors;
 
@@ -124,7 +125,7 @@ private void completeShareFetchRequest(LinkedHashMap<TopicIdPartition, FetchRequ
         try {
             LinkedHashMap<TopicIdPartition, LogReadResult> responseData;
             if (partitionsAlreadyFetched.isEmpty())
-                responseData = readFromLog(topicPartitionData);
+                responseData = readFromLog(topicPartitionData, shareFetch.fetchParams().maxBytes / topicPartitionData.size());
             else
                 // There shouldn't be a case when we have a partitionsAlreadyFetched value here and this variable is getting
                 // updated in a different tryComplete thread.
@@ -206,7 +207,6 @@ LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> acquirablePartitions
         LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData = new LinkedHashMap<>();
 
         sharePartitions.forEach((topicIdPartition, sharePartition) -> {
-            int partitionMaxBytes = shareFetch.partitionMaxBytes().getOrDefault(topicIdPartition, 0);
             // Add the share partition to the list of partitions to be fetched only if we can
             // acquire the fetch lock on it.
             if (sharePartition.maybeAcquireFetchLock()) {
@@ -219,7 +219,6 @@ LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> acquirablePartitions
                                 topicIdPartition.topicId(),
                                 sharePartition.nextFetchOffset(),
                                 0,
-                                partitionMaxBytes,
                                 Optional.empty()
                             )
                         );
@@ -250,7 +249,7 @@ private LinkedHashMap<TopicIdPartition, LogReadResult> maybeReadFromLog(LinkedHa
             return new LinkedHashMap<>();
         }
         // We fetch data from replica manager corresponding to the topic partitions that have missing fetch offset metadata.
-        return readFromLog(partitionsNotMatchingFetchOffsetMetadata);
+        return readFromLog(partitionsNotMatchingFetchOffsetMetadata, shareFetch.fetchParams().maxBytes / topicPartitionData.size());
     }
 
     private void maybeUpdateFetchOffsetMetadata(
@@ -311,7 +310,7 @@ private boolean isMinBytesSatisfied(LinkedHashMap<TopicIdPartition, FetchRequest
                     return true;
                 } else if (fetchOffsetMetadata.onSameSegment(endOffsetMetadata)) {
                     // we take the partition fetch size as upper bound when accumulating the bytes.
-                    long bytesAvailable = Math.min(endOffsetMetadata.positionDiff(fetchOffsetMetadata), partitionData.maxBytes);
+                    long bytesAvailable = Math.min(endOffsetMetadata.positionDiff(fetchOffsetMetadata), shareFetch.fetchParams().maxBytes / topicPartitionData.size());
                     accumulatedSize += bytesAvailable;
                 }
             }
@@ -334,12 +333,13 @@ else if (isolationType == FetchIsolation.HIGH_WATERMARK)
 
     }
 
-    private LinkedHashMap<TopicIdPartition, LogReadResult> readFromLog(LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData) {
+    private LinkedHashMap<TopicIdPartition, LogReadResult> readFromLog(LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData, int partitionMaxBytes) {
         // Filter if there already exists any erroneous topic partition.
         Set<TopicIdPartition> partitionsToFetch = shareFetch.filterErroneousTopicPartitions(topicPartitionData.keySet());
         if (partitionsToFetch.isEmpty()) {
             return new LinkedHashMap<>();
         }
+        topicPartitionData.values().forEach(partitionData -> partitionData.updateMaxBytes(partitionMaxBytes));
 
         Seq<Tuple2<TopicIdPartition, LogReadResult>> responseLogResult = replicaManager.readFromLog(
             shareFetch.fetchParams(),
@@ -392,15 +392,21 @@ private void handleFetchException(
     LinkedHashMap<TopicIdPartition, LogReadResult> combineLogReadResponse(LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData,
                                                                 LinkedHashMap<TopicIdPartition, LogReadResult> existingFetchedData) {
         LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> missingLogReadTopicPartitions = new LinkedHashMap<>();
+        AtomicInteger totalPartitionMaxBytesUsed = new AtomicInteger();
         topicPartitionData.forEach((topicIdPartition, partitionData) -> {
             if (!existingFetchedData.containsKey(topicIdPartition)) {
                 missingLogReadTopicPartitions.put(topicIdPartition, partitionData);
+            } else {
+                totalPartitionMaxBytesUsed.addAndGet(partitionData.maxBytes);
             }
         });
         if (missingLogReadTopicPartitions.isEmpty()) {
             return existingFetchedData;
         }
-        LinkedHashMap<TopicIdPartition, LogReadResult> missingTopicPartitionsLogReadResponse = readFromLog(missingLogReadTopicPartitions);
+        LinkedHashMap<TopicIdPartition, LogReadResult> missingTopicPartitionsLogReadResponse = readFromLog(
+            missingLogReadTopicPartitions,
+            (shareFetch.fetchParams().maxBytes - totalPartitionMaxBytesUsed.get()) / missingLogReadTopicPartitions.size()
+            );
         missingTopicPartitionsLogReadResponse.putAll(existingFetchedData);
         return missingTopicPartitionsLogReadResponse;
     }