Revert "KAFKA-18404: Remove partitionMaxBytes usage from DelayedShareFetch (#17870)" (#18643)

This reverts commit b021b51.

Reviewers: Ismael Juma <[email protected]>, Andrew Schofield <[email protected]>

Author: apoorvmittal10

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

@@ -27,7 +27,6 @@
 import org.apache.kafka.server.purgatory.DelayedOperation;
 import org.apache.kafka.server.share.SharePartitionKey;
 import org.apache.kafka.server.share.fetch.DelayedShareFetchGroupKey;
-import org.apache.kafka.server.share.fetch.PartitionMaxBytesStrategy;
 import org.apache.kafka.server.share.fetch.ShareFetch;
 import org.apache.kafka.server.storage.log.FetchIsolation;
 import org.apache.kafka.server.storage.log.FetchPartitionData;
@@ -61,35 +60,24 @@ public class DelayedShareFetch extends DelayedOperation {
     private final ShareFetch shareFetch;
     private final ReplicaManager replicaManager;
     private final BiConsumer<SharePartitionKey, Throwable> exceptionHandler;
-    private final PartitionMaxBytesStrategy partitionMaxBytesStrategy;
     // The topic partitions that need to be completed for the share fetch request are given by sharePartitions.
     // sharePartitions is a subset of shareFetchData. The order of insertion/deletion of entries in sharePartitions is important.
     private final LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions;
-    private LinkedHashMap<TopicIdPartition, Long> partitionsAcquired;
+    private LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> partitionsAcquired;
     private LinkedHashMap<TopicIdPartition, LogReadResult> partitionsAlreadyFetched;
 
     DelayedShareFetch(
             ShareFetch shareFetch,
             ReplicaManager replicaManager,
             BiConsumer<SharePartitionKey, Throwable> exceptionHandler,
             LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions) {
-        this(shareFetch, replicaManager, exceptionHandler, sharePartitions, PartitionMaxBytesStrategy.type(PartitionMaxBytesStrategy.StrategyType.UNIFORM));
-    }
-
-    DelayedShareFetch(
-        ShareFetch shareFetch,
-        ReplicaManager replicaManager,
-        BiConsumer<SharePartitionKey, Throwable> exceptionHandler,
-        LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions,
-        PartitionMaxBytesStrategy partitionMaxBytesStrategy) {
         super(shareFetch.fetchParams().maxWaitMs, Optional.empty());
         this.shareFetch = shareFetch;
         this.replicaManager = replicaManager;
         this.partitionsAcquired = new LinkedHashMap<>();
         this.partitionsAlreadyFetched = new LinkedHashMap<>();
         this.exceptionHandler = exceptionHandler;
         this.sharePartitions = sharePartitions;
-        this.partitionMaxBytesStrategy = partitionMaxBytesStrategy;
     }
 
     @Override
@@ -111,7 +99,7 @@ public void onComplete() {
             partitionsAcquired.keySet());
 
         try {
-            LinkedHashMap<TopicIdPartition, Long> topicPartitionData;
+            LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData;
             // tryComplete did not invoke forceComplete, so we need to check if we have any partitions to fetch.
             if (partitionsAcquired.isEmpty())
                 topicPartitionData = acquirablePartitions();
@@ -133,13 +121,11 @@ public void onComplete() {
         }
     }
 
-    private void completeShareFetchRequest(LinkedHashMap<TopicIdPartition, Long> topicPartitionData) {
+    private void completeShareFetchRequest(LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData) {
         try {
             LinkedHashMap<TopicIdPartition, LogReadResult> responseData;
             if (partitionsAlreadyFetched.isEmpty())
-                responseData = readFromLog(
-                    topicPartitionData,
-                    partitionMaxBytesStrategy.maxBytes(shareFetch.fetchParams().maxBytes, topicPartitionData.keySet(), topicPartitionData.size()));
+                responseData = readFromLog(topicPartitionData);
             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.
@@ -172,7 +158,7 @@ private void completeShareFetchRequest(LinkedHashMap<TopicIdPartition, Long> top
      */
     @Override
     public boolean tryComplete() {
-        LinkedHashMap<TopicIdPartition, Long> topicPartitionData = acquirablePartitions();
+        LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData = acquirablePartitions();
 
         try {
             if (!topicPartitionData.isEmpty()) {
@@ -181,7 +167,7 @@ public boolean tryComplete() {
                 // those topic partitions.
                 LinkedHashMap<TopicIdPartition, LogReadResult> replicaManagerReadResponse = maybeReadFromLog(topicPartitionData);
                 maybeUpdateFetchOffsetMetadata(topicPartitionData, replicaManagerReadResponse);
-                if (anyPartitionHasLogReadError(replicaManagerReadResponse) || isMinBytesSatisfied(topicPartitionData, partitionMaxBytesStrategy.maxBytes(shareFetch.fetchParams().maxBytes, topicPartitionData.keySet(), topicPartitionData.size()))) {
+                if (anyPartitionHasLogReadError(replicaManagerReadResponse) || isMinBytesSatisfied(topicPartitionData)) {
                     partitionsAcquired = topicPartitionData;
                     partitionsAlreadyFetched = replicaManagerReadResponse;
                     boolean completedByMe = forceComplete();
@@ -216,18 +202,28 @@ public boolean tryComplete() {
      * Prepare fetch request structure for partitions in the share fetch request for which we can acquire records.
      */
     // Visible for testing
-    LinkedHashMap<TopicIdPartition, Long> acquirablePartitions() {
+    LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> acquirablePartitions() {
         // Initialize the topic partitions for which the fetch should be attempted.
-        LinkedHashMap<TopicIdPartition, Long> topicPartitionData = new LinkedHashMap<>();
+        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()) {
                 try {
                     // If the share partition is already at capacity, we should not attempt to fetch.
                     if (sharePartition.canAcquireRecords()) {
-                        topicPartitionData.put(topicIdPartition, sharePartition.nextFetchOffset());
+                        topicPartitionData.put(
+                            topicIdPartition,
+                            new FetchRequest.PartitionData(
+                                topicIdPartition.topicId(),
+                                sharePartition.nextFetchOffset(),
+                                0,
+                                partitionMaxBytes,
+                                Optional.empty()
+                            )
+                        );
                     } else {
                         sharePartition.releaseFetchLock();
                         log.trace("Record lock partition limit exceeded for SharePartition {}, " +
@@ -243,28 +239,24 @@ LinkedHashMap<TopicIdPartition, Long> acquirablePartitions() {
         return topicPartitionData;
     }
 
-    private LinkedHashMap<TopicIdPartition, LogReadResult> maybeReadFromLog(LinkedHashMap<TopicIdPartition, Long> topicPartitionData) {
-        LinkedHashMap<TopicIdPartition, Long> partitionsNotMatchingFetchOffsetMetadata = new LinkedHashMap<>();
-        topicPartitionData.forEach((topicIdPartition, fetchOffset) -> {
+    private LinkedHashMap<TopicIdPartition, LogReadResult> maybeReadFromLog(LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData) {
+        LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> partitionsNotMatchingFetchOffsetMetadata = new LinkedHashMap<>();
+        topicPartitionData.forEach((topicIdPartition, partitionData) -> {
             SharePartition sharePartition = sharePartitions.get(topicIdPartition);
-            if (sharePartition.fetchOffsetMetadata(fetchOffset).isEmpty()) {
-                partitionsNotMatchingFetchOffsetMetadata.put(topicIdPartition, fetchOffset);
+            if (sharePartition.fetchOffsetMetadata(partitionData.fetchOffset).isEmpty()) {
+                partitionsNotMatchingFetchOffsetMetadata.put(topicIdPartition, partitionData);
             }
         });
         if (partitionsNotMatchingFetchOffsetMetadata.isEmpty()) {
             return new LinkedHashMap<>();
         }
         // We fetch data from replica manager corresponding to the topic partitions that have missing fetch offset metadata.
-        // Although we are fetching partition max bytes for partitionsNotMatchingFetchOffsetMetadata,
-        // we will take acquired partitions size = topicPartitionData.size() because we do not want to let the
-        // leftover partitions to starve which will be fetched later.
-        return readFromLog(
-            partitionsNotMatchingFetchOffsetMetadata,
-            partitionMaxBytesStrategy.maxBytes(shareFetch.fetchParams().maxBytes, partitionsNotMatchingFetchOffsetMetadata.keySet(), topicPartitionData.size()));
+        return readFromLog(partitionsNotMatchingFetchOffsetMetadata);
     }
 
-    private void maybeUpdateFetchOffsetMetadata(LinkedHashMap<TopicIdPartition, Long> topicPartitionData,
-                                                LinkedHashMap<TopicIdPartition, LogReadResult> replicaManagerReadResponseData) {
+    private void maybeUpdateFetchOffsetMetadata(
+        LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData,
+        LinkedHashMap<TopicIdPartition, LogReadResult> replicaManagerReadResponseData) {
         for (Map.Entry<TopicIdPartition, LogReadResult> entry : replicaManagerReadResponseData.entrySet()) {
             TopicIdPartition topicIdPartition = entry.getKey();
             SharePartition sharePartition = sharePartitions.get(topicIdPartition);
@@ -275,18 +267,17 @@ private void maybeUpdateFetchOffsetMetadata(LinkedHashMap<TopicIdPartition, Long
                 continue;
             }
             sharePartition.updateFetchOffsetMetadata(
-                topicPartitionData.get(topicIdPartition),
+                topicPartitionData.get(topicIdPartition).fetchOffset,
                 replicaManagerLogReadResult.info().fetchOffsetMetadata);
         }
     }
 
     // minByes estimation currently assumes the common case where all fetched data is acquirable.
-    private boolean isMinBytesSatisfied(LinkedHashMap<TopicIdPartition, Long> topicPartitionData,
-                                        LinkedHashMap<TopicIdPartition, Integer> partitionMaxBytes) {
+    private boolean isMinBytesSatisfied(LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData) {
         long accumulatedSize = 0;
-        for (Map.Entry<TopicIdPartition, Long> entry : topicPartitionData.entrySet()) {
+        for (Map.Entry<TopicIdPartition, FetchRequest.PartitionData> entry : topicPartitionData.entrySet()) {
             TopicIdPartition topicIdPartition = entry.getKey();
-            long fetchOffset = entry.getValue();
+            FetchRequest.PartitionData partitionData = entry.getValue();
 
             LogOffsetMetadata endOffsetMetadata;
             try {
@@ -303,7 +294,7 @@ private boolean isMinBytesSatisfied(LinkedHashMap<TopicIdPartition, Long> topicP
 
             SharePartition sharePartition = sharePartitions.get(topicIdPartition);
 
-            Optional<LogOffsetMetadata> optionalFetchOffsetMetadata = sharePartition.fetchOffsetMetadata(fetchOffset);
+            Optional<LogOffsetMetadata> optionalFetchOffsetMetadata = sharePartition.fetchOffsetMetadata(partitionData.fetchOffset);
             if (optionalFetchOffsetMetadata.isEmpty() || optionalFetchOffsetMetadata.get() == LogOffsetMetadata.UNKNOWN_OFFSET_METADATA)
                 continue;
             LogOffsetMetadata fetchOffsetMetadata = optionalFetchOffsetMetadata.get();
@@ -321,7 +312,7 @@ private boolean isMinBytesSatisfied(LinkedHashMap<TopicIdPartition, Long> topicP
                     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), partitionMaxBytes.get(topicIdPartition));
+                    long bytesAvailable = Math.min(endOffsetMetadata.positionDiff(fetchOffsetMetadata), partitionData.maxBytes);
                     accumulatedSize += bytesAvailable;
                 }
             }
@@ -344,25 +335,13 @@ else if (isolationType == FetchIsolation.HIGH_WATERMARK)
 
     }
 
-    private LinkedHashMap<TopicIdPartition, LogReadResult> readFromLog(LinkedHashMap<TopicIdPartition, Long> topicPartitionFetchOffsets,
-                                                                       LinkedHashMap<TopicIdPartition, Integer> partitionMaxBytes) {
+    private LinkedHashMap<TopicIdPartition, LogReadResult> readFromLog(LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData) {
         // Filter if there already exists any erroneous topic partition.
-        Set<TopicIdPartition> partitionsToFetch = shareFetch.filterErroneousTopicPartitions(topicPartitionFetchOffsets.keySet());
+        Set<TopicIdPartition> partitionsToFetch = shareFetch.filterErroneousTopicPartitions(topicPartitionData.keySet());
         if (partitionsToFetch.isEmpty()) {
             return new LinkedHashMap<>();
         }
 
-        LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData = new LinkedHashMap<>();
-
-        topicPartitionFetchOffsets.forEach((topicIdPartition, fetchOffset) -> topicPartitionData.put(topicIdPartition,
-            new FetchRequest.PartitionData(
-                topicIdPartition.topicId(),
-                fetchOffset,
-                0,
-                partitionMaxBytes.get(topicIdPartition),
-                Optional.empty())
-        ));
-
         Seq<Tuple2<TopicIdPartition, LogReadResult>> responseLogResult = replicaManager.readFromLog(
             shareFetch.fetchParams(),
             CollectionConverters.asScala(
@@ -411,21 +390,18 @@ private void handleFetchException(
     }
 
     // Visible for testing.
-    LinkedHashMap<TopicIdPartition, LogReadResult> combineLogReadResponse(LinkedHashMap<TopicIdPartition, Long> topicPartitionData,
-                                                                          LinkedHashMap<TopicIdPartition, LogReadResult> existingFetchedData) {
-        LinkedHashMap<TopicIdPartition, Long> missingLogReadTopicPartitions = new LinkedHashMap<>();
-        topicPartitionData.forEach((topicIdPartition, fetchOffset) -> {
+    LinkedHashMap<TopicIdPartition, LogReadResult> combineLogReadResponse(LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> topicPartitionData,
+                                                                LinkedHashMap<TopicIdPartition, LogReadResult> existingFetchedData) {
+        LinkedHashMap<TopicIdPartition, FetchRequest.PartitionData> missingLogReadTopicPartitions = new LinkedHashMap<>();
+        topicPartitionData.forEach((topicIdPartition, partitionData) -> {
             if (!existingFetchedData.containsKey(topicIdPartition)) {
-                missingLogReadTopicPartitions.put(topicIdPartition, fetchOffset);
+                missingLogReadTopicPartitions.put(topicIdPartition, partitionData);
             }
         });
         if (missingLogReadTopicPartitions.isEmpty()) {
             return existingFetchedData;
         }
-
-        LinkedHashMap<TopicIdPartition, LogReadResult> missingTopicPartitionsLogReadResponse = readFromLog(
-            missingLogReadTopicPartitions,
-            partitionMaxBytesStrategy.maxBytes(shareFetch.fetchParams().maxBytes, missingLogReadTopicPartitions.keySet(), topicPartitionData.size()));
+        LinkedHashMap<TopicIdPartition, LogReadResult> missingTopicPartitionsLogReadResponse = readFromLog(missingLogReadTopicPartitions);
         missingTopicPartitionsLogReadResponse.putAll(existingFetchedData);
         return missingTopicPartitionsLogReadResponse;
     }