diff --git a/core/src/main/java/kafka/server/share/SharePartition.java b/core/src/main/java/kafka/server/share/SharePartition.java index 7823b76a80afd..a8b0f881a28ee 100644 --- a/core/src/main/java/kafka/server/share/SharePartition.java +++ b/core/src/main/java/kafka/server/share/SharePartition.java @@ -23,6 +23,7 @@ import org.apache.kafka.common.errors.FencedStateEpochException; import org.apache.kafka.common.errors.InvalidRecordStateException; import org.apache.kafka.common.errors.InvalidRequestException; +import org.apache.kafka.common.errors.LeaderNotAvailableException; import org.apache.kafka.common.errors.NotLeaderOrFollowerException; import org.apache.kafka.common.errors.UnknownServerException; import org.apache.kafka.common.message.ShareFetchResponseData.AcquiredRecords; @@ -38,7 +39,6 @@ import org.apache.kafka.server.group.share.Persister; import org.apache.kafka.server.group.share.PersisterStateBatch; import org.apache.kafka.server.group.share.ReadShareGroupStateParameters; -import org.apache.kafka.server.group.share.ReadShareGroupStateResult; import org.apache.kafka.server.group.share.TopicData; import org.apache.kafka.server.group.share.WriteShareGroupStateParameters; import org.apache.kafka.server.share.ShareAcknowledgementBatch; @@ -59,7 +59,6 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentSkipListMap; -import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; @@ -78,6 +77,31 @@ public class SharePartition { */ static final String EMPTY_MEMBER_ID = Uuid.ZERO_UUID.toString(); + /** + * The SharePartitionState is used to track the state of the share partition. The state of the + * share partition determines if the partition is ready to receive requests, be initialized with + * persisted state, or has failed to initialize. + */ + // Visible for testing + enum SharePartitionState { + /** + * The share partition is empty and has not been initialized with persisted state. + */ + EMPTY, + /** + * The share partition is initializing with persisted state. + */ + INITIALIZING, + /** + * The share partition is active and ready to serve requests. + */ + ACTIVE, + /** + * The share partition failed to initialize with persisted state. + */ + FAILED + } + /** * The RecordState is used to track the state of a record that has been fetched from the leader. * The state of the records determines if the records should be re-delivered, move the next fetch @@ -232,6 +256,11 @@ public static RecordState forId(byte id) { */ private int stateEpoch; + /** + * The partition state is used to track the state of the share partition. + */ + private SharePartitionState partitionState; + SharePartition( String groupId, TopicIdPartition topicIdPartition, @@ -254,8 +283,126 @@ public static RecordState forId(byte id) { this.timer = timer; this.time = time; this.persister = persister; - // Initialize the partition. - initialize(); + this.partitionState = SharePartitionState.EMPTY; + } + + /** + * May initialize the share partition by reading the state from the persister. The share partition + * is initialized only if the state is in the EMPTY state. If the share partition is in ACTIVE state, + * the method completes the future successfully. For other states, the method completes the future + * with exception, which might be re-triable. + * + * @return The method returns a future which is completed when the share partition is initialized + * or completes with an exception if the share partition is in non-initializable state. + */ + public CompletableFuture maybeInitialize() { + log.debug("Maybe initialize share partition: {}-{}", groupId, topicIdPartition); + CompletableFuture future = new CompletableFuture<>(); + // Check if the share partition is already initialized. + maybeCompleteInitialization(future); + if (future.isDone()) { + return future; + } + + // All the pending requests should wait to get completed before the share partition is initialized. + // Attain lock to avoid any concurrent requests to be processed. + lock.writeLock().lock(); + try { + // Re-check the state to verify if previous requests has already initialized the share partition. + maybeCompleteInitialization(future); + if (future.isDone()) { + return future; + } + + // Update state to initializing to avoid any concurrent requests to be processed. + partitionState = SharePartitionState.INITIALIZING; + // Initialize the share partition by reading the state from the persister. + persister.readState(new ReadShareGroupStateParameters.Builder() + .setGroupTopicPartitionData(new GroupTopicPartitionData.Builder() + .setGroupId(this.groupId) + .setTopicsData(Collections.singletonList(new TopicData<>(topicIdPartition.topicId(), + Collections.singletonList(PartitionFactory.newPartitionIdLeaderEpochData(topicIdPartition.partition(), 0))))) + .build()) + .build() + ).whenComplete((result, exception) -> { + if (exception != null) { + log.error("Failed to initialize the share partition: {}-{}", groupId, topicIdPartition, exception); + completeInitializationWithException(future, exception); + return; + } + + if (result == null || result.topicsData() == null || result.topicsData().size() != 1) { + log.error("Failed to initialize the share partition: {}-{}. Invalid state found: {}.", + groupId, topicIdPartition, result); + completeInitializationWithException(future, new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition))); + return; + } + + TopicData state = result.topicsData().get(0); + if (state.topicId() != topicIdPartition.topicId() || state.partitions().size() != 1) { + log.error("Failed to initialize the share partition: {}-{}. Invalid topic partition response: {}.", + groupId, topicIdPartition, result); + completeInitializationWithException(future, new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition))); + return; + } + + PartitionAllData partitionData = state.partitions().get(0); + if (partitionData.partition() != topicIdPartition.partition()) { + log.error("Failed to initialize the share partition: {}-{}. Invalid partition response: {}.", + groupId, topicIdPartition, partitionData); + completeInitializationWithException(future, new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition))); + return; + } + + if (partitionData.errorCode() != Errors.NONE.code()) { + KafkaException ex = fetchPersisterError(partitionData.errorCode(), partitionData.errorMessage()); + log.error("Failed to initialize the share partition: {}-{}. Exception occurred: {}.", + groupId, topicIdPartition, partitionData); + completeInitializationWithException(future, ex); + return; + } + + // Set the state epoch and end offset from the persisted state. + startOffset = partitionData.startOffset() != -1 ? partitionData.startOffset() : 0; + stateEpoch = partitionData.stateEpoch(); + + List stateBatches = partitionData.stateBatches(); + for (PersisterStateBatch stateBatch : stateBatches) { + if (stateBatch.firstOffset() < startOffset) { + log.error("Invalid state batch found for the share partition: {}-{}. The base offset: {}" + + " is less than the start offset: {}.", groupId, topicIdPartition, + stateBatch.firstOffset(), startOffset); + completeInitializationWithException(future, new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition))); + return; + } + InFlightBatch inFlightBatch = new InFlightBatch(EMPTY_MEMBER_ID, stateBatch.firstOffset(), + stateBatch.lastOffset(), RecordState.forId(stateBatch.deliveryState()), stateBatch.deliveryCount(), null); + cachedState.put(stateBatch.firstOffset(), inFlightBatch); + } + // Update the endOffset of the partition. + if (!cachedState.isEmpty()) { + // If the cachedState is not empty, findNextFetchOffset flag is set to true so that any AVAILABLE records + // in the cached state are not missed + findNextFetchOffset.set(true); + endOffset = cachedState.lastEntry().getValue().lastOffset(); + // In case the persister read state RPC result contains no AVAILABLE records, we can update cached state + // and start/end offsets. + maybeUpdateCachedStateAndOffsets(); + } else { + endOffset = partitionData.startOffset(); + } + // Set the partition state to Active and complete the future. + partitionState = SharePartitionState.ACTIVE; + future.complete(null); + }); + } catch (Exception e) { + log.error("Failed to initialize the share partition: {}-{}", groupId, topicIdPartition, e); + completeInitializationWithException(future, e); + } finally { + lock.writeLock().unlock(); + } + + return future; } /** @@ -858,79 +1005,28 @@ void releaseFetchLock() { fetchLock.set(false); } - private void initialize() { - log.debug("Initializing share partition: {}-{}", groupId, topicIdPartition); - // Initialize the share partition by reading the state from the persister. - ReadShareGroupStateResult response; - try { - response = persister.readState(new ReadShareGroupStateParameters.Builder() - .setGroupTopicPartitionData(new GroupTopicPartitionData.Builder() - .setGroupId(this.groupId) - .setTopicsData(Collections.singletonList(new TopicData<>(topicIdPartition.topicId(), - Collections.singletonList(PartitionFactory.newPartitionIdLeaderEpochData(topicIdPartition.partition(), 0))))) - .build()) - .build() - ).get(); - } catch (InterruptedException | ExecutionException e) { - log.error("Failed to initialize the share partition: {}-{}", groupId, topicIdPartition, e); - throw new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition), e); - } - - if (response == null || response.topicsData() == null || response.topicsData().size() != 1) { - log.error("Failed to initialize the share partition: {}-{}. Invalid state found: {}.", - groupId, topicIdPartition, response); - throw new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition)); - } - - TopicData state = response.topicsData().get(0); - if (state.topicId() != topicIdPartition.topicId() || state.partitions().size() != 1) { - log.error("Failed to initialize the share partition: {}-{}. Invalid topic partition response: {}.", - groupId, topicIdPartition, response); - throw new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition)); - } - - PartitionAllData partitionData = state.partitions().get(0); - if (partitionData.partition() != topicIdPartition.partition()) { - log.error("Failed to initialize the share partition: {}-{}. Invalid partition response: {}.", - groupId, topicIdPartition, partitionData); - throw new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", - groupId, topicIdPartition)); - } - - if (partitionData.errorCode() != Errors.NONE.code()) { - KafkaException ex = fetchPersisterError(partitionData.errorCode(), partitionData.errorMessage()); - log.error("Failed to initialize the share partition: {}-{}. Exception occurred: {}.", - groupId, topicIdPartition, partitionData); - throw ex; - } - - // Set the state epoch and end offset from the persisted state. - startOffset = partitionData.startOffset() != -1 ? partitionData.startOffset() : 0; - stateEpoch = partitionData.stateEpoch(); + private void completeInitializationWithException(CompletableFuture future, Throwable exception) { + partitionState = SharePartitionState.FAILED; + future.completeExceptionally(exception); + } - List stateBatches = partitionData.stateBatches(); - for (PersisterStateBatch stateBatch : stateBatches) { - if (stateBatch.firstOffset() < startOffset) { - log.error("Invalid state batch found for the share partition: {}-{}. The base offset: {}" - + " is less than the start offset: {}.", groupId, topicIdPartition, - stateBatch.firstOffset(), startOffset); - throw new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition)); - } - InFlightBatch inFlightBatch = new InFlightBatch(EMPTY_MEMBER_ID, stateBatch.firstOffset(), - stateBatch.lastOffset(), RecordState.forId(stateBatch.deliveryState()), stateBatch.deliveryCount(), null); - cachedState.put(stateBatch.firstOffset(), inFlightBatch); - } - // Update the endOffset of the partition. - if (!cachedState.isEmpty()) { - // If the cachedState is not empty, findNextFetchOffset flag is set to true so that any AVAILABLE records - // in the cached state are not missed - findNextFetchOffset.set(true); - endOffset = cachedState.lastEntry().getValue().lastOffset(); - // In case the persister read state RPC result contains no AVAILABLE records, we can update cached state - // and start/end offsets. - maybeUpdateCachedStateAndOffsets(); - } else { - endOffset = partitionData.startOffset(); + private void maybeCompleteInitialization(CompletableFuture future) { + SharePartitionState currentState = partitionState(); + switch (currentState) { + case ACTIVE: + future.complete(null); + return; + case FAILED: + future.completeExceptionally(new IllegalStateException(String.format("Share partition failed to load %s-%s", groupId, topicIdPartition))); + return; + case INITIALIZING: + future.completeExceptionally(new LeaderNotAvailableException(String.format("Share partition is already initializing %s-%s", groupId, topicIdPartition))); + return; + case EMPTY: + // Do not complete the future as the share partition is not yet initialized. + break; + default: + throw new IllegalStateException("Unknown share partition state: " + currentState); } } @@ -1352,6 +1448,16 @@ private Optional acknowledgeCompleteBatch( return Optional.empty(); } + // Visible for testing + SharePartitionState partitionState() { + lock.readLock().lock(); + try { + return partitionState; + } finally { + lock.readLock().unlock(); + } + } + // Visible for testing void rollbackOrProcessStateUpdates( CompletableFuture future, diff --git a/core/src/main/java/kafka/server/share/SharePartitionManager.java b/core/src/main/java/kafka/server/share/SharePartitionManager.java index 2f6ee6b999d3d..c14c1ea15c35c 100644 --- a/core/src/main/java/kafka/server/share/SharePartitionManager.java +++ b/core/src/main/java/kafka/server/share/SharePartitionManager.java @@ -23,6 +23,9 @@ import org.apache.kafka.common.TopicIdPartition; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.errors.FencedStateEpochException; +import org.apache.kafka.common.errors.LeaderNotAvailableException; +import org.apache.kafka.common.errors.NotLeaderOrFollowerException; import org.apache.kafka.common.message.ShareAcknowledgeResponseData; import org.apache.kafka.common.message.ShareFetchResponseData.PartitionData; import org.apache.kafka.common.metrics.Metrics; @@ -546,42 +549,49 @@ void maybeProcessFetchQueue() { shareFetchPartitionData.groupId, topicIdPartition ); - SharePartition sharePartition = partitionCacheMap.computeIfAbsent(sharePartitionKey, - k -> { - long start = time.hiResClockMs(); - SharePartition partition = new SharePartition(shareFetchPartitionData.groupId, topicIdPartition, maxInFlightMessages, maxDeliveryCount, - recordLockDurationMs, timer, time, persister); - this.shareGroupMetrics.partitionLoadTime(start); - return partition; - }); - int partitionMaxBytes = shareFetchPartitionData.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()) { - // If the share partition is already at capacity, we should not attempt to fetch. - if (sharePartition.canAcquireRecords()) { - topicPartitionData.put( - topicIdPartition, - new FetchRequest.PartitionData( - topicIdPartition.topicId(), - sharePartition.nextFetchOffset(), - 0, - partitionMaxBytes, - Optional.empty() - ) - ); - } else { - sharePartition.releaseFetchLock(); - log.info("Record lock partition limit exceeded for SharePartition with key {}, " + - "cannot acquire more records", sharePartitionKey); + SharePartition sharePartition = fetchSharePartition(sharePartitionKey); + + // The share partition is initialized asynchronously, so we need to wait for it to be initialized. + // But if the share partition is already initialized, then the future will be completed immediately. + // Hence, it's safe to call the maybeInitialize method and then wait for the future to be completed. + // TopicPartitionData list will be populated only if the share partition is already initialized. + sharePartition.maybeInitialize().whenComplete((result, throwable) -> { + if (throwable != null) { + maybeCompleteInitializationWithException(sharePartitionKey, shareFetchPartitionData.future, throwable); + return; } - } + // Fetch messages for the partition only if it is active. + int partitionMaxBytes = shareFetchPartitionData.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()) { + // If the share partition is already at capacity, we should not attempt to fetch. + if (sharePartition.canAcquireRecords()) { + topicPartitionData.put( + topicIdPartition, + new FetchRequest.PartitionData( + topicIdPartition.topicId(), + sharePartition.nextFetchOffset(), + 0, + partitionMaxBytes, + Optional.empty() + ) + ); + } else { + sharePartition.releaseFetchLock(); + log.info("Record lock partition limit exceeded for SharePartition with key {}, " + + "cannot acquire more records", sharePartitionKey); + } + } + }); }); if (topicPartitionData.isEmpty()) { - // No locks for share partitions could be acquired, so we complete the request and - // will re-fetch for the client in next poll. - shareFetchPartitionData.future.complete(Collections.emptyMap()); + // No locks for share partitions could be acquired, so we complete the request if not + // already. The records shall be re-fetched in next poll by the client. + if (!shareFetchPartitionData.future.isDone()) { + shareFetchPartitionData.future.complete(Collections.emptyMap()); + } // Though if no partitions can be locked then there must be some other request which // is in-flight and should release the lock. But it's safe to release the lock as // the lock on share partition already exists which facilitates correct behaviour @@ -595,30 +605,7 @@ void maybeProcessFetchQueue() { log.trace("Fetchable share partitions data: {} with groupId: {} fetch params: {}", topicPartitionData, shareFetchPartitionData.groupId, shareFetchPartitionData.fetchParams); - replicaManager.fetchMessages( - shareFetchPartitionData.fetchParams, - CollectionConverters.asScala( - topicPartitionData.entrySet().stream().map(entry -> - new Tuple2<>(entry.getKey(), entry.getValue())).collect(Collectors.toList()) - ), - QuotaFactory.UnboundedQuota$.MODULE$, - responsePartitionData -> { - log.trace("Data successfully retrieved by replica manager: {}", responsePartitionData); - List> responseData = CollectionConverters.asJava( - responsePartitionData); - processFetchResponse(shareFetchPartitionData, responseData).whenComplete( - (result, throwable) -> { - if (throwable != null) { - log.error("Error processing fetch response for share partitions", throwable); - shareFetchPartitionData.future.completeExceptionally(throwable); - } else { - shareFetchPartitionData.future.complete(result); - } - // Releasing the lock to move ahead with the next request in queue. - releaseFetchQueueAndPartitionsLock(shareFetchPartitionData.groupId, topicPartitionData.keySet()); - }); - return BoxedUnit.UNIT; - }); + processReplicaManagerFetch(shareFetchPartitionData, topicPartitionData); // If there are more requests in the queue, then process them. if (!fetchQueue.isEmpty()) @@ -632,6 +619,85 @@ void maybeProcessFetchQueue() { } } + private SharePartition fetchSharePartition(SharePartitionKey sharePartitionKey) { + return partitionCacheMap.computeIfAbsent(sharePartitionKey, + k -> { + long start = time.hiResClockMs(); + SharePartition partition = new SharePartition( + sharePartitionKey.groupId, + sharePartitionKey.topicIdPartition, + maxInFlightMessages, + maxDeliveryCount, + recordLockDurationMs, + timer, + time, + persister + ); + this.shareGroupMetrics.partitionLoadTime(start); + return partition; + }); + } + + private void maybeCompleteInitializationWithException( + SharePartitionKey sharePartitionKey, + CompletableFuture> future, + Throwable throwable) { + if (throwable instanceof LeaderNotAvailableException) { + log.debug("The share partition with key {} is not initialized yet", sharePartitionKey); + // Do not process the fetch request for this partition as the leader is not initialized yet. + // The fetch request will be retried in the next poll. + // TODO: Add the request to delayed fetch purgatory. + return; + } + + if (throwable instanceof NotLeaderOrFollowerException || throwable instanceof FencedStateEpochException) { + log.info("The share partition with key {} is fenced: {}", sharePartitionKey, throwable.getMessage()); + // The share partition is fenced hence remove the partition from map and let the client retry. + // But surface the error to the client so client might take some action i.e. re-fetch + // the metadata and retry the fetch on new leader. + partitionCacheMap.remove(sharePartitionKey); + future.completeExceptionally(throwable); + return; + } + + // The partition initialization failed, so complete the request with the exception. + // The server should not be in this state, so log the error on broker and surface the same + // to the client. As of now this state is in-recoverable for the broker, and we should + // investigate the root cause of the error. + log.error("Error initializing share partition with key {}", sharePartitionKey, throwable); + future.completeExceptionally(throwable); + } + + private void processReplicaManagerFetch( + ShareFetchPartitionData shareFetchPartitionData, + Map topicPartitionData + ) { + replicaManager.fetchMessages( + shareFetchPartitionData.fetchParams, + CollectionConverters.asScala( + topicPartitionData.entrySet().stream().map(entry -> + new Tuple2<>(entry.getKey(), entry.getValue())).collect(Collectors.toList()) + ), + QuotaFactory.UnboundedQuota$.MODULE$, + responsePartitionData -> { + log.trace("Data successfully retrieved by replica manager: {}", responsePartitionData); + List> responseData = CollectionConverters.asJava( + responsePartitionData); + processFetchResponse(shareFetchPartitionData, responseData).whenComplete( + (result, throwable) -> { + if (throwable != null) { + log.error("Error processing fetch response for share partitions", throwable); + shareFetchPartitionData.future.completeExceptionally(throwable); + } else { + shareFetchPartitionData.future.complete(result); + } + // Releasing the lock to move ahead with the next request in queue. + releaseFetchQueueAndPartitionsLock(shareFetchPartitionData.groupId, topicPartitionData.keySet()); + }); + return BoxedUnit.UNIT; + }); + } + // Visible for testing. CompletableFuture> processFetchResponse( ShareFetchPartitionData shareFetchPartitionData, diff --git a/core/src/test/java/kafka/server/share/SharePartitionManagerTest.java b/core/src/test/java/kafka/server/share/SharePartitionManagerTest.java index 3d88a8aa5b03b..5bf0412fd3c6c 100644 --- a/core/src/test/java/kafka/server/share/SharePartitionManagerTest.java +++ b/core/src/test/java/kafka/server/share/SharePartitionManagerTest.java @@ -26,9 +26,13 @@ import org.apache.kafka.common.Uuid; import org.apache.kafka.common.compress.Compression; import org.apache.kafka.common.errors.BrokerNotAvailableException; +import org.apache.kafka.common.errors.CoordinatorNotAvailableException; +import org.apache.kafka.common.errors.FencedStateEpochException; import org.apache.kafka.common.errors.InvalidRecordStateException; import org.apache.kafka.common.errors.InvalidRequestException; import org.apache.kafka.common.errors.InvalidShareSessionEpochException; +import org.apache.kafka.common.errors.LeaderNotAvailableException; +import org.apache.kafka.common.errors.NotLeaderOrFollowerException; import org.apache.kafka.common.errors.ShareSessionNotFoundException; import org.apache.kafka.common.message.ShareAcknowledgeResponseData; import org.apache.kafka.common.message.ShareFetchResponseData; @@ -1299,6 +1303,7 @@ public void testReplicaManagerFetchShouldNotProceed() { SharePartition sp0 = mock(SharePartition.class); when(sp0.maybeAcquireFetchLock()).thenReturn(true); when(sp0.canAcquireRecords()).thenReturn(false); + when(sp0.maybeInitialize()).thenReturn(CompletableFuture.completedFuture(null)); Map partitionCacheMap = new HashMap<>(); partitionCacheMap.put(new SharePartitionManager.SharePartitionKey(groupId, tp0), sp0); @@ -1329,6 +1334,7 @@ public void testReplicaManagerFetchShouldProceed() { SharePartition sp0 = mock(SharePartition.class); when(sp0.maybeAcquireFetchLock()).thenReturn(true); when(sp0.canAcquireRecords()).thenReturn(true); + when(sp0.maybeInitialize()).thenReturn(CompletableFuture.completedFuture(null)); Map partitionCacheMap = new HashMap<>(); partitionCacheMap.put(new SharePartitionManager.SharePartitionKey(groupId, tp0), sp0); @@ -1892,6 +1898,122 @@ public void testFetchQueueProcessingWhenFrontItemIsEmpty() { verify(replicaManager, times(1)).fetchMessages(any(), any(), any(ReplicaQuota.class), any()); } + @Test + public void testPendingInitializationShouldCompleteFetchRequest() throws Exception { + String groupId = "grp"; + Uuid memberId = Uuid.randomUuid(); + FetchParams fetchParams = new FetchParams(ApiKeys.SHARE_FETCH.latestVersion(), FetchRequest.ORDINARY_CONSUMER_ID, -1, 0, + 1, 1024 * 1024, FetchIsolation.HIGH_WATERMARK, Optional.empty()); + Uuid fooId = Uuid.randomUuid(); + TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0)); + Map partitionMaxBytes = Collections.singletonMap(tp0, PARTITION_MAX_BYTES); + + SharePartition sp0 = mock(SharePartition.class); + Map partitionCacheMap = new HashMap<>(); + partitionCacheMap.put(new SharePartitionManager.SharePartitionKey(groupId, tp0), sp0); + + // Keep the initialization future pending, so fetch request is stuck. + CompletableFuture pendingInitializationFuture = new CompletableFuture<>(); + when(sp0.maybeInitialize()).thenReturn(pendingInitializationFuture); + + // Mock replica manager to verify no calls are made to fetchMessages. + ReplicaManager replicaManager = mock(ReplicaManager.class); + + SharePartitionManager sharePartitionManager = SharePartitionManagerBuilder.builder() + .withPartitionCacheMap(partitionCacheMap).withReplicaManager(replicaManager).build(); + + CompletableFuture> future = + sharePartitionManager.fetchMessages(groupId, memberId.toString(), fetchParams, partitionMaxBytes); + // Verify that the fetch request is completed. + assertTrue(future.isDone()); + assertTrue(future.join().isEmpty()); + // Verify that replica manager fetch is not called. + Mockito.verify(replicaManager, times(0)).fetchMessages( + any(), any(), any(ReplicaQuota.class), any()); + // Complete the pending initialization future. + pendingInitializationFuture.complete(null); + } + + @Test + public void testSharePartitionInitializationExceptions() throws Exception { + String groupId = "grp"; + Uuid memberId = Uuid.randomUuid(); + FetchParams fetchParams = new FetchParams(ApiKeys.SHARE_FETCH.latestVersion(), FetchRequest.ORDINARY_CONSUMER_ID, -1, 0, + 1, 1024 * 1024, FetchIsolation.HIGH_WATERMARK, Optional.empty()); + Uuid fooId = Uuid.randomUuid(); + TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0)); + Map partitionMaxBytes = Collections.singletonMap(tp0, PARTITION_MAX_BYTES); + + SharePartition sp0 = mock(SharePartition.class); + Map partitionCacheMap = new HashMap<>(); + partitionCacheMap.put(new SharePartitionManager.SharePartitionKey(groupId, tp0), sp0); + + SharePartitionManager sharePartitionManager = SharePartitionManagerBuilder.builder() + .withPartitionCacheMap(partitionCacheMap).build(); + + // Return LeaderNotAvailableException to simulate initialization failure. + when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new LeaderNotAvailableException("Leader not available"))); + CompletableFuture> future = + sharePartitionManager.fetchMessages(groupId, memberId.toString(), fetchParams, partitionMaxBytes); + assertTrue(future.isDone()); + // Exception for client should not occur for LeaderNotAvailableException, this exception is to communicate + // between SharePartitionManager and SharePartition to retry the request as SharePartition is not yet ready. + assertFalse(future.isCompletedExceptionally()); + assertTrue(future.join().isEmpty()); + + // Return IllegalStateException to simulate initialization failure. + when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new IllegalStateException("Illegal state"))); + future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), fetchParams, partitionMaxBytes); + assertTrue(future.isDone()); + assertTrue(future.isCompletedExceptionally()); + assertFutureThrows(future, IllegalStateException.class); + + // Return CoordinatorNotAvailableException to simulate initialization failure. + when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new CoordinatorNotAvailableException("Coordinator not available"))); + future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), fetchParams, partitionMaxBytes); + assertTrue(future.isDone()); + assertTrue(future.isCompletedExceptionally()); + assertFutureThrows(future, CoordinatorNotAvailableException.class); + + // Return InvalidRequestException to simulate initialization failure. + when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new InvalidRequestException("Invalid request"))); + future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), fetchParams, partitionMaxBytes); + assertTrue(future.isDone()); + assertTrue(future.isCompletedExceptionally()); + assertFutureThrows(future, InvalidRequestException.class); + + // Return FencedStateEpochException to simulate initialization failure. + when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new FencedStateEpochException("Fenced state epoch"))); + // Assert that partitionCacheMap contains instance before the fetch request. + assertEquals(1, partitionCacheMap.size()); + future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), fetchParams, partitionMaxBytes); + assertTrue(future.isDone()); + assertTrue(future.isCompletedExceptionally()); + assertFutureThrows(future, FencedStateEpochException.class); + // Verify that the share partition is removed from the cache. + assertTrue(partitionCacheMap.isEmpty()); + + // The last exception removes the share partition from the cache hence re-add the share partition to cache. + partitionCacheMap.put(new SharePartitionManager.SharePartitionKey(groupId, tp0), sp0); + // Return NotLeaderOrFollowerException to simulate initialization failure. + when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new NotLeaderOrFollowerException("Not leader or follower"))); + future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), fetchParams, partitionMaxBytes); + assertTrue(future.isDone()); + assertTrue(future.isCompletedExceptionally()); + assertFutureThrows(future, NotLeaderOrFollowerException.class); + // Verify that the share partition is removed from the cache. + assertTrue(partitionCacheMap.isEmpty()); + + // The last exception removes the share partition from the cache hence re-add the share partition to cache. + partitionCacheMap.put(new SharePartitionManager.SharePartitionKey(groupId, tp0), sp0); + // Return RuntimeException to simulate initialization failure. + when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new RuntimeException("Runtime exception"))); + future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), fetchParams, partitionMaxBytes); + assertTrue(future.isDone()); + assertTrue(future.isCompletedExceptionally()); + assertFutureThrows(future, RuntimeException.class); + } + private ShareFetchResponseData.PartitionData noErrorShareFetchResponse() { return new ShareFetchResponseData.PartitionData().setPartitionIndex(0); } diff --git a/core/src/test/java/kafka/server/share/SharePartitionTest.java b/core/src/test/java/kafka/server/share/SharePartitionTest.java index 9754b31c16566..51405f4a1fc96 100644 --- a/core/src/test/java/kafka/server/share/SharePartitionTest.java +++ b/core/src/test/java/kafka/server/share/SharePartitionTest.java @@ -18,6 +18,7 @@ import kafka.server.share.SharePartition.InFlightState; import kafka.server.share.SharePartition.RecordState; +import kafka.server.share.SharePartition.SharePartitionState; import org.apache.kafka.common.TopicIdPartition; import org.apache.kafka.common.Uuid; @@ -66,6 +67,9 @@ import java.util.OptionalInt; import java.util.OptionalLong; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import static kafka.server.share.SharePartition.EMPTY_MEMBER_ID; import static org.apache.kafka.test.TestUtils.assertFutureThrows; @@ -141,7 +145,7 @@ public void testRecordStateForId() { } @Test - public void testInitialize() { + public void testMaybeInitialize() { Persister persister = Mockito.mock(Persister.class); ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -153,6 +157,11 @@ public void testInitialize() { Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertFalse(result.isCompletedExceptionally()); + + assertEquals(SharePartitionState.ACTIVE, sharePartition.partitionState()); assertFalse(sharePartition.cachedState().isEmpty()); assertEquals(5, sharePartition.startOffset()); assertEquals(15, sharePartition.endOffset()); @@ -175,7 +184,73 @@ public void testInitialize() { } @Test - public void testInitializeWithEmptyStateBatches() { + public void testMaybeInitializeSharePartitionAgain() { + Persister persister = Mockito.mock(Persister.class); + ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); + Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( + new TopicData<>(TOPIC_ID_PARTITION.topicId(), Collections.singletonList( + PartitionFactory.newPartitionAllData(0, 3, 5L, Errors.NONE.code(), Errors.NONE.message(), + Arrays.asList( + new PersisterStateBatch(5L, 10L, RecordState.AVAILABLE.id, (short) 2), + new PersisterStateBatch(11L, 15L, RecordState.ARCHIVED.id, (short) 3))))))); + Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertFalse(result.isCompletedExceptionally()); + assertEquals(SharePartitionState.ACTIVE, sharePartition.partitionState()); + + // Initialize again, no need to send mock persister response again as the state is already initialized. + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertFalse(result.isCompletedExceptionally()); + assertEquals(SharePartitionState.ACTIVE, sharePartition.partitionState()); + + // Verify the persister read state is called only once. + Mockito.verify(persister, Mockito.times(1)).readState(Mockito.any()); + } + + @Test + public void testMaybeInitializeSharePartitionAgainConcurrentRequests() throws InterruptedException { + Persister persister = Mockito.mock(Persister.class); + ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); + Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( + new TopicData<>(TOPIC_ID_PARTITION.topicId(), Collections.singletonList( + PartitionFactory.newPartitionAllData(0, 3, 5L, Errors.NONE.code(), Errors.NONE.message(), + Arrays.asList( + new PersisterStateBatch(5L, 10L, RecordState.AVAILABLE.id, (short) 2), + new PersisterStateBatch(11L, 15L, RecordState.ARCHIVED.id, (short) 3))))))); + Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + // No need to send mock persister response again as only 1 thread should read state from persister. + ExecutorService executorService = Executors.newFixedThreadPool(10); + List> results = new ArrayList<>(10); + + try { + for (int i = 0; i < 10; i++) { + executorService.submit(() -> { + results.add(sharePartition.maybeInitialize()); + }); + } + } finally { + if (!executorService.awaitTermination(30, TimeUnit.MILLISECONDS)) + executorService.shutdown(); + } + + for (CompletableFuture result : results) { + assertTrue(result.isDone()); + assertFalse(result.isCompletedExceptionally()); + } + + assertEquals(SharePartitionState.ACTIVE, sharePartition.partitionState()); + // Verify the persister read state is called only once. + Mockito.verify(persister, Mockito.times(1)).readState(Mockito.any()); + } + + @Test + public void testMaybeInitializeWithEmptyStateBatches() { Persister persister = Mockito.mock(Persister.class); ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -185,6 +260,11 @@ public void testInitializeWithEmptyStateBatches() { Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertFalse(result.isCompletedExceptionally()); + + assertEquals(SharePartitionState.ACTIVE, sharePartition.partitionState()); assertTrue(sharePartition.cachedState().isEmpty()); assertEquals(10, sharePartition.startOffset()); assertEquals(10, sharePartition.endOffset()); @@ -193,7 +273,7 @@ public void testInitializeWithEmptyStateBatches() { } @Test - public void testInitializeWithErrorPartitionResponse() { + public void testMaybeInitializeWithErrorPartitionResponse() { Persister persister = Mockito.mock(Persister.class); ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); @@ -203,7 +283,13 @@ public void testInitializeWithErrorPartitionResponse() { PartitionFactory.newPartitionAllData(0, 5, 10L, Errors.NOT_COORDINATOR.code(), Errors.NOT_COORDINATOR.message(), Collections.emptyList()))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(CoordinatorNotAvailableException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, CoordinatorNotAvailableException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); // Mock COORDINATOR_NOT_AVAILABLE error. Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -211,7 +297,13 @@ public void testInitializeWithErrorPartitionResponse() { PartitionFactory.newPartitionAllData(0, 5, 10L, Errors.COORDINATOR_NOT_AVAILABLE.code(), Errors.COORDINATOR_NOT_AVAILABLE.message(), Collections.emptyList()))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(CoordinatorNotAvailableException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, CoordinatorNotAvailableException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); // Mock COORDINATOR_LOAD_IN_PROGRESS error. Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -219,7 +311,13 @@ public void testInitializeWithErrorPartitionResponse() { PartitionFactory.newPartitionAllData(0, 5, 10L, Errors.COORDINATOR_LOAD_IN_PROGRESS.code(), Errors.COORDINATOR_LOAD_IN_PROGRESS.message(), Collections.emptyList()))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(CoordinatorNotAvailableException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, CoordinatorNotAvailableException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); // Mock GROUP_ID_NOT_FOUND error. Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -227,7 +325,13 @@ public void testInitializeWithErrorPartitionResponse() { PartitionFactory.newPartitionAllData(0, 5, 10L, Errors.GROUP_ID_NOT_FOUND.code(), Errors.GROUP_ID_NOT_FOUND.message(), Collections.emptyList()))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(InvalidRequestException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, InvalidRequestException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); // Mock UNKNOWN_TOPIC_OR_PARTITION error. Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -235,7 +339,13 @@ public void testInitializeWithErrorPartitionResponse() { PartitionFactory.newPartitionAllData(0, 5, 10L, Errors.UNKNOWN_TOPIC_OR_PARTITION.code(), Errors.UNKNOWN_TOPIC_OR_PARTITION.message(), Collections.emptyList()))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(InvalidRequestException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, InvalidRequestException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); // Mock FENCED_STATE_EPOCH error. Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -243,7 +353,13 @@ public void testInitializeWithErrorPartitionResponse() { PartitionFactory.newPartitionAllData(0, 5, 10L, Errors.FENCED_STATE_EPOCH.code(), Errors.FENCED_STATE_EPOCH.message(), Collections.emptyList()))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(FencedStateEpochException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, FencedStateEpochException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); // Mock FENCED_LEADER_EPOCH error. Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -251,7 +367,13 @@ public void testInitializeWithErrorPartitionResponse() { PartitionFactory.newPartitionAllData(0, 5, 10L, Errors.FENCED_LEADER_EPOCH.code(), Errors.FENCED_LEADER_EPOCH.message(), Collections.emptyList()))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(NotLeaderOrFollowerException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, NotLeaderOrFollowerException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); // Mock UNKNOWN_SERVER_ERROR error. Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -259,11 +381,17 @@ public void testInitializeWithErrorPartitionResponse() { PartitionFactory.newPartitionAllData(0, 5, 10L, Errors.UNKNOWN_SERVER_ERROR.code(), Errors.UNKNOWN_SERVER_ERROR.message(), Collections.emptyList()))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(UnknownServerException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, UnknownServerException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); } @Test - public void testInitializeWithInvalidStartOffsetStateBatches() { + public void testMaybeInitializeWithInvalidStartOffsetStateBatches() { Persister persister = Mockito.mock(Persister.class); ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -273,11 +401,17 @@ public void testInitializeWithInvalidStartOffsetStateBatches() { new PersisterStateBatch(5L, 10L, RecordState.AVAILABLE.id, (short) 2), new PersisterStateBatch(11L, 15L, RecordState.ARCHIVED.id, (short) 3))))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(IllegalStateException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, IllegalStateException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); } @Test - public void testInitializeWithInvalidTopicIdResponse() { + public void testMaybeInitializeWithInvalidTopicIdResponse() { Persister persister = Mockito.mock(Persister.class); ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -287,11 +421,17 @@ public void testInitializeWithInvalidTopicIdResponse() { new PersisterStateBatch(5L, 10L, RecordState.AVAILABLE.id, (short) 2), new PersisterStateBatch(11L, 15L, RecordState.ARCHIVED.id, (short) 3))))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(IllegalStateException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, IllegalStateException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); } @Test - public void testInitializeWithInvalidPartitionResponse() { + public void testMaybeInitializeWithInvalidPartitionResponse() { Persister persister = Mockito.mock(Persister.class); ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.singletonList( @@ -301,50 +441,96 @@ public void testInitializeWithInvalidPartitionResponse() { new PersisterStateBatch(5L, 10L, RecordState.AVAILABLE.id, (short) 2), new PersisterStateBatch(11L, 15L, RecordState.ARCHIVED.id, (short) 3))))))); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(IllegalStateException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, IllegalStateException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); } @Test - public void testInitializeWithNoOpShareStatePersister() { + public void testMaybeInitializeWithNoOpShareStatePersister() { SharePartition sharePartition = SharePartitionBuilder.builder().build(); + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertFalse(result.isCompletedExceptionally()); assertTrue(sharePartition.cachedState().isEmpty()); assertEquals(0, sharePartition.startOffset()); assertEquals(0, sharePartition.endOffset()); assertEquals(0, sharePartition.stateEpoch()); assertEquals(0, sharePartition.nextFetchOffset()); + assertEquals(SharePartitionState.ACTIVE, sharePartition.partitionState()); } @Test - public void testInitializeWithNullResponse() { + public void testMaybeInitializeWithNullResponse() { Persister persister = Mockito.mock(Persister.class); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(null)); - assertThrows(IllegalStateException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, IllegalStateException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); } @Test - public void testInitializeWithNullTopicsData() { + public void testMaybeInitializeWithNullTopicsData() { Persister persister = Mockito.mock(Persister.class); ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(null); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(IllegalStateException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, IllegalStateException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); } @Test - public void testInitializeWithEmptyTopicsData() { + public void testMaybeInitializeWithEmptyTopicsData() { Persister persister = Mockito.mock(Persister.class); ReadShareGroupStateResult readShareGroupStateResult = Mockito.mock(ReadShareGroupStateResult.class); Mockito.when(readShareGroupStateResult.topicsData()).thenReturn(Collections.emptyList()); Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); - assertThrows(IllegalStateException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, IllegalStateException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); } @Test - public void testInitializeWithReadException() { + public void testMaybeInitializeWithReadException() { Persister persister = Mockito.mock(Persister.class); + // Complete the future exceptionally for read state. Mockito.when(persister.readState(Mockito.any())).thenReturn(FutureUtils.failedFuture(new RuntimeException("Read exception"))); - assertThrows(IllegalStateException.class, () -> SharePartitionBuilder.builder().withPersister(persister).build()); + SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, RuntimeException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); + + persister = Mockito.mock(Persister.class); + // Throw exception for read state. + Mockito.when(persister.readState(Mockito.any())).thenThrow(new RuntimeException("Read exception")); + sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + + result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertTrue(result.isCompletedExceptionally()); + assertFutureThrows(result, RuntimeException.class); + assertEquals(SharePartitionState.FAILED, sharePartition.partitionState()); } @Test @@ -4660,7 +4846,7 @@ public void testReacquireSubsetWithAnotherMember() { } @Test - public void testInitializeWhenReadStateRpcReturnsZeroAvailableRecords() { + public void testMaybeInitializeWhenReadStateRpcReturnsZeroAvailableRecords() { List stateBatches = new ArrayList<>(); for (int i = 0; i < 500; i++) { stateBatches.add(new PersisterStateBatch(234L + i, 234L + i, RecordState.ACKNOWLEDGED.id, (short) 1)); @@ -4676,6 +4862,10 @@ public void testInitializeWhenReadStateRpcReturnsZeroAvailableRecords() { Mockito.when(persister.readState(Mockito.any())).thenReturn(CompletableFuture.completedFuture(readShareGroupStateResult)); SharePartition sharePartition = SharePartitionBuilder.builder().withPersister(persister).build(); + CompletableFuture result = sharePartition.maybeInitialize(); + assertTrue(result.isDone()); + assertFalse(result.isCompletedExceptionally()); + assertTrue(sharePartition.cachedState().isEmpty()); assertEquals(734, sharePartition.nextFetchOffset()); assertEquals(734, sharePartition.startOffset());