Addressed Jun's round 3 review comments

Author: adixitconfluent

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

@@ -618,16 +618,18 @@ private boolean maybeProcessRemoteFetch(
         LinkedHashMap<TopicIdPartition, Long> topicPartitionData,
         TopicPartitionRemoteFetchInfo topicPartitionRemoteFetchInfo
     ) throws Exception {
+        Set<TopicIdPartition> nonRemoteFetchTopicPartitions = new LinkedHashSet<>();
         topicPartitionData.keySet().forEach(topicIdPartition -> {
             // topic partitions for which fetch would not be happening in this share fetch request.
             if (!topicPartitionRemoteFetchInfo.topicIdPartition().equals(topicIdPartition)) {
-                // Release acquisition lock for the topic partitions that were acquired but were not a part of remote fetch.
-                releasePartitionLocksAndAddToActionQueue(Set.of(topicIdPartition));
+                nonRemoteFetchTopicPartitions.add(topicIdPartition);
             }
         });
+        // Release fetch lock for the topic partitions that were acquired but were not a part of remote fetch and add
+        // them to the delayed actions queue.
+        releasePartitionLocksAndAddToActionQueue(nonRemoteFetchTopicPartitions);
         Optional<Exception> exceptionOpt = processRemoteFetchOrException(topicPartitionRemoteFetchInfo);
         if (exceptionOpt.isPresent()) {
-            remoteStorageFetchException = exceptionOpt;
             throw exceptionOpt.get();
         }
         // Check if remote fetch can be completed.
@@ -657,8 +659,10 @@ private Optional<Exception> processRemoteFetchOrException(
         } catch (RejectedExecutionException e) {
             // Return the error if any in scheduling the remote fetch task.
             log.warn("Unable to fetch data from remote storage", e);
+            remoteStorageFetchException = Optional.of(e);
             return Optional.of(e);
         } catch (Exception e) {
+            remoteStorageFetchException = Optional.of(e);
             return Optional.of(e);
         }
         remoteFetchOpt = Optional.of(new RemoteFetch(remoteFetchTopicIdPartition, topicPartitionRemoteFetchInfo.logReadResult(), remoteFetchTask, remoteFetchResult, remoteStorageFetchInfo));
@@ -730,7 +734,7 @@ private void releasePartitionLocksAndAddToActionQueue(Set<TopicIdPartition> topi
 
     /**
      * This function completes a share fetch request for which we have identified remoteFetch during tryComplete()
-     * Note - This function should only be called when we know that there is remote fetch in-flight/completed/expired.
+     * Note - This function should only be called when we know that there is remote fetch.
      */
     private void completeRemoteStorageShareFetchRequest() {
         LinkedHashMap<TopicIdPartition, Long> acquiredNonRemoteFetchTopicPartitionData = new LinkedHashMap<>();
@@ -842,7 +846,7 @@ private void handleExceptionInCompletingRemoteStorageShareFetchRequest(
     /**
      * Cancel the remote storage read task, if it has not been executed yet and avoid interrupting the task if it is
      * already running as it may force closing opened/cached resources as transaction index.
-     * Note - This function should only be called when we know that there is a remote fetch in-flight/expired.
+     * Note - This function should only be called when we know that there is remote fetch.
      */
     private void cancelRemoteFetchTask() {
         boolean cancelled = remoteFetchOpt.get().remoteFetchTask().cancel(false);

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

@@ -1228,8 +1228,7 @@ public void testRemoteStorageFetchTryCompleteReturnsFalse() {
         // Remote fetch object gets created for delayed share fetch object.
         assertNotNull(delayedShareFetch.remoteFetch());
         // Verify the locks are released for local log read topic partitions tp0 and tp1.
-        Mockito.verify(delayedShareFetch, times(1)).releasePartitionLocks(Set.of(tp0));
-        Mockito.verify(delayedShareFetch, times(1)).releasePartitionLocks(Set.of(tp1));
+        Mockito.verify(delayedShareFetch, times(1)).releasePartitionLocks(Set.of(tp0, tp1));
         assertTrue(delayedShareFetch.lock().tryLock());
         delayedShareFetch.lock().unlock();
     }
@@ -1459,24 +1458,19 @@ public void testRemoteStorageFetchRequestCompletionOnFutureCompletionFailure() {
     public void testRemoteStorageFetchRequestCompletionOnFutureCompletionSuccessfully() {
         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));
 
         SharePartition sp0 = mock(SharePartition.class);
-        SharePartition sp1 = mock(SharePartition.class);
 
         // sp0 is acquirable, sp1 is not acquirable.
         when(sp0.maybeAcquireFetchLock()).thenReturn(true);
-        when(sp1.maybeAcquireFetchLock()).thenReturn(false);
         when(sp0.canAcquireRecords()).thenReturn(true);
-        when(sp1.canAcquireRecords()).thenReturn(false);
 
         LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
         sharePartitions.put(tp0, sp0);
-        sharePartitions.put(tp1, sp1);
 
         CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future = new CompletableFuture<>();
         ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, "grp", Uuid.randomUuid().toString(),
-            future, List.of(tp0, tp1), BATCH_SIZE, MAX_FETCH_RECORDS,
+            future, List.of(tp0), BATCH_SIZE, MAX_FETCH_RECORDS,
             BROKER_TOPIC_STATS);
 
         when(sp0.nextFetchOffset()).thenReturn(10L);
@@ -1504,7 +1498,7 @@ public void testRemoteStorageFetchRequestCompletionOnFutureCompletionSuccessfull
             .withShareFetchData(shareFetch)
             .withSharePartitions(sharePartitions)
             .withReplicaManager(replicaManager)
-            .withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp0, tp1)))
+            .withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp0)))
             .build());
 
         when(sp0.acquire(any(), anyInt(), anyInt(), anyLong(), any(), any())).thenReturn(
@@ -1531,18 +1525,15 @@ public void testRemoteStorageFetchRequestCompletionAlongWithLocalLogRead() {
         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));
