This is an automated email from the ASF dual-hosted git repository.
apoorvmittal10 pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new 45443a6a6f2 KAFKA-20783: Fixes for locks release for async handling in
DelayedShareFetch (6/N) (#22881)
45443a6a6f2 is described below
commit 45443a6a6f21d6d25b5a1104f4c0d644c6d2dd6a
Author: Apoorv Mittal <[email protected]>
AuthorDate: Mon Jul 20 18:29:53 2026 +0100
KAFKA-20783: Fixes for locks release for async handling in
DelayedShareFetch (6/N) (#22881)
Hardens partition fetch-lock release across the async/remote completion
paths so an acquired lock is released exactly once on every success,
timeout, and exception path.
Tests: added DelayedShareFetchTest cases for each fix covering fast-path
failure, slow-path remote transition and
exception-during-error-handling.
Reviewers: Abhinav Dixit <[email protected]>, Andrew Schofield
<[email protected]>
---
.../java/kafka/server/share/DelayedShareFetch.java | 63 ++-
.../kafka/server/share/DelayedShareFetchTest.java | 423 +++++++++++++++++++++
2 files changed, 469 insertions(+), 17 deletions(-)
diff --git a/core/src/main/java/kafka/server/share/DelayedShareFetch.java
b/core/src/main/java/kafka/server/share/DelayedShareFetch.java
index 9c410261a29..35483252bbb 100644
--- a/core/src/main/java/kafka/server/share/DelayedShareFetch.java
+++ b/core/src/main/java/kafka/server/share/DelayedShareFetch.java
@@ -265,11 +265,14 @@ public class DelayedShareFetch extends DelayedOperation {
} catch (Exception e) {
log.error("Error getting async fetch result for group {}, member
{}",
shareFetch.groupId(), shareFetch.memberId(), e);
- recordTopicPartitionsFetchRatioMetric(partitionsAcquired);
- handleFetchException(shareFetch, partitionsAcquired.keySet(), e);
- // No data could be fetched as the async read failed, hence pass
an empty set of partitions
- // with data to the action queue.
-
releasePartitionLocksAndAddToActionQueue(partitionsAcquired.keySet(), Set.of());
+ try {
+ recordTopicPartitionsFetchRatioMetric(partitionsAcquired);
+ handleFetchException(shareFetch, partitionsAcquired.keySet(),
e);
+ } finally {
+ // Release the locks in a finally block so that a throw from
the metric recording or the
+ // exception handling above cannot leak the partition locks.
+ releasePartitionLocks(partitionsAcquired.keySet());
+ }
return;
} finally {
// Clear the pending future regardless of completion status
@@ -336,7 +339,7 @@ public class DelayedShareFetch extends DelayedOperation {
log.error("Error initiating async fetch during request completion
for group {}, member {}",
shareFetch.groupId(), shareFetch.memberId(), e);
handleFetchException(shareFetch, topicPartitionData.keySet(), e);
-
releasePartitionLocksAndAddToActionQueue(topicPartitionData.keySet(), new
HashSet<>());
+ releasePartitionLocks(topicPartitionData.keySet());
return;
}
// If the future is successfully created then complete the request
asynchronously. The completion
@@ -346,7 +349,7 @@ public class DelayedShareFetch extends DelayedOperation {
log.error("Async fetch failed during request completion for
group {}, member {}",
shareFetch.groupId(), shareFetch.memberId(), throwable);
handleFetchException(shareFetch, topicPartitionData.keySet(),
throwable);
-
releasePartitionLocksAndAddToActionQueue(topicPartitionData.keySet(), new
HashSet<>());
+ releasePartitionLocks(topicPartitionData.keySet());
return;
}
// Complete the share fetch request and release the locks as
completion handler will complete
@@ -724,9 +727,17 @@ public class DelayedShareFetch extends DelayedOperation {
// Remove pending fetch as the request is completed now.
pendingFetch = null;
} catch (Exception e) {
- // Force complete to send error response to client as the pending
future is not marked null
- // yet hence force complete will try to retrieve the data again
which shall fail and complete
- // the request.
+ // On the slow path pendingFetch is set, so forceComplete() ->
onComplete() re-enters
+ // completeShareFetchAsyncRequest(), which retrieves the data
again (failing the same way),
+ // surfaces the error and releases the partition locks held under
partitionsAcquired.
+ //
+ // On the fast path (invoked directly from tryComplete() when the
read completed
+ // synchronously) set the partitionsAcquired so onComplete() can
release the held partitions.
+ if (pendingFetch == null) {
+ log.error("Error retrieving async fetch result for group {},
member {}",
+ shareFetch.groupId(), shareFetch.memberId(), e);
+ partitionsAcquired = topicPartitionData;
+ }
return forceComplete();
}
// Process the successful read response - check for remote fetch,
minBytes, etc.
@@ -757,7 +768,10 @@ public class DelayedShareFetch extends DelayedOperation {
if (anyPartitionHasLogReadError(readResponse) ||
isMinBytesSatisfied(topicPartitionData,
partitionMaxBytesStrategy.maxBytes(shareFetch.fetchParams().maxBytes,
topicPartitionData.keySet(), topicPartitionData.size()))) {
- // Request can be completed - set state and force completion
+ // Request can be completed - set state and force completion.
Though on-slow path when
+ // readFuture is not syncronously completed the partitionsAcquired
will be already set
+ // but on fast-path when readFuture is completed synchronously, we
need to set the
+ // partitionsAcquired.
partitionsAcquired = topicPartitionData;
localPartitionsAlreadyFetched = readResponse;
return forceComplete();
@@ -797,7 +811,7 @@ public class DelayedShareFetch extends DelayedOperation {
} finally {
// Release locks for all partitions acquired for this request. The
async fetch did not
// complete, so no partition has data, hence pass an empty set as
the partitions with data.
-
releasePartitionLocksAndAddToActionQueue(topicPartitionData.keySet(), Set.of());
+ releasePartitionLocks(topicPartitionData.keySet());
}
}
@@ -893,6 +907,15 @@ public class DelayedShareFetch extends DelayedOperation {
// 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,
nonRemoteFetchTopicPartitions);
+ // Remove the just-released non-remote partitions from
partitionsAcquired so it tracks only the
+ // remote-fetch partitions whose locks are still held. On the slow
path partitionsAcquired was
+ // populated in tryComplete with ALL acquired partitions (remote and
non-remote); without this
+ // removal the remote completion path would later release these
already-released locks a second
+ // time (releaseFetchLock force-clears the lock, so this can steal a
lock re-acquired by another
+ // request) and would skip the additional local read for these
partitions (they would still look
+ // acquired). On the fast path partitionsAcquired only holds the
remote partitions, so this is a
+ // no-op there.
+ nonRemoteFetchTopicPartitions.forEach(partitionsAcquired::remove);
processRemoteFetchOrException(remoteStorageFetchInfoMap);
// Check if remote fetch can be completed.
return maybeCompletePendingRemoteFetch();
@@ -1142,14 +1165,14 @@ public class DelayedShareFetch extends DelayedOperation
{
// Release locks on error
Set<TopicIdPartition> topicIdPartitions = new
LinkedHashSet<>(partitionsAcquired.keySet());
topicIdPartitions.addAll(acquiredNonRemoteFetchTopicPartitionData.keySet());
- releasePartitionLocksAndAddToActionQueue(topicIdPartitions, new
HashSet<>());
+ releasePartitionLocks(topicIdPartitions);
} catch (Exception e) {
log.error("Unexpected error in processing delayed share fetch
request", e);
handleExceptionInCompletingRemoteStorageShareFetchRequest(acquiredNonRemoteFetchTopicPartitionData.keySet(),
e);
// Release locks on error
Set<TopicIdPartition> topicIdPartitions = new
LinkedHashSet<>(partitionsAcquired.keySet());
topicIdPartitions.addAll(acquiredNonRemoteFetchTopicPartitionData.keySet());
- releasePartitionLocksAndAddToActionQueue(topicIdPartitions, new
HashSet<>());
+ releasePartitionLocks(topicIdPartitions);
}
}
@@ -1176,7 +1199,13 @@ public class DelayedShareFetch extends DelayedOperation {
// locks are release in finally block in
readAndProcessFetchResultCompletion hence share
// fetch will be completed in the current try block prior locks
getting released. Hence,
// need to release the locks for the remote fetch partitions here.
-
releasePartitionLocksAndAddToActionQueue(partitionsAcquired.keySet(),
partitionsWithData(shareFetchPartitionDataList));
+ //
+ // shareFetchPartitionDataList can also contain the local
partitions read alongside the remote
+ // fetch (added via the readAndProcessFetchResultCompletion
callback), whose action queue
+ // entries were already handled by processFetchResultAndComplete.
+ Set<TopicIdPartition> remotePartitionsWithData = new
HashSet<>(partitionsWithData(shareFetchPartitionDataList));
+ remotePartitionsWithData.retainAll(partitionsAcquired.keySet());
+
releasePartitionLocksAndAddToActionQueue(partitionsAcquired.keySet(),
remotePartitionsWithData);
}
}
@@ -1201,7 +1230,7 @@ public class DelayedShareFetch extends DelayedOperation {
// processFetchResultAndComplete won't be invoked in this error
path, so it won't release the
// locks for partitionsToFetch. Release them here to avoid leaking
the local partition locks
// before completing with just the remote data.
-
releasePartitionLocksAndAddToActionQueue(partitionsToFetch.keySet(), Set.of());
+ releasePartitionLocks(partitionsToFetch.keySet());
// Continue to complete with just remote data.
completionConsumer.accept(List.of());
return;
@@ -1217,7 +1246,7 @@ public class DelayedShareFetch extends DelayedOperation {
// processFetchResultAndComplete won't be invoked in this
error path, so it won't release
// the locks for partitionsToFetch. Release them here to avoid
leaking the local partition
// locks before completing with just the remote data.
-
releasePartitionLocksAndAddToActionQueue(partitionsToFetch.keySet(), Set.of());
+ releasePartitionLocks(partitionsToFetch.keySet());
// Continue to complete with just remote data.
completionConsumer.accept(List.of());
return;
diff --git a/core/src/test/java/kafka/server/share/DelayedShareFetchTest.java
b/core/src/test/java/kafka/server/share/DelayedShareFetchTest.java
index ff2bb171dae..de997acf1e8 100644
--- a/core/src/test/java/kafka/server/share/DelayedShareFetchTest.java
+++ b/core/src/test/java/kafka/server/share/DelayedShareFetchTest.java
@@ -77,6 +77,7 @@ import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Collectors;
@@ -96,6 +97,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyBoolean;
@@ -953,6 +955,63 @@ public class DelayedShareFetchTest {
Mockito.verify(exceptionHandler, times(1)).accept(any(), any());
}
+ @Test
+ public void testTryCompleteReleasesLocksWhenAsyncReadImmediatelyThrows() {
+ // Reproduces the fast-path lock check: when maybeReadFromLog's future
is already completed
+ // but exceptionally.
+ ReplicaManager replicaManager = mock(ReplicaManager.class);
+ TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new
TopicPartition("foo", 0));
+
+ SharePartition sp0 = mock(SharePartition.class);
+ when(sp0.canAcquireRecords()).thenReturn(true);
+ // Empty fetch offset metadata forces maybeReadFromLog to actually
read from the log.
+ when(sp0.fetchOffsetMetadata(anyLong())).thenReturn(Optional.empty());
+
+ LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new
LinkedHashMap<>();
+ sharePartitions.put(tp0, sp0);
+
+ CompletableFuture<Map<TopicIdPartition,
ShareFetchResponseData.PartitionData>> future = new CompletableFuture<>();
+ ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, "grp",
Uuid.randomUuid().toString(),
+ future, List.of(tp0), BATCH_OPTIMIZED, BATCH_SIZE,
MAX_FETCH_RECORDS,
+ BROKER_TOPIC_STATS);
+
+ // A synchronous throw from readFromLog makes
ReplicaManagerLogReader.readAsync return an
+ // already-completed (failed) future, so tryComplete takes the fast
path with a completed future.
+ when(replicaManager.readFromLog(any(), any(), any(ReplicaQuota.class),
anyBoolean()))
+ .thenThrow(new RuntimeException("Read failed"));
+
+ PartitionMaxBytesStrategy partitionMaxBytesStrategy =
mockPartitionMaxBytes(Set.of(tp0));
+ BiConsumer<SharePartitionKey, Throwable> exceptionHandler =
mockExceptionHandler();
+ Uuid fetchId = Uuid.randomUuid();
+ DelayedShareFetch delayedShareFetch =
spy(DelayedShareFetchBuilder.builder()
+ .withShareFetchData(shareFetch)
+ .withSharePartitions(sharePartitions)
+ .withReplicaManagerLogReader(replicaManager)
+ .withExceptionHandler(exceptionHandler)
+ .withPartitionMaxBytesStrategy(partitionMaxBytesStrategy)
+ .withFetchId(fetchId)
+ .build());
+
+ // The first acquire (in tryComplete) succeeds; a subsequent acquire
attempt fails, mirroring the
+ // real SharePartition where the lock is already held by this fetchId
and cannot be re-acquired.
+
when(sp0.maybeAcquireFetchLock(fetchId)).thenReturn(true).thenReturn(false);
+
+ assertFalse(delayedShareFetch.isCompleted());
+ // The fast-path exception handling completes the request with the
error.
+ assertTrue(delayedShareFetch.tryComplete());
+ assertTrue(delayedShareFetch.isCompleted());
+ assertTrue(shareFetch.errorInAllPartitions());
+ assertTrue(future.isDone());
+
+ // The fetch lock must be released (exactly once).
+ Mockito.verify(sp0, times(1)).releaseFetchLock(fetchId);
+ Mockito.verify(exceptionHandler, times(1)).accept(any(), any());
+ Map<TopicIdPartition, ShareFetchResponseData.PartitionData> response =
future.join();
+ assertEquals(Set.of(tp0), response.keySet());
+ assertEquals(Errors.UNKNOWN_SERVER_ERROR.code(),
response.get(tp0).errorCode());
+ assertTrue(response.get(tp0).errorMessage().contains("Read failed"));
+ }
+
@Test
public void testTryCompleteLocksReleasedOnCompleteException() {
ReplicaManager replicaManager = mock(ReplicaManager.class);
@@ -1488,6 +1547,87 @@ public class DelayedShareFetchTest {
delayedShareFetch.lock().unlock();
}
+ /**
+ * Validates the tryComplete() that handles an in-flight remote fetch:
when tryComplete() is
+ * re-invoked (e.g. via a purgatory checkAndComplete) while a remote fetch
initiated by an earlier
+ * tryComplete() is still pending, it must route to
maybeCompletePendingRemoteFetch() rather than
+ * re-acquiring partitions or starting a new read. It stays in purgatory
while the remote fetch is
+ * pending and completes once the remote fetch is done.
+ */
+ @Test
+ public void testTryCompleteWithPendingRemoteFetchInFlight() {
+ ReplicaManager replicaManager = mock(ReplicaManager.class);
+ TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new
TopicPartition("foo", 0));
+ SharePartition sp0 = mock(SharePartition.class);
+ when(sp0.canAcquireRecords()).thenReturn(true);
+ when(sp0.nextFetchOffset()).thenReturn(10L);
+ when(sp0.fetchOffsetMetadata(anyLong())).thenReturn(Optional.empty());
+
+ LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new
LinkedHashMap<>();
+ sharePartitions.put(tp0, sp0);
+
+ CompletableFuture<Map<TopicIdPartition,
ShareFetchResponseData.PartitionData>> future = new CompletableFuture<>();
+ ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, "grp",
Uuid.randomUuid().toString(),
+ future, List.of(tp0), BATCH_OPTIMIZED, BATCH_SIZE,
MAX_FETCH_RECORDS, BROKER_TOPIC_STATS);
+
+ // tp0 is served from remote storage.
+ doAnswer(invocation -> buildLocalAndRemoteFetchResult(Set.of(),
Set.of(tp0)))
+ .when(replicaManager).readFromLog(any(), any(),
any(ReplicaQuota.class), anyBoolean());
+
+ // Capture the remote read completion callback but do not invoke it
yet, so the remote fetch stays
+ // in flight after the first tryComplete().
+ AtomicReference<Consumer<RemoteLogReadResult>> remoteCallback = new
AtomicReference<>();
+ RemoteLogManager remoteLogManager = mock(RemoteLogManager.class);
+ when(remoteLogManager.asyncRead(any(), any())).thenAnswer(invocation
-> {
+ remoteCallback.set(invocation.getArgument(1));
+ return mock(Future.class);
+ });
+
when(replicaManager.remoteLogManager()).thenReturn(Option.apply(remoteLogManager));
+
+ Partition p0 = mock(Partition.class);
+ when(p0.isLeader()).thenReturn(true);
+
when(replicaManager.getPartitionOrException(tp0.topicPartition())).thenReturn(p0);
+
+ Uuid fetchId = Uuid.randomUuid();
+ DelayedShareFetch delayedShareFetch =
spy(DelayedShareFetchBuilder.builder()
+ .withShareFetchData(shareFetch)
+ .withSharePartitions(sharePartitions)
+ .withReplicaManagerLogReader(replicaManager)
+ .withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp0)))
+ .withFetchId(fetchId)
+ .build());
+
+ when(sp0.maybeAcquireFetchLock(fetchId)).thenReturn(true);
+
+ try (MockedStatic<ShareFetchUtils> mockedShareFetchUtils =
Mockito.mockStatic(ShareFetchUtils.class)) {
+ mockedShareFetchUtils.when(() ->
ShareFetchUtils.processFetchResponse(any(), any(), any(), any(), any()))
+ .thenReturn(Map.of(tp0,
mock(ShareFetchResponseData.PartitionData.class)));
+
+ // First tryComplete initiates the remote fetch which does not
complete synchronously.
+ assertFalse(delayedShareFetch.tryComplete());
+ assertNotNull(delayedShareFetch.pendingRemoteFetches());
+ assertFalse(delayedShareFetch.isCompleted());
+
+ // Second tryComplete re-enters with a pending remote fetch in
flight (the branch under test).
+ // The remote fetch is still not done, so it stays in purgatory
and does not re-acquire the
+ // partition or start a new read.
+ assertFalse(delayedShareFetch.tryComplete());
+ assertFalse(delayedShareFetch.isCompleted());
+ Mockito.verify(sp0, times(1)).maybeAcquireFetchLock(fetchId);
+ Mockito.verify(sp0, times(0)).releaseFetchLock(fetchId);
+
+ // Complete the in-flight remote fetch, then a subsequent
tryComplete completes via the same
+ // branch (maybeCompletePendingRemoteFetch now sees the remote
fetch as done).
+ remoteCallback.get().accept(new
RemoteLogReadResult(Optional.of(REMOTE_FETCH_INFO), Optional.empty()));
+ assertTrue(delayedShareFetch.tryComplete());
+ assertTrue(delayedShareFetch.isCompleted());
+ assertTrue(shareFetch.isCompleted());
+ assertEquals(Set.of(tp0), future.join().keySet());
+ Mockito.verify(sp0, times(1)).maybeAcquireFetchLock(fetchId);
+ Mockito.verify(sp0, times(1)).releaseFetchLock(fetchId);
+ }
+ }
+
@SuppressWarnings("unchecked")
@Test
public void testRemoteStorageFetchPartitionLeaderChanged() {
@@ -2020,6 +2160,224 @@ public class DelayedShareFetchTest {
}
}
+ /**
+ * Tests that when a remote fetch completes alongside a successful
additional local read (with data),
+ * the remote-fetch completion path does not re-enqueue the local
partitions to the action queue.
+ * The local partitions' completion is separately triggered by
processFetchResultAndComplete, so
+ * completeRemoteFetchRequest must only enqueue the remote (acquired)
partitions.
+ */
+ @Test
+ public void
testRemoteStorageFetchCompletionDoesNotDuplicateLocalPartitionsInActionQueue() {
+ 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));
+
+ SharePartition sp0 = mock(SharePartition.class);
+ SharePartition sp1 = mock(SharePartition.class);
+ SharePartition sp2 = mock(SharePartition.class);
+ when(sp0.canAcquireRecords()).thenReturn(true);
+ when(sp1.canAcquireRecords()).thenReturn(true);
+ when(sp2.canAcquireRecords()).thenReturn(true);
+
+ LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new
LinkedHashMap<>();
+ sharePartitions.put(tp0, sp0);
+ sharePartitions.put(tp1, sp1);
+ sharePartitions.put(tp2, sp2);
+
+ CompletableFuture<Map<TopicIdPartition,
ShareFetchResponseData.PartitionData>> future = new CompletableFuture<>();
+ ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, "grp",
Uuid.randomUuid().toString(),
+ future, List.of(tp0, tp1, tp2), BATCH_OPTIMIZED, BATCH_SIZE,
MAX_FETCH_RECORDS,
+ BROKER_TOPIC_STATS);
+
+ when(sp0.nextFetchOffset()).thenReturn(10L);
+ when(sp1.nextFetchOffset()).thenReturn(20L);
+ when(sp2.nextFetchOffset()).thenReturn(30L);
+ when(sp0.fetchOffsetMetadata(anyLong())).thenReturn(Optional.empty());
+ when(sp1.fetchOffsetMetadata(anyLong())).thenReturn(Optional.empty());
+ when(sp2.fetchOffsetMetadata(anyLong())).thenReturn(Optional.empty());
+
+ try (MockedStatic<ShareFetchUtils> mockedShareFetchUtils =
Mockito.mockStatic(ShareFetchUtils.class)) {
+ Map<TopicIdPartition, ShareFetchResponseData.PartitionData>
partitionDataMap = new LinkedHashMap<>();
+ partitionDataMap.put(tp0,
mock(ShareFetchResponseData.PartitionData.class));
+ partitionDataMap.put(tp1,
mock(ShareFetchResponseData.PartitionData.class));
+ partitionDataMap.put(tp2,
mock(ShareFetchResponseData.PartitionData.class));
+ mockedShareFetchUtils.when(() ->
ShareFetchUtils.processFetchResponse(any(), any(), any(), any(),
any())).thenReturn(partitionDataMap);
+
+ // First read (tryComplete): tp0, tp1 local and tp2 remote. Second
read (additional local read
+ // in onComplete): tp0, tp1 returned WITH data so that they appear
in partitionsWithData.
+ doAnswer(invocation -> buildLocalAndRemoteFetchResult(Set.of(tp0,
tp1), Set.of(tp2))
+ ).doAnswer(invocation -> buildLogReadResult(List.of(tp0, tp1))
+ ).when(replicaManager).readFromLog(any(), any(),
any(ReplicaQuota.class), anyBoolean());
+
+ // Remote fetch object completes within tryComplete so the request
moves on to onComplete.
+ RemoteLogReadResult remoteFetchResult = new
RemoteLogReadResult(Optional.of(REMOTE_FETCH_INFO), Optional.empty());
+ RemoteLogManager remoteLogManager = mock(RemoteLogManager.class);
+ doAnswer(invocationOnMock -> {
+ Consumer<RemoteLogReadResult> callback =
invocationOnMock.getArgument(1);
+ callback.accept(remoteFetchResult);
+ return CompletableFuture.completedFuture(remoteFetchResult);
+ }).when(remoteLogManager).asyncRead(any(), any());
+
when(replicaManager.remoteLogManager()).thenReturn(Option.apply(remoteLogManager));
+
+ Partition p0 = mock(Partition.class);
+ when(p0.isLeader()).thenReturn(true);
+ Partition p1 = mock(Partition.class);
+ when(p1.isLeader()).thenReturn(true);
+ Partition p2 = mock(Partition.class);
+ when(p2.isLeader()).thenReturn(true);
+
when(replicaManager.getPartitionOrException(tp0.topicPartition())).thenReturn(p0);
+
when(replicaManager.getPartitionOrException(tp1.topicPartition())).thenReturn(p1);
+
when(replicaManager.getPartitionOrException(tp2.topicPartition())).thenReturn(p2);
+
+ // Execute the action queue runnables synchronously so
completeDelayedShareFetchRequest is invoked.
+ doAnswer(invocation -> {
+ Runnable runnable = invocation.getArgument(0);
+ runnable.run();
+ return null;
+ }).when(replicaManager).addToActionQueue(any());
+
+ Uuid fetchId = Uuid.randomUuid();
+ DelayedShareFetch delayedShareFetch =
spy(DelayedShareFetchBuilder.builder()
+ .withShareFetchData(shareFetch)
+ .withReplicaManagerLogReader(replicaManager)
+ .withSharePartitions(sharePartitions)
+
.withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp0, tp1, tp2)))
+ .withFetchId(fetchId)
+ .build());
+
+ when(sp0.maybeAcquireFetchLock(fetchId)).thenReturn(true);
+ when(sp1.maybeAcquireFetchLock(fetchId)).thenReturn(true);
+ when(sp2.maybeAcquireFetchLock(fetchId)).thenReturn(true);
+
+ assertTrue(delayedShareFetch.tryComplete());
+ assertTrue(delayedShareFetch.isCompleted());
+ assertTrue(shareFetch.isCompleted());
+
+ // The local partitions tp0 and tp1 are enqueued twice: once when
they are released as
+ // non-remote partitions while the remote fetch is prepared in
tryComplete, and again
+ // after they are re-acquired for the additional local read in
onComplete (via
+ // processFetchResultAndComplete). The remote-fetch completion
(completeRemoteFetchRequest)
+ // must not enqueue them a third time - it only enqueues the
remote (acquired) partitions.
+ Mockito.verify(replicaManager,
times(2)).completeDelayedShareFetchRequest(
+ new DelayedShareFetchGroupKey(shareFetch.groupId(),
tp0.topicId(), tp0.partition()));
+ Mockito.verify(replicaManager,
times(2)).completeDelayedShareFetchRequest(
+ new DelayedShareFetchGroupKey(shareFetch.groupId(),
tp1.topicId(), tp1.partition()));
+ // The remote partition tp2 is enqueued once, by
completeRemoteFetchRequest.
+ Mockito.verify(replicaManager,
times(1)).completeDelayedShareFetchRequest(
+ new DelayedShareFetchGroupKey(shareFetch.groupId(),
tp2.topicId(), tp2.partition()));
+
+ // Validate acquire and release locks for the partitions.
+ Mockito.verify(sp0, times(2)).maybeAcquireFetchLock(fetchId);
+ Mockito.verify(sp1, times(2)).maybeAcquireFetchLock(fetchId);
+ Mockito.verify(sp2, times(1)).maybeAcquireFetchLock(fetchId);
+ Mockito.verify(sp0, times(2)).releaseFetchLock(fetchId);
+ Mockito.verify(sp1, times(2)).releaseFetchLock(fetchId);
+ Mockito.verify(sp2, times(1)).releaseFetchLock(fetchId);
+ }
+ }
+
+ /**
+ * On the slow path, tryComplete records partitionsAcquired = ALL acquired
partitions (remote and
+ * non-remote) before the async local read completes. When the completed
read then routes some
+ * partitions to a remote fetch, maybeProcessRemoteFetch releases the
non-remote partitions and must
+ * also remove them from partitionsAcquired.
+ */
+ @Test
+ public void
testAsyncReadRemoteFetchRemovesNonRemotePartitionsFromAcquiredOnRelease() {
+ String groupId = "grp";
+ LogReader logReader = mock(LogReader.class);
+ ReplicaManager replicaManager = mock(ReplicaManager.class);
+ PartitionMetadataProvider metadataProvider =
mock(PartitionMetadataProvider.class);
+ TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new
TopicPartition("foo", 0)); // remote
+ TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new
TopicPartition("foo", 1)); // non-remote
+
+ SharePartition sp0 = mock(SharePartition.class);
+ SharePartition sp1 = mock(SharePartition.class);
+ when(sp0.canAcquireRecords()).thenReturn(true);
+ when(sp1.canAcquireRecords()).thenReturn(true);
+ when(sp0.nextFetchOffset()).thenReturn(10L);
+ when(sp1.nextFetchOffset()).thenReturn(20L);
+ when(sp0.fetchOffsetMetadata(anyLong())).thenReturn(Optional.empty());
+ when(sp1.fetchOffsetMetadata(anyLong())).thenReturn(Optional.empty());
+
+ 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, groupId,
Uuid.randomUuid().toString(),
+ future, List.of(tp0, tp1), BATCH_OPTIMIZED, BATCH_SIZE,
MAX_FETCH_RECORDS, BROKER_TOPIC_STATS);
+
+ // Slow path: the first readAsync stays pending, so tryComplete
records partitionsAcquired = {tp0, tp1}
+ // and stores the pending fetch. The second readAsync is the
additional local read for tp1 in onComplete.
+ CompletableFuture<LinkedHashMap<TopicIdPartition, LogReadResult>>
asyncReadFuture = new CompletableFuture<>();
+ when(logReader.readAsync(any(), any(), any(), any(), anyBoolean()))
+ .thenReturn(asyncReadFuture)
+
.thenReturn(CompletableFuture.completedFuture(buildLogReadResultMap(List.of(tp1))));
+
+ Uuid fetchId = Uuid.randomUuid();
+ when(sp0.maybeAcquireFetchLock(fetchId)).thenReturn(true);
+ when(sp1.maybeAcquireFetchLock(fetchId)).thenReturn(true);
+
+ // Remote fetch for tp0 completes synchronously so onComplete runs
during the second tryComplete.
+ RemoteLogReadResult remoteFetchResult = new
RemoteLogReadResult(Optional.of(REMOTE_FETCH_INFO), Optional.empty());
+ RemoteLogManager remoteLogManager = mock(RemoteLogManager.class);
+ doAnswer(inv -> {
+ Consumer<RemoteLogReadResult> callback = inv.getArgument(1);
+ callback.accept(remoteFetchResult);
+ return CompletableFuture.completedFuture(remoteFetchResult);
+ }).when(remoteLogManager).asyncRead(any(), any());
+
when(replicaManager.remoteLogManager()).thenReturn(Option.apply(remoteLogManager));
+
+ Partition p0 = mock(Partition.class);
+ when(p0.isLeader()).thenReturn(true);
+
when(replicaManager.getPartitionOrException(tp0.topicPartition())).thenReturn(p0);
+
+ DelayedShareFetch delayedShareFetch =
spy(DelayedShareFetchBuilder.builder()
+ .withShareFetchData(shareFetch)
+ .withSharePartitions(sharePartitions)
+ .withReplicaManager(replicaManager)
+ .withLogReader(logReader)
+ .withMetadataProvider(metadataProvider)
+ .withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp0,
tp1)))
+ .withFetchId(fetchId)
+ .build());
+
+ try (MockedStatic<ShareFetchUtils> mockedShareFetchUtils =
Mockito.mockStatic(ShareFetchUtils.class)) {
+ Map<TopicIdPartition, ShareFetchResponseData.PartitionData>
partitionDataMap = new LinkedHashMap<>();
+ partitionDataMap.put(tp0,
mock(ShareFetchResponseData.PartitionData.class));
+ partitionDataMap.put(tp1,
mock(ShareFetchResponseData.PartitionData.class));
+ mockedShareFetchUtils.when(() ->
ShareFetchUtils.processFetchResponse(any(), any(), any(), any(),
any())).thenReturn(partitionDataMap);
+
+ // First tryComplete: slow path, records partitionsAcquired =
{tp0, tp1}.
+ assertFalse(delayedShareFetch.tryComplete());
+ assertNotNull(delayedShareFetch.pendingFetch());
+
+ Mockito.verify(sp0, times(1)).maybeAcquireFetchLock(fetchId);
+ Mockito.verify(sp1, times(1)).maybeAcquireFetchLock(fetchId);
+ Mockito.verify(sp0, times(0)).releaseFetchLock(fetchId);
+ Mockito.verify(sp1, times(0)).releaseFetchLock(fetchId);
+ // Complete the async local read: tp0 is tiered (remote), tp1 is a
local partition.
+
asyncReadFuture.complete(buildLocalAndRemoteFetchResultMap(Set.of(tp1),
Set.of(tp0)));
+
+ // Next tryComplete transitions tp0 to a remote fetch (releasing
tp1 as a
+ // non-remote partition and removing it from partitionsAcquired),
and completes.
+ assertTrue(delayedShareFetch.tryComplete());
+ assertTrue(delayedShareFetch.isCompleted());
+ assertTrue(shareFetch.isCompleted());
+
+ Mockito.verify(sp0, times(1)).maybeAcquireFetchLock(fetchId);
+ // tp1 was released as a non-remote partition and removed from
partitionsAcquired, so the remote
+ // completion re-acquires it for the additional local read (a
second maybeAcquireFetchLock) and
+ // performs a second readAsync - rather than leaving tp1 in
partitionsAcquired.
+ Mockito.verify(sp1, times(2)).maybeAcquireFetchLock(fetchId);
+ Mockito.verify(sp0, times(1)).releaseFetchLock(fetchId);
+ Mockito.verify(sp1, times(2)).releaseFetchLock(fetchId);
+ Mockito.verify(logReader, times(2)).readAsync(any(), any(), any(),
any(), anyBoolean());
+ }
+ }
+
@Test
public void testRemoteStorageFetchHappensForAllTopicPartitions() {
ReplicaManager replicaManager = mock(ReplicaManager.class);
@@ -2505,6 +2863,71 @@ public class DelayedShareFetchTest {
Mockito.verify(sp0).releaseFetchLock(fetchId);
}
+ /**
+ * Tests that when the async read fails and the error handling in
completeShareFetchAsyncRequest
+ * itself throws (here the exception handler callback), the partition
locks are still released
+ * This guards the try-finally around the metric recording and exception
handling in the get()-failure
+ * path.
+ */
+ @Test
+ public void
testCompleteShareFetchAsyncRequestReleasesLocksWhenExceptionHandlingThrows() {
+ String groupId = "grp";
+ Uuid topicId = Uuid.randomUuid();
+ LogReader logReader = mock(LogReader.class);
+ PartitionMetadataProvider metadataProvider =
mock(PartitionMetadataProvider.class);
+ ReplicaManager replicaManager = mock(ReplicaManager.class);
+ TopicIdPartition tp0 = new TopicIdPartition(topicId, new
TopicPartition("foo", 0));
+
+ SharePartition sp0 = mock(SharePartition.class);
+ LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new
LinkedHashMap<>();
+ sharePartitions.put(tp0, sp0);
+
+ CompletableFuture<Map<TopicIdPartition,
ShareFetchResponseData.PartitionData>> future = new CompletableFuture<>();
+ ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId,
Uuid.randomUuid().toString(),
+ future, List.of(tp0), BATCH_OPTIMIZED, BATCH_SIZE,
MAX_FETCH_RECORDS, BROKER_TOPIC_STATS);
+
+ CompletableFuture<LinkedHashMap<TopicIdPartition, LogReadResult>>
asyncReadFuture = new CompletableFuture<>();
+ when(logReader.readAsync(any(), any(), any(), any(),
anyBoolean())).thenReturn(asyncReadFuture);
+
+ when(sp0.maybeAcquireFetchLock(any())).thenReturn(true);
+ when(sp0.canAcquireRecords()).thenReturn(true);
+ when(sp0.nextFetchOffset()).thenReturn(0L);
+ when(sp0.fetchOffsetMetadata(anyLong()))
+ .thenReturn(Optional.empty())
+ .thenReturn(Optional.of(new LogOffsetMetadata(0, 1, 0)));
+ when(metadataProvider.endOffsetMetadata(any(), any())).thenReturn(new
LogOffsetMetadata(1, 1, 1));
+
+ PartitionMaxBytesStrategy partitionMaxBytesStrategy =
mockPartitionMaxBytes(Set.of(tp0));
+
+ // The exception handler throws while handling the async read failure,
simulating a failure
+ // inside the catch block's error handling. Locks must still be
released via the finally block.
+ BiConsumer<SharePartitionKey, Throwable> exceptionHandler =
mockExceptionHandler();
+ doThrow(new RuntimeException("handler
boom")).when(exceptionHandler).accept(any(), any());
+
+ Uuid fetchId = Uuid.randomUuid();
+ DelayedShareFetch delayedShareFetch =
DelayedShareFetchBuilder.builder()
+ .withShareFetchData(shareFetch)
+ .withReplicaManager(replicaManager)
+ .withLogReader(logReader)
+ .withMetadataProvider(metadataProvider)
+ .withPartitionMaxBytesStrategy(partitionMaxBytesStrategy)
+ .withSharePartitions(sharePartitions)
+ .withExceptionHandler(exceptionHandler)
+ .withFetchId(fetchId)
+ .build();
+
+ // First tryComplete triggers the async read, which stays pending
(slow path).
+ assertFalse(delayedShareFetch.tryComplete());
+ assertNotNull(delayedShareFetch.pendingFetch());
+ // Complete the async read exceptionally.
+ asyncReadFuture.completeExceptionally(new RuntimeException("Simulated
read failure"));
+
+ // Purgatory re-invokes tryComplete; get() throws, and the error
handling (exception handler)
+ // then throws too. The finally block must still release the partition
locks before the throw propagates.
+ assertThrows(RuntimeException.class, delayedShareFetch::tryComplete);
+ Mockito.verify(sp0).releaseFetchLock(fetchId);
+ }
+
/**
* Tests that onComplete handles the case where async fetch is still
pending.
* This simulates timeout scenario where async read hasn't completed yet.
