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 ac0d7e245b2 KAFKA-20783: Completion handling and more tests (3/N)
(#22804)
ac0d7e245b2 is described below
commit ac0d7e245b213c20ab66bcea211d1cbd5ed2c95a
Author: Apoorv Mittal <[email protected]>
AuthorDate: Fri Jul 10 17:24:39 2026 +0100
KAFKA-20783: Completion handling and more tests (3/N) (#22804)
Minor change and more tests. I ll continue adding more tests for async
handling.
Reviewers: Andrew Schofield <[email protected]>
---
.../java/kafka/server/share/DelayedShareFetch.java | 16 +-
.../kafka/server/share/DelayedShareFetchTest.java | 219 ++++++++++++++++++++-
2 files changed, 231 insertions(+), 4 deletions(-)
diff --git a/core/src/main/java/kafka/server/share/DelayedShareFetch.java
b/core/src/main/java/kafka/server/share/DelayedShareFetch.java
index e60d3979b20..312c1d1c102 100644
--- a/core/src/main/java/kafka/server/share/DelayedShareFetch.java
+++ b/core/src/main/java/kafka/server/share/DelayedShareFetch.java
@@ -63,6 +63,7 @@ import java.util.Optional;
import java.util.OptionalInt;
import java.util.OptionalLong;
import java.util.Set;
+import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
@@ -476,6 +477,13 @@ public class DelayedShareFetch extends DelayedOperation {
TopicIdPartition topicIdPartition =
topicPartitionData.keySet().iterator().next();
readFuture.whenComplete((result, throwable) -> {
if (throwable != null) {
+ if (throwable instanceof CancellationException) {
+ // The completion handler execution has been
cancelled due to the operation
+ // being completed in onComplete.
+ log.trace("Async read was cancelled for group
{}, member {}",
+ shareFetch.groupId(),
shareFetch.memberId());
+ return;
+ }
log.error("Async read failed for group {}, member
{}",
shareFetch.groupId(), shareFetch.memberId(),
throwable);
// Error is handled in maybeCompleteAsyncFetch()
via get() exception handling
@@ -774,6 +782,8 @@ public class DelayedShareFetch extends DelayedOperation {
try {
log.warn("Completing async fetch for group {}, member {},
partitions {} with empty response",
shareFetch.groupId(), shareFetch.memberId(),
topicPartitionData.keySet());
+ // Cancel the future invocation of pending fetch.
+ pendingFetch.cancel(true);
pendingFetch = null;
// Update fetch ratio metric and complete the request with empty
response.
recordTopicPartitionsFetchRatioMetric(topicPartitionData);
@@ -884,10 +894,10 @@ public class DelayedShareFetch extends DelayedOperation {
private boolean maybeRegisterCallbackPendingRemoteFetch() {
log.trace("Registering callback pending remote fetch");
- PendingRemoteFetches pendingFetch = pendingRemoteFetchesOpt.get();
- if (!pendingFetch.isDone() && shareFetch.fetchParams().maxWaitMs <
remoteFetchMaxWaitMs) {
+ PendingRemoteFetches pendingRemoteFetches =
pendingRemoteFetchesOpt.get();
+ if (!pendingRemoteFetches.isDone() &&
shareFetch.fetchParams().maxWaitMs < remoteFetchMaxWaitMs) {
TimerTask timerTask = new PendingRemoteFetchTimerTask();
- pendingFetch.invokeCallbackOnCompletion(((ignored, throwable) -> {
+ pendingRemoteFetches.invokeCallbackOnCompletion(((ignored,
throwable) -> {
timerTask.cancel();
log.trace("Invoked remote storage fetch callback for group {},
member {}, "
+ "topic partitions {}", shareFetch.groupId(),
shareFetch.memberId(),
diff --git a/core/src/test/java/kafka/server/share/DelayedShareFetchTest.java
b/core/src/test/java/kafka/server/share/DelayedShareFetchTest.java
index e04274f0751..d89594807f4 100644
--- a/core/src/test/java/kafka/server/share/DelayedShareFetchTest.java
+++ b/core/src/test/java/kafka/server/share/DelayedShareFetchTest.java
@@ -2378,8 +2378,9 @@ public class DelayedShareFetchTest {
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(),
- new CompletableFuture<>(), List.of(tp0), BATCH_OPTIMIZED,
BATCH_SIZE, MAX_FETCH_RECORDS, BROKER_TOPIC_STATS);
+ 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);
@@ -2416,6 +2417,8 @@ public class DelayedShareFetchTest {
// Invoke tryComplete again to simulate async read completion but
should return false as future
// is not yet completed.
assertFalse(delayedShareFetch.tryComplete());
+ // The result future should still not be completed.
+ assertFalse(shareFetch.isCompleted());
// Complete future to simulate async read completion. The registered
callback should trigger a
// purgatory check rather than completing the request directly on the
completing thread.
asyncReadFuture.complete(buildLogReadResultMap(List.of(tp0)));
@@ -2431,6 +2434,220 @@ public class DelayedShareFetchTest {
// Verify readFromLog was called (which is used by async readAsync
internally)
Mockito.verify(logReader).readAsync(any(), any(), any(), any(),
anyBoolean());
Mockito.verify(sp0).releaseFetchLock(fetchId);
+ // The future of shareFetch completes.
+ assertTrue(shareFetch.isCompleted());
+ assertEquals(Set.of(tp0), future.join().keySet());
+ }
+
+ /**
+ * Tests that when async read encounters an exception, request completes
with error.
+ */
+ @Test
+ public void testAsyncReadExceptionReleasesLocks() {
+ 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));
+
+ Uuid fetchId = Uuid.randomUuid();
+ DelayedShareFetch delayedShareFetch =
DelayedShareFetchBuilder.builder()
+ .withShareFetchData(shareFetch)
+ .withReplicaManager(replicaManager)
+ .withLogReader(logReader)
+ .withMetadataProvider(metadataProvider)
+ .withPartitionMaxBytesStrategy(partitionMaxBytesStrategy)
+ .withSharePartitions(sharePartitions)
+ .withFetchId(fetchId)
+ .build();
+
+ // Initially no pending fetch
+ assertNull(delayedShareFetch.pendingFetch());
+ // First tryComplete triggers async read
+ assertFalse(delayedShareFetch.tryComplete());
+ assertNotNull(delayedShareFetch.pendingFetch());
+ // Complete with exception.
+ asyncReadFuture.completeExceptionally(new RuntimeException("Simulated
read failure"));
+ Mockito.verify(replicaManager).completeDelayedShareFetchRequest(any());
+ // The request should not be completed yet as the purgatory check
(mocked) is a no-op here.
+ assertFalse(delayedShareFetch.isCompleted());
+
+ // Simulate the purgatory re-invoking tryComplete. get() throws, so
the request is force-completed
+ // with an error and the partition locks are released.
+ assertTrue(delayedShareFetch.tryComplete());
+ assertTrue(delayedShareFetch.isCompleted());
+ 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("Simulated read
failure"));
+ 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.
+ */
+ @Test
+ public void testOnCompleteWithPendingAsyncFetch() {
+ String groupId = "grp";
+ Uuid topicId = Uuid.randomUuid();
+ LogReader logReader = mock(LogReader.class);
+ ReplicaManager replicaManager = mock(ReplicaManager.class);
+ PartitionMetadataProvider metadataProvider =
mock(PartitionMetadataProvider.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));
+
+ Uuid fetchId = Uuid.randomUuid();
+ DelayedShareFetch delayedShareFetch =
DelayedShareFetchBuilder.builder()
+ .withShareFetchData(shareFetch)
+ .withLogReader(logReader)
+ .withMetadataProvider(metadataProvider)
+ .withPartitionMaxBytesStrategy(partitionMaxBytesStrategy)
+ .withSharePartitions(sharePartitions)
+ .withReplicaManager(replicaManager)
+ .withFetchId(fetchId)
+ .build();
+
+ // Initially no pending fetch
+ assertNull(delayedShareFetch.pendingFetch());
+ // First tryComplete triggers async read
+ assertFalse(delayedShareFetch.tryComplete());
+ assertNotNull(delayedShareFetch.pendingFetch());
+ // Simulate timeout: onComplete is invoked while the async read is
still in-flight (not done).
+ delayedShareFetch.forceComplete();
+ // As the async read had not completed, the request is completed with
an empty response, the
+ // pending fetch is cleared and the partition locks are released.
+ assertNull(delayedShareFetch.pendingFetch());
+ assertTrue(shareFetch.isCompleted());
+ assertTrue(future.isDone());
+ assertEquals(Set.of(), future.join().keySet());
+ Mockito.verify(sp0).releaseFetchLock(fetchId);
+ // Pending fetch is marked to null.
+ assertNull(delayedShareFetch.pendingFetch());
+
+ // Complete the pending fetch successfully to validate that it doesn't
invoke the callback.
+ asyncReadFuture.complete(buildLogReadResultMap(List.of(tp0)));
+ // Verify the callback handler should not be executed on the request
completion now.
+ Mockito.verify(replicaManager,
times(0)).completeDelayedShareFetchRequest(any());
+ }
+
+ /**
+ * Tests that onComplete handles the case where async fetch is just
completed, on timeout.
+ */
+ @SuppressWarnings("unchecked")
+ @Test
+ public void testOnCompleteWithCompletedAsyncFetch() throws Exception {
+ String groupId = "grp";
+ Uuid topicId = Uuid.randomUuid();
+ LogReader logReader = mock(LogReader.class);
+ ReplicaManager replicaManager = mock(ReplicaManager.class);
+ PartitionMetadataProvider metadataProvider =
mock(PartitionMetadataProvider.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 = mock(CompletableFuture.class);
+ 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));
+
+ // Run false for tryComplete and true for onComplete.
+ when(asyncReadFuture.isDone()).thenReturn(false).thenReturn(true);
+ // Teurn successful response when fetched in onComplete.
+
when(asyncReadFuture.get()).thenReturn(buildLogReadResultMap(List.of(tp0)));
+
+ Uuid fetchId = Uuid.randomUuid();
+ DelayedShareFetch delayedShareFetch =
DelayedShareFetchBuilder.builder()
+ .withShareFetchData(shareFetch)
+ .withLogReader(logReader)
+ .withMetadataProvider(metadataProvider)
+ .withPartitionMaxBytesStrategy(partitionMaxBytesStrategy)
+ .withSharePartitions(sharePartitions)
+ .withReplicaManager(replicaManager)
+ .withFetchId(fetchId)
+ .build();
+
+ // Initially no pending fetch
+ assertNull(delayedShareFetch.pendingFetch());
+ // First tryComplete triggers async read
+ assertFalse(delayedShareFetch.tryComplete());
+ assertNotNull(delayedShareFetch.pendingFetch());
+ try (MockedStatic<ShareFetchUtils> mockedShareFetchUtils =
Mockito.mockStatic(ShareFetchUtils.class, Mockito.CALLS_REAL_METHODS)) {
+ mockedShareFetchUtils.when(
+ () -> ShareFetchUtils.processFetchResponse(any(), any(),
any(), any(), any()))
+ .thenReturn(Map.of(tp0,
mock(ShareFetchResponseData.PartitionData.class)));
+ // Simulate timeout: onComplete is invoked while the async read is
still in-flight (not done).
+ delayedShareFetch.forceComplete();
+ // As the async read has completed now since second invocation to
isDone returns true from
+ // mocked response, the request is completed with data.
+ assertNull(delayedShareFetch.pendingFetch());
+ assertTrue(shareFetch.isCompleted());
+ assertTrue(future.isDone());
+ Map<TopicIdPartition, ShareFetchResponseData.PartitionData>
response = future.join();
+ assertEquals(Set.of(tp0), response.keySet());
+ // Request should be completed without any error. Do not validate
the response as it's returned
+ // from mocked ShareFetchUtils.
+ assertEquals(Errors.NONE.code(), response.get(tp0).errorCode());
+ Mockito.verify(sp0).releaseFetchLock(fetchId);
+ // Pending fetch is marked to null.
+ assertNull(delayedShareFetch.pendingFetch());
+ }
}
static void mockTopicIdPartitionToReturnDataEqualToMinBytes(ReplicaManager
replicaManager, TopicIdPartition topicIdPartition, int minBytes) {