-        TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 3));
 
         SharePartition sp0 = mock(SharePartition.class);
         SharePartition sp1 = mock(SharePartition.class);
         SharePartition sp2 = mock(SharePartition.class);
-        SharePartition sp3 = mock(SharePartition.class);
 
         // Except tp3, all the topic partitions are acquirable.
         when(sp0.maybeAcquireFetchLock()).thenReturn(true);
         when(sp1.maybeAcquireFetchLock()).thenReturn(true);
         when(sp2.maybeAcquireFetchLock()).thenReturn(true);
-        when(sp3.maybeAcquireFetchLock()).thenReturn(false);
         when(sp0.canAcquireRecords()).thenReturn(true);
         when(sp1.canAcquireRecords()).thenReturn(true);
         when(sp2.canAcquireRecords()).thenReturn(true);
@@ -1551,11 +1542,10 @@ public void testRemoteStorageFetchRequestCompletionAlongWithLocalLogRead() {
         sharePartitions.put(tp0, sp0);
         sharePartitions.put(tp1, sp1);
         sharePartitions.put(tp2, sp2);
-        sharePartitions.put(tp3, sp3);
 
         CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future = new CompletableFuture<>();
         ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, "grp", Uuid.randomUuid().toString(),
-            future, List.of(tp0, tp1, tp2, tp3), BATCH_SIZE, MAX_FETCH_RECORDS,
+            future, List.of(tp0, tp1, tp2), BATCH_SIZE, MAX_FETCH_RECORDS,
             BROKER_TOPIC_STATS);
 
         when(sp0.nextFetchOffset()).thenReturn(10L);
@@ -1679,8 +1669,6 @@ public void testRemoteStorageFetchOnlyHappensForFirstTopicPartition() {
 
         when(sp0.acquire(any(), anyInt(), anyInt(), anyLong(), any(), any())).thenReturn(
             createShareAcquiredRecords(new ShareFetchResponseData.AcquiredRecords().setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1)));
-        when(sp1.acquire(any(), anyInt(), anyInt(), anyLong(), any(), any())).thenReturn(
-            createShareAcquiredRecords(new ShareFetchResponseData.AcquiredRecords().setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1)));
 
         assertFalse(delayedShareFetch.isCompleted());
         assertTrue(delayedShareFetch.tryComplete());