diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractMembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractMembershipManager.java index c6aa70d805e0c..4a7ec8ade0eb8 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractMembershipManager.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractMembershipManager.java @@ -650,7 +650,7 @@ public void transitionToSendingLeaveGroup(boolean dueToExpiredPollTimer) { * the group, this will be invoked with empty epoch. */ void notifyEpochChange(Optional epoch) { - stateUpdatesListeners.forEach(stateListener -> stateListener.onMemberEpochUpdated(epoch, memberId)); + stateUpdatesListeners.forEach(stateListener -> stateListener.onMemberEpochUpdated(epoch, Optional.of(memberId))); } /** diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java index 34daa6d5b2ccf..f96a6319d5f58 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java @@ -266,7 +266,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) { private final MemberStateListener memberStateListener = new MemberStateListener() { @Override - public void onMemberEpochUpdated(Optional memberEpoch, String memberId) { + public void onMemberEpochUpdated(Optional memberEpoch, Optional memberId) { updateGroupMetadata(memberEpoch, memberId); } @@ -387,7 +387,9 @@ public AsyncKafkaConsumer(final ConsumerConfig config, applicationEventProcessorSupplier, networkClientDelegateSupplier, requestManagersSupplier); - + streamsAssignmentInterface.ifPresent( + sai -> sai.setApplicationEventHandler(applicationEventHandler) + ); this.rebalanceListenerInvoker = new ConsumerRebalanceListenerInvoker( logContext, subscriptions, @@ -477,8 +479,7 @@ public AsyncKafkaConsumer(final ConsumerConfig config, Deserializer valueDeserializer, KafkaClient client, SubscriptionState subscriptions, - ConsumerMetadata metadata, - Optional streamsInstanceMetadata) { + ConsumerMetadata metadata) { this.log = logContext.logger(getClass()); this.subscriptions = subscriptions; this.clientId = config.getString(ConsumerConfig.CLIENT_ID_CONFIG); @@ -546,8 +547,8 @@ public AsyncKafkaConsumer(final ConsumerConfig config, clientTelemetryReporter, metrics, offsetCommitCallbackInvoker, - memberStateListener, - streamsInstanceMetadata + this::updateGroupMetadata, + Optional.empty() ); Supplier applicationEventProcessorSupplier = ApplicationEventProcessor.supplier( logContext, @@ -651,13 +652,13 @@ private ConsumerGroupMetadata initializeConsumerGroupMetadata(final String group ); } - private void updateGroupMetadata(final Optional memberEpoch, final String memberId) { + private void updateGroupMetadata(final Optional memberEpoch, final Optional memberId) { memberEpoch.ifPresent(epoch -> groupMetadata.updateAndGet( oldGroupMetadataOptional -> oldGroupMetadataOptional.map( oldGroupMetadata -> new ConsumerGroupMetadata( oldGroupMetadata.groupId(), memberEpoch.orElse(oldGroupMetadata.generationId()), - memberId, + memberId.orElse(oldGroupMetadata.memberId()), oldGroupMetadata.groupInstanceId() ) ) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerDelegateCreator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerDelegateCreator.java index 3f5926080fcbc..d4ded4377b290 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerDelegateCreator.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerDelegateCreator.java @@ -92,8 +92,7 @@ public ConsumerDelegate create(LogContext logContext, valueDeserializer, client, subscriptions, - metadata, - Optional.empty() + metadata ); else return new ClassicKafkaConsumer<>( diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberStateListener.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberStateListener.java index 98b6271fcc0a5..62d82dcf93760 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberStateListener.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberStateListener.java @@ -36,7 +36,7 @@ public interface MemberStateListener { * not part of the group anymore. * @param memberId Current member ID. It won't change until the process is terminated. */ - void onMemberEpochUpdated(Optional memberEpoch, String memberId); + void onMemberEpochUpdated(Optional memberEpoch, Optional memberId); /** * This callback is invoked when a group member's assigned set of partitions changes. Assignments can change via diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java index 4058bd727253f..27dd432b599e3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java @@ -54,6 +54,7 @@ public class RequestManagers implements Closeable { public final Optional shareHeartbeatRequestManager; public final Optional consumerMembershipManager; public final Optional shareMembershipManager; + public final Optional streamsMembershipManager; public final OffsetsRequestManager offsetsRequestManager; public final TopicMetadataRequestManager topicMetadataRequestManager; public final FetchRequestManager fetchRequestManager; @@ -70,7 +71,8 @@ public RequestManagers(LogContext logContext, Optional commitRequestManager, Optional heartbeatRequestManager, Optional membershipManager, - Optional streamsGroupHeartbeatRequestManager) { + Optional streamsGroupHeartbeatRequestManager, + Optional streamsMembershipManager) { this.log = logContext.logger(RequestManagers.class); this.offsetsRequestManager = requireNonNull(offsetsRequestManager, "OffsetsRequestManager cannot be null"); this.coordinatorRequestManager = coordinatorRequestManager; @@ -83,6 +85,7 @@ public RequestManagers(LogContext logContext, this.consumerMembershipManager = membershipManager; this.shareMembershipManager = Optional.empty(); this.streamsGroupHeartbeatRequestManager = streamsGroupHeartbeatRequestManager; + this.streamsMembershipManager = streamsMembershipManager; List> list = new ArrayList<>(); list.add(coordinatorRequestManager); @@ -93,6 +96,7 @@ public RequestManagers(LogContext logContext, list.add(Optional.of(topicMetadataRequestManager)); list.add(Optional.of(fetchRequestManager)); list.add(streamsGroupHeartbeatRequestManager); + list.add(streamsMembershipManager); entries = Collections.unmodifiableList(list); } @@ -110,6 +114,7 @@ public RequestManagers(LogContext logContext, this.consumerMembershipManager = Optional.empty(); this.shareMembershipManager = shareMembershipManager; this.streamsGroupHeartbeatRequestManager = Optional.empty(); + this.streamsMembershipManager = Optional.empty(); this.offsetsRequestManager = null; this.topicMetadataRequestManager = null; this.fetchRequestManager = null; @@ -194,6 +199,7 @@ protected RequestManagers create() { CoordinatorRequestManager coordinator = null; CommitRequestManager commitRequestManager = null; StreamsGroupHeartbeatRequestManager streamsGroupHeartbeatRequestManager = null; + StreamsMembershipManager streamsMembershipManager = null; if (groupRebalanceConfig != null && groupRebalanceConfig.groupId != null) { Optional serverAssignor = Optional.ofNullable(config.getString(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG)); @@ -240,17 +246,18 @@ protected RequestManagers create() { if (streamsInstanceMetadata.isPresent()) { streamsGroupHeartbeatRequestManager = new StreamsGroupHeartbeatRequestManager( - logContext, - time, - config, - coordinator, - membershipManager, - backgroundEventHandler, - metrics, - streamsInstanceMetadata.get(), - metadata + logContext, + time, + config, + coordinator, + streamsMembershipManager, + backgroundEventHandler, + metrics, + streamsInstanceMetadata.get() ); } else { + membershipManager.registerStateListener(commitRequestManager); + membershipManager.registerStateListener(applicationThreadMemberStateListener); heartbeatRequestManager = new ConsumerHeartbeatRequestManager( logContext, time, @@ -284,7 +291,8 @@ protected RequestManagers create() { Optional.ofNullable(commitRequestManager), Optional.ofNullable(heartbeatRequestManager), Optional.ofNullable(membershipManager), - Optional.ofNullable(streamsGroupHeartbeatRequestManager) + Optional.ofNullable(streamsGroupHeartbeatRequestManager), + Optional.ofNullable(streamsMembershipManager) ); } }; diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsAssignmentInterface.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsAssignmentInterface.java index de6cb9e209833..8bafa619627cb 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsAssignmentInterface.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsAssignmentInterface.java @@ -16,20 +16,37 @@ */ package org.apache.kafka.clients.consumer.internals; +import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler; +import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent; +import org.apache.kafka.clients.consumer.internals.events.ErrorEvent; +import org.apache.kafka.clients.consumer.internals.events.EventProcessor; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnAllTasksLostCallbackCompletedEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnAllTasksLostCallbackNeededEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksAssignedCallbackCompletedEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksAssignedCallbackNeededEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksRevokedCallbackCompletedEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksRevokedCallbackNeededEvent; +import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.TopicPartition; import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.UUID; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.function.Supplier; /** * Defines a self-contained object to exchange assignment-related metadata with the Kafka Streams instance. @@ -128,6 +145,8 @@ public String toString() { public static class Assignment { + public static final Assignment EMPTY = new Assignment(); + public final Set activeTasks = new HashSet<>(); public final Set standbyTasks = new HashSet<>(); @@ -281,17 +300,173 @@ public String toString() { } } + private final BlockingQueue onCallbackRequests = new LinkedBlockingQueue<>(); + + private ApplicationEventHandler applicationEventHandler = null; + + private Optional, Optional>> onTasksRevokedCallback = null; + private Optional>> onTasksAssignedCallback = null; + private Optional>> onAllTasksLostCallback = null; + + private final StreamsRebalanceEventProcessor streamsRebalanceEventProcessor; + + private class StreamsRebalanceEventProcessor implements EventProcessor { + + @Override + public void process(final BackgroundEvent event) { + switch (event.type()) { + case ERROR: + process((ErrorEvent) event); + break; + + case STREAMS_ON_TASKS_REVOKED_CALLBACK_NEEDED: + process((StreamsOnTasksRevokedCallbackNeededEvent) event); + break; + + case STREAMS_ON_TASKS_ASSIGNED_CALLBACK_NEEDED: + process((StreamsOnTasksAssignedCallbackNeededEvent) event); + break; + + case STREAMS_ON_ALL_TASKS_LOST_CALLBACK_NEEDED: + process((StreamsOnAllTasksLostCallbackNeededEvent) event); + break; + + default: + throw new IllegalArgumentException("Background event type " + event.type() + " was not expected"); + + } + } + + private void process(final ErrorEvent event) { + throw event.error(); + } + + private void process(final StreamsOnTasksRevokedCallbackNeededEvent event) { + StreamsOnTasksRevokedCallbackCompletedEvent invokedEvent = invokeOnTasksRevokedCallback(event.activeTasksToRevoke(), event.future()); + applicationEventHandler.add(invokedEvent); + if (invokedEvent.error().isPresent()) { + throw invokedEvent.error().get(); + } + } + + private void process(final StreamsOnTasksAssignedCallbackNeededEvent event) { + StreamsOnTasksAssignedCallbackCompletedEvent invokedEvent = invokeOnTasksAssignedCallback(event.assignment(), event.future()); + applicationEventHandler.add(invokedEvent); + if (invokedEvent.error().isPresent()) { + throw invokedEvent.error().get(); + } + } + + private void process(final StreamsOnAllTasksLostCallbackNeededEvent event) { + StreamsOnAllTasksLostCallbackCompletedEvent invokedEvent = invokeOnAllTasksLostCallback(event.future()); + applicationEventHandler.add(invokedEvent); + if (invokedEvent.error().isPresent()) { + throw invokedEvent.error().get(); + } + } + + private StreamsOnTasksRevokedCallbackCompletedEvent invokeOnTasksRevokedCallback(final Set activeTasksToRevoke, + final CompletableFuture future) { + final Optional exceptionFromCallback = onTasksRevokedCallback + .orElseThrow(() -> new IllegalStateException("No tasks assignment callback set!")).apply(activeTasksToRevoke); + + return exceptionFromCallback + .map(exception -> + new StreamsOnTasksRevokedCallbackCompletedEvent( + future, + Optional.of(ConsumerUtils.maybeWrapAsKafkaException(exception, "Task revocation callback throws an error")) + )) + .orElseGet(() -> new StreamsOnTasksRevokedCallbackCompletedEvent(future, Optional.empty())); + } + + private StreamsOnTasksAssignedCallbackCompletedEvent invokeOnTasksAssignedCallback(final StreamsAssignmentInterface.Assignment assignment, + final CompletableFuture future) { + Optional error = Optional.empty(); + // ToDo: Can we avoid the following check? + if (!assignment.equals(reconciledAssignment.get())) { + + final Optional exceptionFromCallback = onTasksAssignedCallback + .orElseThrow(() -> new IllegalStateException("No tasks assignment callback set!")).apply(assignment); + + if (exceptionFromCallback.isPresent()) { + error = Optional.of(ConsumerUtils.maybeWrapAsKafkaException(exceptionFromCallback.get(), "Task assignment callback throws an error")); + } else { + reconciledAssignment.set(assignment); + } + } + return new StreamsOnTasksAssignedCallbackCompletedEvent(future, error); + } + + private StreamsOnAllTasksLostCallbackCompletedEvent invokeOnAllTasksLostCallback(final CompletableFuture future) { + final Optional exceptionFromCallback = onAllTasksLostCallback + .orElseThrow(() -> new IllegalStateException("No tasks assignment callback set!")).get(); + + final Optional error; + + if (exceptionFromCallback.isPresent()) { + error = Optional.of(ConsumerUtils.maybeWrapAsKafkaException(exceptionFromCallback.get(), "Task assignment callback throws an error")); + } else { + error = Optional.empty(); + reconciledAssignment.set(Assignment.EMPTY); + } + + return new StreamsOnAllTasksLostCallbackCompletedEvent(future, error); + } + } + public StreamsAssignmentInterface(UUID processId, Optional endpoint, Map subtopologyMap, - Map clientTags - ) { + Map clientTags) { this.processId = processId; this.endpoint = endpoint; this.subtopologyMap = subtopologyMap; this.taskLags = new HashMap<>(); this.shutdownRequested = new AtomicBoolean(false); this.clientTags = clientTags; + this.streamsRebalanceEventProcessor = new StreamsRebalanceEventProcessor(); + } + + public void setOnTasksRevokedCallback(final Function, Optional> onTasksRevokedCallback) { + this.onTasksRevokedCallback = Optional.ofNullable(onTasksRevokedCallback); + } + + public void setOnTasksAssignedCallback(final Function> onTasksAssignedCallback) { + this.onTasksAssignedCallback = Optional.ofNullable(onTasksAssignedCallback); + } + + public void setOnAllTasksLostCallback(final Supplier> onAllTasksLostCallback) { + this.onAllTasksLostCallback = Optional.ofNullable(onAllTasksLostCallback); + } + + public void setApplicationEventHandler(final ApplicationEventHandler applicationEventHandler) { + this.applicationEventHandler = applicationEventHandler; + } + + public CompletableFuture requestOnTasksAssignedCallbackInvocation(final Assignment assignment) { + final StreamsOnTasksAssignedCallbackNeededEvent onTasksAssignedCallbackNeededEvent = new StreamsOnTasksAssignedCallbackNeededEvent(assignment); + onCallbackRequests.add(onTasksAssignedCallbackNeededEvent); + return onTasksAssignedCallbackNeededEvent.future(); + } + + public CompletableFuture requestOnTasksRevokedCallbackInvocation(final Set activeTasksToRevoke) { + final StreamsOnTasksRevokedCallbackNeededEvent onTasksRevokedCallbackNeededEvent = new StreamsOnTasksRevokedCallbackNeededEvent(activeTasksToRevoke); + onCallbackRequests.add(onTasksRevokedCallbackNeededEvent); + return onTasksRevokedCallbackNeededEvent.future(); + } + + public CompletableFuture requestOnAllTasksLostCallbackInvocation() { + final StreamsOnAllTasksLostCallbackNeededEvent onAllTasksLostCallbackNeededEvent = new StreamsOnAllTasksLostCallbackNeededEvent(); + onCallbackRequests.add(onAllTasksLostCallbackNeededEvent); + return onAllTasksLostCallbackNeededEvent.future(); + } + + public void processStreamsRebalanceEvents() { + LinkedList events = new LinkedList<>(); + onCallbackRequests.drainTo(events); + for (BackgroundEvent event : events) { + streamsRebalanceEventProcessor.process(event); + } } @Override diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsGroupHeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsGroupHeartbeatRequestManager.java index 7506a49793949..f435b1cd21b94 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsGroupHeartbeatRequestManager.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsGroupHeartbeatRequestManager.java @@ -23,11 +23,8 @@ import org.apache.kafka.clients.consumer.internals.events.ErrorEvent; import org.apache.kafka.clients.consumer.internals.metrics.HeartbeatMetricsManager; import org.apache.kafka.common.TopicPartition; -import org.apache.kafka.common.Uuid; import org.apache.kafka.common.errors.GroupAuthorizationException; import org.apache.kafka.common.errors.RetriableException; -import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData; -import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData.TopicPartitions; import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData; import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.CopartitionGroup; import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.TaskIds; @@ -35,7 +32,6 @@ import org.apache.kafka.common.message.StreamsGroupHeartbeatResponseData.Endpoint; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.protocol.Errors; -import org.apache.kafka.common.requests.ConsumerGroupHeartbeatResponse; import org.apache.kafka.common.requests.StreamsGroupHeartbeatRequest; import org.apache.kafka.common.requests.StreamsGroupHeartbeatResponse; import org.apache.kafka.common.utils.LogContext; @@ -51,11 +47,9 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.IntStream; -import java.util.stream.Stream; public class StreamsGroupHeartbeatRequestManager implements RequestManager { @@ -69,7 +63,7 @@ public class StreamsGroupHeartbeatRequestManager implements RequestManager { private final StreamsGroupHeartbeatRequestManager.HeartbeatState heartbeatState; - private final ConsumerMembershipManager membershipManager; + private final StreamsMembershipManager membershipManager; private final BackgroundEventHandler backgroundEventHandler; @@ -79,20 +73,15 @@ public class StreamsGroupHeartbeatRequestManager implements RequestManager { private StreamsAssignmentInterface streamsInterface; - private final Map assignedTopicIdCache; - - private final ConsumerMetadata metadata; - public StreamsGroupHeartbeatRequestManager( final LogContext logContext, final Time time, final ConsumerConfig config, final CoordinatorRequestManager coordinatorRequestManager, - final ConsumerMembershipManager membershipManager, + final StreamsMembershipManager membershipManager, final BackgroundEventHandler backgroundEventHandler, final Metrics metrics, - final StreamsAssignmentInterface streamsAssignmentInterface, - final ConsumerMetadata metadata + final StreamsAssignmentInterface streamsAssignmentInterface ) { this.coordinatorRequestManager = coordinatorRequestManager; this.logger = logContext.logger(getClass()); @@ -108,14 +97,12 @@ public StreamsGroupHeartbeatRequestManager( this.pollTimer = time.timer(maxPollIntervalMs); this.metricsManager = new HeartbeatMetricsManager(metrics); this.streamsInterface = streamsAssignmentInterface; - this.assignedTopicIdCache = new HashMap<>(); - this.metadata = metadata; } @Override public NetworkClientDelegate.PollResult poll(long currentTimeMs) { if (!coordinatorRequestManager.coordinator().isPresent() || membershipManager.shouldSkipHeartbeat()) { - membershipManager.onHeartbeatRequestSkipped(); + membershipManager.transitionToUnsubscribeIfLeaving(); return NetworkClientDelegate.PollResult.EMPTY; } pollTimer.update(currentTimeMs); @@ -126,7 +113,7 @@ public NetworkClientDelegate.PollResult poll(long currentTimeMs) { "messages. You can address this either by increasing max.poll.interval.ms or by " + "reducing the maximum size of batches returned in poll() with max.poll.records."); - membershipManager.transitionToSendingLeaveGroup(true); + membershipManager.onPollTimerExpired(); NetworkClientDelegate.UnsentRequest leaveHeartbeat = makeHeartbeatRequest(currentTimeMs, true); // We can ignore the leave response because we can join before or after receiving the response. @@ -145,7 +132,7 @@ public NetworkClientDelegate.PollResult poll(long currentTimeMs) { return new NetworkClientDelegate.PollResult(heartbeatRequestState.heartbeatIntervalMs, Collections.singletonList(request)); } - public ConsumerMembershipManager membershipManager() { + public StreamsMembershipManager membershipManager() { return membershipManager; } @@ -252,14 +239,6 @@ private void onSuccessResponse(final StreamsGroupHeartbeatResponse response, fin streamsInterface.partitionsByHost.set(convertHostInfoMap(data)); } - ConsumerGroupHeartbeatResponseData cgData = new ConsumerGroupHeartbeatResponseData(); - cgData.setMemberId(data.memberId()); - cgData.setMemberEpoch(data.memberEpoch()); - cgData.setErrorCode(data.errorCode()); - cgData.setErrorMessage(data.errorMessage()); - cgData.setThrottleTimeMs(data.throttleTimeMs()); - cgData.setHeartbeatIntervalMs(data.heartbeatIntervalMs()); - List statuses = data.status(); if (statuses != null && !statuses.isEmpty()) { @@ -269,54 +248,7 @@ private void onSuccessResponse(final StreamsGroupHeartbeatResponse response, fin logger.warn("Membership is in the following statuses: {}.", statusDetails); } - if (data.activeTasks() != null && data.standbyTasks() != null && data.warmupTasks() != null) { - - setTargetAssignment(data); - setTargetAssignmentForConsumerGroup(data, cgData); - - } else { - if (data.activeTasks() != null || data.standbyTasks() != null || data.warmupTasks() != null) { - throw new IllegalStateException("Invalid response data, task collections must be all null or all non-null: " + data); - } - } - - membershipManager.onHeartbeatSuccess(new ConsumerGroupHeartbeatResponse(cgData)); - } - - private void setTargetAssignmentForConsumerGroup(final StreamsGroupHeartbeatResponseData data, - final ConsumerGroupHeartbeatResponseData cgData) { - Map tps = new HashMap<>(); - data.activeTasks().forEach(taskId -> Stream.concat( - streamsInterface.subtopologyMap().get(taskId.subtopologyId()).sourceTopics.stream(), - streamsInterface.subtopologyMap().get(taskId.subtopologyId()).repartitionSourceTopics.keySet().stream() - ) - .forEach(topic -> { - final TopicPartitions toInsert = tps.computeIfAbsent(topic, k -> { - final Optional uuid = findTopicIdInGlobalOrLocalCache(topic); - if (uuid.isPresent()) { - TopicPartitions t = - new TopicPartitions(); - t.setTopicId(uuid.get()); - return t; - } else { - return null; - } - }); - if (toInsert != null) { - toInsert.partitions().addAll(taskId.partitions()); - } - })); - ConsumerGroupHeartbeatResponseData.Assignment cgAssignment = new ConsumerGroupHeartbeatResponseData.Assignment(); - cgAssignment.setTopicPartitions(new ArrayList<>(tps.values())); - cgData.setAssignment(cgAssignment); - } - - private void setTargetAssignment(final StreamsGroupHeartbeatResponseData data) { - Assignment targetAssignment = new Assignment(); - updateTaskIdCollection(data.activeTasks(), targetAssignment.activeTasks); - updateTaskIdCollection(data.standbyTasks(), targetAssignment.standbyTasks); - updateTaskIdCollection(data.warmupTasks(), targetAssignment.warmupTasks); - streamsInterface.targetAssignment.set(targetAssignment); + membershipManager.onHeartbeatSuccess(response); } private static Map> convertHostInfoMap( @@ -410,7 +342,7 @@ private void onErrorResponse(final StreamsGroupHeartbeatResponse response, message = String.format("StreamsGroupHeartbeatRequest failed for member %s because epoch %s is fenced.", membershipManager.memberId(), membershipManager.memberEpoch()); logInfo(message, response, currentTimeMs); - membershipManager.transitionToFenced(); + membershipManager.onFenced(); // Skip backoff so that a next HB to rejoin is sent as soon as the fenced member releases its assignment heartbeatRequestState.reset(); break; @@ -419,7 +351,7 @@ private void onErrorResponse(final StreamsGroupHeartbeatResponse response, message = String.format("StreamsGroupHeartbeatRequest failed because member %s is unknown.", membershipManager.memberId()); logInfo(message, response, currentTimeMs); - membershipManager.transitionToFenced(); + membershipManager.onFenced(); // Skip backoff so that a next HB to rejoin is sent as soon as the fenced member releases its assignment heartbeatRequestState.reset(); break; @@ -507,22 +439,9 @@ private void updateHeartbeatIntervalMs(final long heartbeatIntervalMs) { } } - private Optional findTopicIdInGlobalOrLocalCache(String topicName) { - Uuid idFromMetadataCache = metadata.topicIds().getOrDefault(topicName, null); - if (idFromMetadataCache != null) { - // Add topic name to local cache, so it can be reused if included in a next target - // assignment if metadata cache not available. - assignedTopicIdCache.put(topicName, idFromMetadataCache); - return Optional.of(idFromMetadataCache); - } else { - Uuid idFromLocalCache = assignedTopicIdCache.getOrDefault(topicName, null); - return Optional.ofNullable(idFromLocalCache); - } - } - static class HeartbeatState { - private final ConsumerMembershipManager membershipManager; + private final StreamsMembershipManager membershipManager; private final int rebalanceTimeoutMs; private final StreamsGroupHeartbeatRequestManager.HeartbeatState.SentFields sentFields; @@ -533,7 +452,7 @@ static class HeartbeatState { public HeartbeatState( final StreamsAssignmentInterface streamsInterface, - final ConsumerMembershipManager membershipManager, + final StreamsMembershipManager membershipManager, final int rebalanceTimeoutMs) { this.membershipManager = membershipManager; diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java index fb34f50847b66..532a95e3959a2 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java @@ -16,9 +16,9 @@ */ package org.apache.kafka.clients.consumer.internals; -import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler; -import org.apache.kafka.clients.consumer.internals.events.StreamsOnAssignmentCallbackCompletedEvent; -import org.apache.kafka.clients.consumer.internals.events.StreamsOnAssignmentCallbackNeededEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnAllTasksLostCallbackCompletedEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksAssignedCallbackCompletedEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksRevokedCallbackCompletedEvent; import org.apache.kafka.clients.consumer.internals.metrics.ConsumerRebalanceMetricsManager; import org.apache.kafka.clients.consumer.internals.metrics.RebalanceMetricsManager; import org.apache.kafka.common.KafkaException; @@ -35,6 +35,7 @@ import org.slf4j.Logger; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; @@ -145,8 +146,6 @@ public int hashCode() { private final SubscriptionState subscriptionState; - private final BackgroundEventHandler backgroundEventHandler; - private MemberState state; private final String groupId; @@ -165,6 +164,8 @@ public int hashCode() { private boolean rejoinedWhileReconciliationInProgress; + private final List stateUpdatesListeners = new ArrayList<>(); + private final Optional clientTelemetryReporter; private LocalAssignment targetAssignment = LocalAssignment.NONE; @@ -184,7 +185,6 @@ public StreamsMembershipManager(final String groupId, final SubscriptionState subscriptionState, final LogContext logContext, final Optional clientTelemetryReporter, - final BackgroundEventHandler backgroundEventHandler, final Time time, final Metrics metrics) { log = logContext.logger(StreamsMembershipManager.class); @@ -193,7 +193,6 @@ public StreamsMembershipManager(final String groupId, this.streamsAssignmentInterface = streamsAssignmentInterface; this.subscriptionState = subscriptionState; this.clientTelemetryReporter = clientTelemetryReporter; - this.backgroundEventHandler = backgroundEventHandler; metricsManager = new ConsumerRebalanceMetricsManager(metrics); this.time = time; } @@ -231,6 +230,14 @@ private boolean isNotInGroup() { state == MemberState.STALE; } + public void registerStateListener(MemberStateListener listener) { + stateUpdatesListeners.add(Objects.requireNonNull(listener, "State updates listener cannot be null")); + } + + void notifyEpochChange(Optional epoch, Optional memberId) { + stateUpdatesListeners.forEach(stateListener -> stateListener.onMemberEpochUpdated(epoch, memberId)); + } + private void transitionToJoining() { if (state == MemberState.FATAL) { log.warn("No action taken to join the group with the updated subscription because " + @@ -276,8 +283,9 @@ private void finalizeLeaving() { private void transitionToStale() { transitionTo(MemberState.STALE); - CompletableFuture onAllTasksRevokedDone = invokeRevokingAllTasksCallback(); - staleMemberAssignmentRelease = onAllTasksRevokedDone.whenComplete((result, error) -> { + final CompletableFuture onAllTasksLostCallbackExecution = + streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation(); + staleMemberAssignmentRelease = onAllTasksLostCallbackExecution.whenComplete((result, error) -> { if (error != null) { log.error("Task revocation callback invocation failed " + "after member left group due to expired poll timer.", error); @@ -293,6 +301,7 @@ public void transitionToFatal() { MemberState previousState = state; transitionTo(MemberState.FATAL); log.error("Member {} with epoch {} transitioned to fatal state", memberIdInfoForLog(), memberEpoch); + notifyEpochChange(Optional.empty(), Optional.empty()); if (previousState == MemberState.UNSUBSCRIBED) { log.debug("Member {} with epoch {} got fatal error from the broker but it already " + @@ -308,8 +317,8 @@ public void transitionToFatal() { return; } - CompletableFuture callbackResult = invokeRevokingAllTasksCallback(); - callbackResult.whenComplete((result, error) -> { + CompletableFuture onAllTasksLostCallbackExecuted = streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation(); + onAllTasksLostCallbackExecuted.whenComplete((result, error) -> { if (error != null) { log.error("onTaskAssignment callback invocation failed while releasing assignment" + "after member failed with fatal error.", error); @@ -359,7 +368,15 @@ private void resetEpoch() { } private void updateMemberEpoch(int newEpoch) { - memberEpoch = newEpoch; + boolean newEpochReceived = this.memberEpoch != newEpoch; + this.memberEpoch = newEpoch; + if (newEpochReceived) { + if (memberEpoch > 0) { + notifyEpochChange(Optional.of(memberEpoch), Optional.ofNullable(memberId)); + } else { + notifyEpochChange(Optional.empty(), Optional.empty()); + } + } } private void clearPendingTaskAssignment() { @@ -519,17 +536,17 @@ public void onFenced() { log.debug("Member {} with epoch {} transitioned to {} state. It will release its " + "assignment and rejoin the group.", memberIdInfoForLog(), memberEpoch, MemberState.FENCED); - CompletableFuture callbackResult = invokeRevokingAllTasksCallback(); + CompletableFuture callbackResult = streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation(); callbackResult.whenComplete((result, error) -> { if (error != null) { - log.error("onTaskAssignment callback invocation failed while releasing assignment" + + log.error("onAllTasksLost callback invocation failed while releasing assignment" + " after member got fenced. Member will rejoin the group anyways.", error); } clearTaskAndPartitionAssignment(); if (state == MemberState.FENCED) { transitionToJoining(); } else { - log.debug("Fenced member onTaskAssignment callback completed but the state has " + + log.debug("Fenced member onAllTasksLost callback completed but the state has " + "already changed to {}, so the member won't rejoin the group", state); } }); @@ -615,15 +632,19 @@ public CompletableFuture leaveGroup() { CompletableFuture onGroupLeft = new CompletableFuture<>(); leaveGroupInProgress = Optional.of(onGroupLeft); - CompletableFuture onAllTasksRevokedDone = prepareLeaving(); - onAllTasksRevokedDone.whenComplete((__, callbackError) -> leaving(callbackError)); + CompletableFuture onAllTasksRevokedCallbackExecuted = prepareLeaving(); + onAllTasksRevokedCallbackExecuted.whenComplete((__, callbackError) -> leaving(callbackError)); return onGroupLeft; } private CompletableFuture prepareLeaving() { transitionTo(MemberState.PREPARE_LEAVING); - return invokeRevokingAllTasksCallback(); + if (memberEpoch > 0) { + return revokeActiveTasks(toTaskIdSet(currentAssignment.activeTasks)); + } else { + return releaseLostActiveTasks(); + } } private void leaving(Throwable callbackError) { @@ -636,8 +657,8 @@ private void leaving(Throwable callbackError) { "to clear its assignment and send a leave group heartbeat", memberIdInfoForLog()); } - subscriptionState.unsubscribe(); clearTaskAndPartitionAssignment(); + subscriptionState.unsubscribe(); transitionToSendingLeaveGroup(false); } @@ -709,41 +730,54 @@ private void maybeReconcile() { // ToDo: add standby and warmup tasks SortedSet assignedActiveTasks = toTaskIdSet(targetAssignment.activeTasks); + SortedSet ownedActiveTasks = toTaskIdSet(currentAssignment.activeTasks); + SortedSet activeTasksToRevoke = new TreeSet<>(ownedActiveTasks); + activeTasksToRevoke.removeAll(assignedActiveTasks); log.info("Assigned tasks with local epoch {}\n" + "\tMember: {}\n" + - "\tActive tasks: {}\n", + "\tAssigned active tasks: {}\n" + + "\tOwned active tasks: {}\n" + + "\tActive tasks to revoke: {}\n", targetAssignment.localEpoch, memberIdInfoForLog(), - assignedActiveTasks + assignedActiveTasks, + ownedActiveTasks, + activeTasksToRevoke ); - SortedSet ownedTopicPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR); - ownedTopicPartitions.addAll(subscriptionState.assignedPartitions()); + SortedSet ownedTopicPartitionsFromSubscriptionState = new TreeSet<>(TOPIC_PARTITION_COMPARATOR); + ownedTopicPartitionsFromSubscriptionState.addAll(subscriptionState.assignedPartitions()); + SortedSet ownedTopicPartitionsFromAssignedTasks = + topicPartitionsForActiveTasks(currentAssignment.activeTasks); + if (!ownedTopicPartitionsFromAssignedTasks.equals(ownedTopicPartitionsFromSubscriptionState)) { + throw new IllegalStateException("Owned partitions from subscription state and owned partitions from " + + "assigned active tasks are not equal. " + + "Owned partitions from subscription state: " + ownedTopicPartitionsFromSubscriptionState + ", " + + "Owned partitions from assigned active tasks: " + ownedTopicPartitionsFromAssignedTasks); + } SortedSet assignedTopicPartitions = topicPartitionsForActiveTasks(targetAssignment.activeTasks); SortedSet assignedTopicPartitionsNotPreviouslyOwned = - assignedTopicPartitionsNotPreviouslyOwned(assignedTopicPartitions, ownedTopicPartitions); + partitionsToAssignNotPreviouslyOwned(assignedTopicPartitions, ownedTopicPartitionsFromSubscriptionState); + SortedSet partitionsToRevoke = new TreeSet<>(ownedTopicPartitionsFromSubscriptionState); + partitionsToRevoke.removeAll(assignedTopicPartitions); - subscriptionState.assignFromSubscribedAwaitingCallback( - assignedTopicPartitions, - assignedTopicPartitionsNotPreviouslyOwned - ); + final CompletableFuture onTasksRevokedCallbackExecuted = revokeActiveTasks(activeTasksToRevoke); + + final CompletableFuture onTasksRevokedAndAssignedCallbacksExecuted = onTasksRevokedCallbackExecuted.thenCompose(__ -> { + if (!maybeAbortReconciliation()) { + return assignActiveTasks(assignedActiveTasks, ownedActiveTasks); + } + return CompletableFuture.completedFuture(null); + }); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - new StreamsOnAssignmentCallbackNeededEvent(new StreamsAssignmentInterface.Assignment( - assignedActiveTasks, - Collections.emptySet(), - Collections.emptySet() - )); - CompletableFuture onTasksAssignmentDone = onAssignmentCallbackNeededEvent.future(); - backgroundEventHandler.add(onAssignmentCallbackNeededEvent); // The current target assignment is captured to ensure that acknowledging the current assignment is done with // the same target assignment that was used when this reconciliation was initiated. LocalAssignment currentTargetAssignment = targetAssignment; - onTasksAssignmentDone.whenComplete((__, callbackError) -> { + onTasksRevokedAndAssignedCallbacksExecuted.whenComplete((__, callbackError) -> { if (callbackError != null) { - log.error("Reconciliation failed: onTasksAssignment callback invocation failed for tasks {}", - targetAssignment, callbackError); + log.error("Reconciliation failed: callback invocation failed for tasks {}", + currentTargetAssignment, callbackError); markReconciliationCompleted(); } else { if (reconciliationInProgress && !maybeAbortReconciliation()) { @@ -756,8 +790,61 @@ private void maybeReconcile() { }); } - private SortedSet assignedTopicPartitionsNotPreviouslyOwned(final SortedSet assignedTopicPartitions, - final SortedSet ownedTopicPartitions) { + private CompletableFuture revokeActiveTasks(final SortedSet activeTasksToRevoke) { + if (activeTasksToRevoke.isEmpty()) { + return CompletableFuture.completedFuture(null); + } + + log.info("Revoking previously assigned active tasks {}", activeTasksToRevoke.stream() + .map(StreamsAssignmentInterface.TaskId::toString) + .collect(Collectors.joining(", "))); + + final SortedSet partitionsToRevoke = topicPartitionsForActiveTasks(activeTasksToRevoke); + log.debug("Marking partitions pending for revocation: {}", partitionsToRevoke); + subscriptionState.markPendingRevocation(partitionsToRevoke); + + return streamsAssignmentInterface.requestOnTasksRevokedCallbackInvocation(activeTasksToRevoke); + } + + private CompletableFuture assignActiveTasks(final SortedSet activeTasksToAssign, + final SortedSet ownedActiveTasks) { + log.info("Assigning active tasks {}", activeTasksToAssign.stream() + .map(StreamsAssignmentInterface.TaskId::toString) + .collect(Collectors.joining(", "))); + + final SortedSet partitionsToAssign = topicPartitionsForActiveTasks(activeTasksToAssign); + final SortedSet partitionsToAssigneNotPreviouslyOwned = + partitionsToAssignNotPreviouslyOwned(partitionsToAssign, topicPartitionsForActiveTasks(ownedActiveTasks)); + + subscriptionState.assignFromSubscribedAwaitingCallback( + partitionsToAssign, + partitionsToAssigneNotPreviouslyOwned + ); + + return streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation( + new StreamsAssignmentInterface.Assignment( + activeTasksToAssign, + Collections.emptySet(), + Collections.emptySet() + ) + ); + } + + private CompletableFuture releaseLostActiveTasks() { + final SortedSet activeTasksToRelease = toTaskIdSet(currentAssignment.activeTasks); + log.info("Revoking previously assigned and now lost active tasks {}", activeTasksToRelease.stream() + .map(StreamsAssignmentInterface.TaskId::toString) + .collect(Collectors.joining(", "))); + + final SortedSet partitionsToRelease = topicPartitionsForActiveTasks(activeTasksToRelease); + log.debug("Marking lost partitions pending for revocation: {}", partitionsToRelease); + subscriptionState.markPendingRevocation(partitionsToRelease); + + return streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation(); + } + + private SortedSet partitionsToAssignNotPreviouslyOwned(final SortedSet assignedTopicPartitions, + final SortedSet ownedTopicPartitions) { SortedSet assignedPartitionsNotPreviouslyOwned = new TreeSet<>(TOPIC_PARTITION_COMPARATOR); assignedPartitionsNotPreviouslyOwned.addAll(assignedTopicPartitions); assignedPartitionsNotPreviouslyOwned.removeAll(ownedTopicPartitions); @@ -779,6 +866,19 @@ private SortedSet topicPartitionsForActiveTasks(final Map topicPartitionsForActiveTasks(final SortedSet activeTasks) { + final SortedSet topicPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR); + activeTasks.forEach(task -> + Stream.concat( + streamsAssignmentInterface.subtopologyMap().get(task.subtopologyId()).sourceTopics.stream(), + streamsAssignmentInterface.subtopologyMap().get(task.subtopologyId()).repartitionSourceTopics.keySet().stream() + ).forEach(topic -> { + topicPartitions.add(new TopicPartition(topic, task.partitionId())); + }) + ); + return topicPartitions; + } + private void markReconciliationCompleted() { reconciliationInProgress = false; rejoinedWhileReconciliationInProgress = false; @@ -801,28 +901,47 @@ private void markReconciliationInProgress() { rejoinedWhileReconciliationInProgress = false; } - private CompletableFuture invokeRevokingAllTasksCallback() { - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - new StreamsOnAssignmentCallbackNeededEvent(new StreamsAssignmentInterface.Assignment( - Collections.emptySet(), - Collections.emptySet(), - Collections.emptySet() - )); - backgroundEventHandler.add(onAssignmentCallbackNeededEvent); - return onAssignmentCallbackNeededEvent.future(); + public void onTasksRevokedCallbackCompleted(final StreamsOnTasksRevokedCallbackCompletedEvent event) { + Optional error = event.error(); + CompletableFuture future = event.future(); + + if (error.isPresent()) { + Exception e = error.get(); + log.warn("The onTasksRevoked callback completed with an error ({}); " + + "signaling to continue to the next phase of rebalance", e.getMessage()); + future.completeExceptionally(e); + } else { + log.debug("The onTasksRevoked callback completed successfully; signaling to continue to the next phase of rebalance"); + future.complete(null); + } + } + + public void onTasksAssignedCallbackCompleted(final StreamsOnTasksAssignedCallbackCompletedEvent event) { + Optional error = event.error(); + CompletableFuture future = event.future(); + + if (error.isPresent()) { + Exception e = error.get(); + log.warn("The onTasksAssigned callback completed with an error ({}); " + + "signaling to continue to the next phase of rebalance", e.getMessage()); + future.completeExceptionally(e); + } else { + log.debug("The onTasksAssigned callback completed successfully; signaling to continue to the next phase of rebalance"); + future.complete(null); + } } - public void onTaskAssignmentCallbackCompleted(StreamsOnAssignmentCallbackCompletedEvent event) { + public void onAllTasksLostCallbackCompleted(final StreamsOnAllTasksLostCallbackCompletedEvent event) { Optional error = event.error(); CompletableFuture future = event.future(); if (error.isPresent()) { Exception e = error.get(); - log.warn("The onTaskAssignment callback completed with an error ({}); " + + log.warn("The onAllTasksLost callback completed with an error ({}); " + "signaling to continue to the next phase of rebalance", e.getMessage()); future.completeExceptionally(e); } else { - log.debug("The onTaskAssignment callback completed successfully; signaling to continue to the next phase of rebalance"); + log.debug("The onAllTasksLost callback completed successfully; signaling to continue to the next phase of rebalance"); future.complete(null); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java index 799501987f0e8..56bccaf2b3f2c 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java @@ -39,7 +39,9 @@ public enum Type { SHARE_ACKNOWLEDGE_ON_CLOSE, SHARE_ACKNOWLEDGEMENT_COMMIT_CALLBACK_REGISTRATION, SEEK_UNVALIDATED, - STREAMS_ON_ASSIGNMENT_CALLBACK_COMPLETED, + STREAMS_ON_TASKS_ASSIGNED_CALLBACK_COMPLETED, + STREAMS_ON_TASKS_REVOKED_CALLBACK_COMPLETED, + STREAMS_ON_ALL_TASKS_LOST_CALLBACK_COMPLETED, } private final Type type; diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java index a4eb0771d53c0..ced9b9544b175 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java @@ -176,6 +176,18 @@ public void process(ApplicationEvent event) { process((SeekUnvalidatedEvent) event); return; + case STREAMS_ON_TASKS_REVOKED_CALLBACK_COMPLETED: + process((StreamsOnTasksRevokedCallbackCompletedEvent) event); + return; + + case STREAMS_ON_TASKS_ASSIGNED_CALLBACK_COMPLETED: + process((StreamsOnTasksAssignedCallbackCompletedEvent) event); + return; + + case STREAMS_ON_ALL_TASKS_LOST_CALLBACK_COMPLETED: + process((StreamsOnAllTasksLostCallbackCompletedEvent) event); + return; + default: log.warn("Application event type {} was not expected", event.type()); } @@ -284,54 +296,19 @@ private void process(final ListOffsetsEvent event) { * consumer join the group if it is not part of it yet, or send the updated subscription if * it is already a member on the next poll. */ - private void process(final TopicSubscriptionChangeEvent event) { - if (requestManagers.consumerHeartbeatRequestManager.isEmpty()) { + private void process(final TopicSubscriptionChangeEvent ignored) { + if (requestManagers.consumerMembershipManager.isPresent()) { + requestManagers.consumerMembershipManager.get().onSubscriptionUpdated(); + } else if (requestManagers.streamsMembershipManager.isPresent()) { + requestManagers.streamsMembershipManager.get().onSubscriptionUpdated(); + } else { log.warn("Group membership manager not present when processing a subscribe event"); - event.future().complete(null); - return; - } - - try { - if (subscriptions.subscribe(event.topics(), event.listener())) - this.metadataVersionSnapshot = metadata.requestUpdateForNewTopics(); - - // Join the group if not already part of it, or just send the new subscription to the broker on the next poll. - requestManagers.consumerHeartbeatRequestManager.get().membershipManager().onSubscriptionUpdated(); - event.future().complete(null); - } catch (Exception e) { - event.future().completeExceptionally(e); - } - } - - /** - * Process event that indicates that the subscription topic pattern changed. This will make the - * consumer join the group if it is not part of it yet, or send the updated subscription if - * it is already a member on the next poll. - */ - private void process(final TopicPatternSubscriptionChangeEvent event) { - try { - subscriptions.subscribe(event.pattern(), event.listener()); - metadata.requestUpdateForNewTopics(); - updatePatternSubscription(metadata.fetch()); - event.future().complete(null); - } catch (Exception e) { - event.future().completeExceptionally(e); - } - } - - /** - * Process event that re-evaluates the subscribed regular expression using the latest topics from metadata, only if metadata changed. - * This will make the consumer send the updated subscription on the next poll. - */ - private void process(final UpdatePatternSubscriptionEvent event) { - if (!subscriptions.hasPatternSubscription()) { - return; } if (this.metadataVersionSnapshot < metadata.updateVersion()) { this.metadataVersionSnapshot = metadata.updateVersion(); updatePatternSubscription(metadata.fetch()); } - event.future().complete(null); + ignored.future().complete(null); } /** @@ -346,6 +323,9 @@ private void process(final UnsubscribeEvent event) { if (requestManagers.consumerMembershipManager.isPresent()) { CompletableFuture future = requestManagers.consumerMembershipManager.get().leaveGroup(); future.whenComplete(complete(event.future())); + } else if (requestManagers.streamsMembershipManager.isPresent()) { + CompletableFuture future = requestManagers.streamsMembershipManager.get().leaveGroup(); + future.whenComplete(complete(event.future())); } else { // If the consumer is not using the group management capabilities, we still need to clear all assignments it may have. subscriptions.unsubscribe(); @@ -520,6 +500,33 @@ private void process(final ShareAcknowledgementCommitCallbackRegistrationEvent e manager.setAcknowledgementCommitCallbackRegistered(event.isCallbackRegistered()); } + private void process(final StreamsOnTasksRevokedCallbackCompletedEvent event) { + if (!requestManagers.streamsMembershipManager.isPresent()) { + log.warn("An internal error occurred; the Streams membership manager was not present, so the notification " + + "of the onTasksRevoked callback execution could not be sent"); + return; + } + requestManagers.streamsMembershipManager.get().onTasksRevokedCallbackCompleted(event); + } + + private void process(final StreamsOnTasksAssignedCallbackCompletedEvent event) { + if (!requestManagers.streamsMembershipManager.isPresent()) { + log.warn("An internal error occurred; the Streams membership manager was not present, so the notification " + + "of the onTasksAssigned callback execution could not be sent"); + return; + } + requestManagers.streamsMembershipManager.get().onTasksAssignedCallbackCompleted(event); + } + + private void process(final StreamsOnAllTasksLostCallbackCompletedEvent event) { + if (!requestManagers.streamsMembershipManager.isPresent()) { + log.warn("An internal error occurred; the Streams membership manager was not present, so the notification " + + "of the onAllTasksLost callback execution could not be sent"); + return; + } + requestManagers.streamsMembershipManager.get().onAllTasksLostCallbackCompleted(event); + } + private BiConsumer complete(final CompletableFuture b) { return (value, exception) -> { if (exception != null) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java index cffcca42f48a5..b18611895f9eb 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java @@ -28,7 +28,8 @@ public abstract class BackgroundEvent { public enum Type { ERROR, CONSUMER_REBALANCE_LISTENER_CALLBACK_NEEDED, SHARE_ACKNOWLEDGEMENT_COMMIT_CALLBACK, - STREAMS_ON_ASSIGNMENT_CALLBACK_NEEDED + STREAMS_ON_TASKS_ASSIGNED_CALLBACK_NEEDED, STREAMS_ON_TASKS_REVOKED_CALLBACK_NEEDED, + STREAMS_ON_ALL_TASKS_LOST_CALLBACK_NEEDED } private final Type type; diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAllTasksLostCallbackCompletedEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAllTasksLostCallbackCompletedEvent.java new file mode 100644 index 0000000000000..b84e9d0c1386d --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAllTasksLostCallbackCompletedEvent.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.consumer.internals.events; + +import org.apache.kafka.common.KafkaException; + +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; + +public class StreamsOnAllTasksLostCallbackCompletedEvent extends ApplicationEvent { + + private final CompletableFuture future; + private final Optional error; + + public StreamsOnAllTasksLostCallbackCompletedEvent(final CompletableFuture future, + final Optional error) { + super(Type.STREAMS_ON_ALL_TASKS_LOST_CALLBACK_COMPLETED); + this.future = Objects.requireNonNull(future); + this.error = Objects.requireNonNull(error); + } + + public CompletableFuture future() { + return future; + } + + public Optional error() { + return error; + } + + @Override + protected String toStringBase() { + return super.toStringBase() + + ", future=" + future + + ", error=" + error; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAllTasksLostCallbackNeededEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAllTasksLostCallbackNeededEvent.java new file mode 100644 index 0000000000000..294cf3e2fde03 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAllTasksLostCallbackNeededEvent.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.consumer.internals.events; + +public class StreamsOnAllTasksLostCallbackNeededEvent extends CompletableBackgroundEvent { + + public StreamsOnAllTasksLostCallbackNeededEvent() { + super(Type.STREAMS_ON_ALL_TASKS_LOST_CALLBACK_NEEDED, Long.MAX_VALUE); + } + + @Override + protected String toStringBase() { + return super.toStringBase(); + } +} + diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksAssignedCallbackCompletedEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksAssignedCallbackCompletedEvent.java new file mode 100644 index 0000000000000..96c2519bb2d33 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksAssignedCallbackCompletedEvent.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.consumer.internals.events; + +import org.apache.kafka.common.KafkaException; + +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; + +public class StreamsOnTasksAssignedCallbackCompletedEvent extends ApplicationEvent { + + private final CompletableFuture future; + private final Optional error; + + public StreamsOnTasksAssignedCallbackCompletedEvent(final CompletableFuture future, + final Optional error) { + super(Type.STREAMS_ON_TASKS_ASSIGNED_CALLBACK_COMPLETED); + this.future = Objects.requireNonNull(future); + this.error = Objects.requireNonNull(error); + } + + public CompletableFuture future() { + return future; + } + + public Optional error() { + return error; + } + + @Override + protected String toStringBase() { + return super.toStringBase() + + ", future=" + future + + ", error=" + error; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAssignmentCallbackNeededEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksAssignedCallbackNeededEvent.java similarity index 82% rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAssignmentCallbackNeededEvent.java rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksAssignedCallbackNeededEvent.java index 5ef6cd9c462a9..276eba3ca4ab0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAssignmentCallbackNeededEvent.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksAssignedCallbackNeededEvent.java @@ -20,12 +20,12 @@ import java.util.Objects; -public class StreamsOnAssignmentCallbackNeededEvent extends CompletableBackgroundEvent { +public class StreamsOnTasksAssignedCallbackNeededEvent extends CompletableBackgroundEvent { private final StreamsAssignmentInterface.Assignment assignment; - public StreamsOnAssignmentCallbackNeededEvent(StreamsAssignmentInterface.Assignment assignment) { - super(Type.STREAMS_ON_ASSIGNMENT_CALLBACK_NEEDED, Long.MAX_VALUE); + public StreamsOnTasksAssignedCallbackNeededEvent(StreamsAssignmentInterface.Assignment assignment) { + super(Type.STREAMS_ON_TASKS_ASSIGNED_CALLBACK_NEEDED, Long.MAX_VALUE); this.assignment = Objects.requireNonNull(assignment); } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAssignmentCallbackCompletedEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksRevokedCallbackCompletedEvent.java similarity index 82% rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAssignmentCallbackCompletedEvent.java rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksRevokedCallbackCompletedEvent.java index 5b702efb62c99..5717012ac4576 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnAssignmentCallbackCompletedEvent.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksRevokedCallbackCompletedEvent.java @@ -22,14 +22,14 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; -public class StreamsOnAssignmentCallbackCompletedEvent extends ApplicationEvent { +public class StreamsOnTasksRevokedCallbackCompletedEvent extends ApplicationEvent { private final CompletableFuture future; private final Optional error; - public StreamsOnAssignmentCallbackCompletedEvent(final CompletableFuture future, - final Optional error) { - super(Type.STREAMS_ON_ASSIGNMENT_CALLBACK_COMPLETED); + public StreamsOnTasksRevokedCallbackCompletedEvent(final CompletableFuture future, + final Optional error) { + super(Type.STREAMS_ON_TASKS_REVOKED_CALLBACK_COMPLETED); this.future = Objects.requireNonNull(future); this.error = Objects.requireNonNull(error); } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksRevokedCallbackNeededEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksRevokedCallbackNeededEvent.java new file mode 100644 index 0000000000000..be7ae13b446bc --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/StreamsOnTasksRevokedCallbackNeededEvent.java @@ -0,0 +1,43 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.consumer.internals.events; + +import org.apache.kafka.clients.consumer.internals.StreamsAssignmentInterface; + +import java.util.Objects; +import java.util.Set; + +public class StreamsOnTasksRevokedCallbackNeededEvent extends CompletableBackgroundEvent { + + private final Set activeTasksToRevoke; + + public StreamsOnTasksRevokedCallbackNeededEvent(final Set activeTasksToRevoke) { + super(Type.STREAMS_ON_TASKS_REVOKED_CALLBACK_NEEDED, Long.MAX_VALUE); + this.activeTasksToRevoke = Objects.requireNonNull(activeTasksToRevoke); + } + + public Set activeTasksToRevoke() { + return activeTasksToRevoke; + } + + @Override + protected String toStringBase() { + return super.toStringBase() + + ", active tasks to revoke=" + activeTasksToRevoke; + } +} + diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsGroupHeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsGroupHeartbeatRequestManagerTest.java index 9e88e2c756954..a9e34cf2067ff 100644 --- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsGroupHeartbeatRequestManagerTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsGroupHeartbeatRequestManagerTest.java @@ -19,22 +19,17 @@ import org.apache.kafka.clients.ClientResponse; import org.apache.kafka.clients.consumer.ConsumerConfig; import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.UnsentRequest; -import org.apache.kafka.clients.consumer.internals.StreamsAssignmentInterface.Assignment; import org.apache.kafka.clients.consumer.internals.StreamsAssignmentInterface.Subtopology; -import org.apache.kafka.clients.consumer.internals.StreamsAssignmentInterface.TaskId; import org.apache.kafka.clients.consumer.internals.StreamsAssignmentInterface.TopicInfo; import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler; import org.apache.kafka.common.Node; import org.apache.kafka.common.Uuid; -import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData; -import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData.TopicPartitions; import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData; import org.apache.kafka.common.message.StreamsGroupHeartbeatResponseData; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.metrics.Sensor; import org.apache.kafka.common.protocol.ApiKeys; import org.apache.kafka.common.protocol.Errors; -import org.apache.kafka.common.requests.ConsumerGroupHeartbeatResponse; import org.apache.kafka.common.requests.RequestHeader; import org.apache.kafka.common.requests.StreamsGroupHeartbeatRequest; import org.apache.kafka.common.requests.StreamsGroupHeartbeatResponse; @@ -59,7 +54,6 @@ import java.util.Properties; import java.util.Set; import java.util.UUID; -import java.util.stream.Collectors; import static org.apache.kafka.common.utils.Utils.mkEntry; import static org.apache.kafka.common.utils.Utils.mkMap; @@ -93,7 +87,7 @@ class StreamsGroupHeartbeatRequestManagerTest { private CoordinatorRequestManager coordinatorRequestManager; @Mock - private ConsumerMembershipManager membershipManager; + private StreamsMembershipManager membershipManager; @Mock private BackgroundEventHandler backgroundEventHandler; @@ -141,8 +135,7 @@ void setUp() { membershipManager, backgroundEventHandler, metrics, - streamsAssignmentInterface, - metadata + streamsAssignmentInterface ); when(membershipManager.groupId()).thenReturn(TEST_GROUP_ID); @@ -160,7 +153,7 @@ void testNoHeartbeatIfCoordinatorUnknown() { NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds()); assertEquals(0, result.unsentRequests.size()); - verify(membershipManager).onHeartbeatRequestSkipped(); + verify(membershipManager).transitionToUnsubscribeIfLeaving(); } @Test @@ -171,7 +164,7 @@ void testNoHeartbeatIfHeartbeatSkipped() { NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds()); assertEquals(0, result.unsentRequests.size()); - verify(membershipManager).onHeartbeatRequestSkipped(); + verify(membershipManager).transitionToUnsubscribeIfLeaving(); } @Test @@ -347,35 +340,35 @@ void testSuccessfulResponse() { mockResponse(data); - ArgumentCaptor captor = ArgumentCaptor.forClass(ConsumerGroupHeartbeatResponse.class); + ArgumentCaptor captor = ArgumentCaptor.forClass(StreamsGroupHeartbeatResponse.class); verify(membershipManager, times(1)).onHeartbeatSuccess(captor.capture()); - ConsumerGroupHeartbeatResponseData response = captor.getValue().data(); + StreamsGroupHeartbeatResponseData response = captor.getValue().data(); assertEquals(Errors.NONE.code(), response.errorCode()); assertEquals(TEST_MEMBER_ID, response.memberId()); assertEquals(TEST_MEMBER_EPOCH, response.memberEpoch()); assertEquals(TEST_THROTTLE_TIME_MS, response.throttleTimeMs()); assertEquals(1000, response.heartbeatIntervalMs()); - final List tps = response.assignment().topicPartitions(); - assertEquals(2, tps.size()); - assertEquals(Set.of(uuid0, uuid1), tps.stream().map(TopicPartitions::topicId).collect(Collectors.toSet())); - assertEquals(Collections.singletonList(0), tps.get(0).partitions()); - assertEquals(Collections.singletonList(0), tps.get(1).partitions()); - - final Assignment targetAssignment = streamsAssignmentInterface.targetAssignment.get(); - assertEquals(1, targetAssignment.activeTasks.size()); - final TaskId activeTaskId = targetAssignment.activeTasks.stream().findFirst().get(); - assertEquals(activeTaskId.subtopologyId(), "0"); - assertEquals(activeTaskId.partitionId(), 0); - - assertEquals(1, targetAssignment.standbyTasks.size()); - final TaskId standbyTaskId = targetAssignment.standbyTasks.stream().findFirst().get(); - assertEquals(standbyTaskId.subtopologyId(), "1"); - assertEquals(standbyTaskId.partitionId(), 1); - - assertEquals(1, targetAssignment.warmupTasks.size()); - final TaskId warmupTaskId = targetAssignment.warmupTasks.stream().findFirst().get(); - assertEquals(warmupTaskId.subtopologyId(), "2"); - assertEquals(warmupTaskId.partitionId(), 2); +// final List tps = response.assign.topicPartitions(); +// assertEquals(2, tps.size()); +// assertEquals(Set.of(uuid0, uuid1), tps.stream().map(TopicPartitions::topicId).collect(Collectors.toSet())); +// assertEquals(Collections.singletonList(0), tps.get(0).partitions()); +// assertEquals(Collections.singletonList(0), tps.get(1).partitions()); + +// final Assignment targetAssignment = streamsAssignmentInterface.targetAssignment.get(); +// assertEquals(1, targetAssignment.activeTasks.size()); +// final TaskId activeTaskId = targetAssignment.activeTasks.stream().findFirst().get(); +// assertEquals(activeTaskId.subtopologyId(), "0"); +// assertEquals(activeTaskId.partitionId(), 0); +// +// assertEquals(1, targetAssignment.standbyTasks.size()); +// final TaskId standbyTaskId = targetAssignment.standbyTasks.stream().findFirst().get(); +// assertEquals(standbyTaskId.subtopologyId(), "1"); +// assertEquals(standbyTaskId.partitionId(), 1); +// +// assertEquals(1, targetAssignment.warmupTasks.size()); +// final TaskId warmupTaskId = targetAssignment.warmupTasks.stream().findFirst().get(); +// assertEquals(warmupTaskId.subtopologyId(), "2"); +// assertEquals(warmupTaskId.partitionId(), 2); } diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java index 7d954477da96a..6f7c66af6a00b 100644 --- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java @@ -18,8 +18,9 @@ import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent; import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler; -import org.apache.kafka.clients.consumer.internals.events.StreamsOnAssignmentCallbackCompletedEvent; -import org.apache.kafka.clients.consumer.internals.events.StreamsOnAssignmentCallbackNeededEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnAllTasksLostCallbackCompletedEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksAssignedCallbackCompletedEvent; +import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksAssignedCallbackNeededEvent; import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.message.StreamsGroupHeartbeatResponseData; @@ -34,6 +35,8 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -54,6 +57,7 @@ import static org.junit.jupiter.api.Assertions.assertInstanceOf; 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.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -85,6 +89,9 @@ public class StreamsMembershipManagerTest { @Mock private StreamsAssignmentInterface streamsAssignmentInterface; + @Captor + private ArgumentCaptor onAssignmentCallbackNeededEventCaptor; + private Queue backgroundEventQueue = new LinkedList<>(); private BackgroundEventHandler backgroundEventHandler = new BackgroundEventHandler(backgroundEventQueue); @@ -96,7 +103,6 @@ public void setup() { subscriptionState, new LogContext("test"), Optional.empty(), - backgroundEventHandler, time, metrics ); @@ -151,62 +157,87 @@ public void testJoining() { @Test public void testReconcilingEmptyToSingleActiveTask() { setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final Set activeTasks = + Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); + final CompletableFuture onTasksAssignedCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecuted); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - final Collection expectedPartitionAssignment = Set.of(new TopicPartition(TOPIC_0, PARTITION_0)); - verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionAssignment, expectedPartitionAssignment); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - final Set activeTasks = - Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); - final StreamsAssignmentInterface.Assignment expectedTaskAssignment = makeTaskAssignment(activeTasks); - assertEquals(expectedTaskAssignment, onAssignmentCallbackNeededEvent.assignment()); - verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionAssignment); + final Collection expectedPartitionsToAssign = Set.of(new TopicPartition(TOPIC_0, PARTITION_0)); + verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionsToAssign, expectedPartitionsToAssign); + verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionsToAssign); verifyInStateReconciling(membershipManager); - onAssignmentCallbackNeededEvent.future().complete(null); - verify(subscriptionState).enablePartitionsAwaitingCallback(expectedPartitionAssignment); + onTasksAssignedCallbackExecuted.complete(null); + verify(subscriptionState).enablePartitionsAwaitingCallback(expectedPartitionsToAssign); verifyInStateAcknowledging(membershipManager); + verify(streamsAssignmentInterface, never()).requestOnTasksRevokedCallbackInvocation(any()); } @Test public void testReconcilingActiveTaskToDifferentActiveTask() { setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + final CompletableFuture onTasksAssignedCallbackExecuted = new CompletableFuture<>(); + final CompletableFuture onTasksRevokedCallbackExecuted = new CompletableFuture<>(); + final Set activeTasksSetup = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0) + ); + final Set activeTasks = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_1) + ); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasksSetup))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); + when(streamsAssignmentInterface.requestOnTasksRevokedCallbackInvocation(activeTasksSetup)) + .thenReturn(onTasksRevokedCallbackExecuted); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecuted); when(subscriptionState.assignedPartitions()) .thenReturn(Collections.emptySet()) .thenReturn(Set.of(new TopicPartition(TOPIC_0, PARTITION_0))); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_1))); - final Collection expectedPartitionAssignment = Set.of(new TopicPartition(TOPIC_0, PARTITION_1)); - verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionAssignment, expectedPartitionAssignment); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - final Set activeTasks = Set.of( - new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_1) - ); - final StreamsAssignmentInterface.Assignment expectedStreamsAssignment = makeTaskAssignment(activeTasks); - assertEquals(expectedStreamsAssignment, onAssignmentCallbackNeededEvent.assignment()); - verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionAssignment); + final Set expectedPartitionsToRevoke = Set.of(new TopicPartition(TOPIC_0, PARTITION_0)); + final Collection expectedPartitionsToAssign = Set.of(new TopicPartition(TOPIC_0, PARTITION_1)); + verify(subscriptionState).markPendingRevocation(expectedPartitionsToRevoke); + verify(subscriptionState, never()).assignFromSubscribedAwaitingCallback(expectedPartitionsToAssign, expectedPartitionsToAssign); + onTasksRevokedCallbackExecuted.complete(null); + verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionsToAssign, expectedPartitionsToAssign); + verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionsToAssign); verifyInStateReconciling(membershipManager); - onAssignmentCallbackNeededEvent.future().complete(null); - verify(subscriptionState).enablePartitionsAwaitingCallback(expectedPartitionAssignment); + onTasksAssignedCallbackExecuted.complete(null); + verify(subscriptionState).enablePartitionsAwaitingCallback(expectedPartitionsToAssign); verifyInStateAcknowledging(membershipManager); } @Test public void testReconcilingSingleActiveTaskToAdditionalActiveTask() { setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + final CompletableFuture onTasksAssignedCallbackExecuted = new CompletableFuture<>(); + final Set activeTasksSetup = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0) + ); + final Set activeTasks = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0), + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_1) + ); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasksSetup))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecuted); when(subscriptionState.assignedPartitions()) .thenReturn(Collections.emptySet()) .thenReturn(Set.of(new TopicPartition(TOPIC_0, PARTITION_0))); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0, PARTITION_1))); @@ -214,57 +245,58 @@ public void testReconcilingSingleActiveTaskToAdditionalActiveTask() { new TopicPartition(TOPIC_0, PARTITION_0), new TopicPartition(TOPIC_0, PARTITION_1) ); - final Collection expectedAdditionalPartitionAssignment = Set.of(new TopicPartition(TOPIC_0, PARTITION_1)); - verify(subscriptionState).assignFromSubscribedAwaitingCallback( - expectedPartitionAssignment, - expectedAdditionalPartitionAssignment - ); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - final Set activeTasks = Set.of( - new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0), - new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_1) - ); - final StreamsAssignmentInterface.Assignment expectedStreamsAssignment = makeTaskAssignment(activeTasks); - assertEquals(expectedStreamsAssignment, onAssignmentCallbackNeededEvent.assignment()); - verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedAdditionalPartitionAssignment); + final Collection expectedAdditionalPartitionsToAssign = Set.of(new TopicPartition(TOPIC_0, PARTITION_1)); + verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionAssignment, expectedAdditionalPartitionsToAssign); + verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionAssignment); verifyInStateReconciling(membershipManager); - onAssignmentCallbackNeededEvent.future().complete(null); - verify(subscriptionState).enablePartitionsAwaitingCallback(expectedAdditionalPartitionAssignment); + onTasksAssignedCallbackExecuted.complete(null); + verify(subscriptionState).enablePartitionsAwaitingCallback(expectedAdditionalPartitionsToAssign); verifyInStateAcknowledging(membershipManager); + verify(streamsAssignmentInterface, never()).requestOnTasksRevokedCallbackInvocation(any()); } @Test public void testReconcilingMultipleActiveTaskToSingleActiveTask() { setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + final CompletableFuture onTasksRevokedCallbackExecuted = new CompletableFuture<>(); + final CompletableFuture onTasksAssignedCallbackExecuted = new CompletableFuture<>(); + final Set activeTasksSetup = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0), + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_1) + ); + final Set activeTasksToRevoke = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0) + ); + final Set activeTasks = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_1) + ); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasksSetup))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); + when(streamsAssignmentInterface.requestOnTasksRevokedCallbackInvocation(activeTasksToRevoke)) + .thenReturn(onTasksRevokedCallbackExecuted); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecuted); when(subscriptionState.assignedPartitions()) .thenReturn(Collections.emptySet()) .thenReturn(Set.of(new TopicPartition(TOPIC_0, PARTITION_0), new TopicPartition(TOPIC_0, PARTITION_1))); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0, PARTITION_1))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_1))); - final Collection expectedPartitionAssignment = Set.of( - new TopicPartition(TOPIC_0, PARTITION_1) - ); - final Collection expectedAdditionalPartitionAssignment = Collections.emptySet(); - verify(subscriptionState).assignFromSubscribedAwaitingCallback( - expectedPartitionAssignment, - expectedAdditionalPartitionAssignment - ); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - final Set activeTasks = Set.of( - new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_1) - ); - final StreamsAssignmentInterface.Assignment expectedStreamsAssignment = makeTaskAssignment(activeTasks); - assertEquals(expectedStreamsAssignment, onAssignmentCallbackNeededEvent.assignment()); - verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedAdditionalPartitionAssignment); + final Set expectedPartitionsToRevoke = Set.of(new TopicPartition(TOPIC_0, PARTITION_0)); + final Collection expectedPartitionsToAssign = Set.of(new TopicPartition(TOPIC_0, PARTITION_1)); + final Collection expectedAdditionalPartitionsToAssign = Collections.emptySet(); + verify(subscriptionState).markPendingRevocation(expectedPartitionsToRevoke); + verify(subscriptionState, never()).assignFromSubscribedAwaitingCallback(expectedPartitionsToAssign, expectedAdditionalPartitionsToAssign); + onTasksRevokedCallbackExecuted.complete(null); + verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionsToAssign, expectedAdditionalPartitionsToAssign); + verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedAdditionalPartitionsToAssign); verifyInStateReconciling(membershipManager); - onAssignmentCallbackNeededEvent.future().complete(null); - verify(subscriptionState).enablePartitionsAwaitingCallback(expectedAdditionalPartitionAssignment); + onTasksAssignedCallbackExecuted.complete(null); + verify(subscriptionState).enablePartitionsAwaitingCallback(expectedAdditionalPartitionsToAssign); verifyInStateAcknowledging(membershipManager); } @@ -274,79 +306,50 @@ public void testReconcilingEmptyToMultipleActiveTaskOfDifferentSubtopologies() { SUB_TOPOLOGY_ID_0, TOPIC_0, SUB_TOPOLOGY_ID_1, TOPIC_1 ); - joining(); - - reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0), SUB_TOPOLOGY_ID_1, List.of(PARTITION_0))); - - final Collection expectedPartitionAssignment = Set.of( - new TopicPartition(TOPIC_0, PARTITION_0), - new TopicPartition(TOPIC_1, PARTITION_0) - ); - verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionAssignment, expectedPartitionAssignment); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); final Set activeTasks = Set.of( new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0), new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_1, PARTITION_0) ); - final StreamsAssignmentInterface.Assignment expectedTaskAssignment = makeTaskAssignment(activeTasks); - assertEquals(expectedTaskAssignment, onAssignmentCallbackNeededEvent.assignment()); - verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionAssignment); - verifyInStateReconciling(membershipManager); - onAssignmentCallbackNeededEvent.future().complete(null); - verify(subscriptionState).enablePartitionsAwaitingCallback(expectedPartitionAssignment); - verifyInStateAcknowledging(membershipManager); - } - - @Test - public void testReconcilingEmptyToMultipleActiveTaskOfConcatenatedSubtopologies() { - setupStreamsAssignmentInterfaceWithTwoConcatenedSubtopologies( - SUB_TOPOLOGY_ID_0, TOPIC_0, - SUB_TOPOLOGY_ID_1, TOPIC_1 - ); + final CompletableFuture onTasksAssignedCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecuted); joining(); - reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0), SUB_TOPOLOGY_ID_1, List.of(PARTITION_0))); + reconcile(makeHeartbeatResponse( + SUB_TOPOLOGY_ID_0, List.of(PARTITION_0), + SUB_TOPOLOGY_ID_1, List.of(PARTITION_0)) + ); - final Collection expectedPartitionAssignment = Set.of( + final Collection expectedPartitionsToAssign = Set.of( new TopicPartition(TOPIC_0, PARTITION_0), new TopicPartition(TOPIC_1, PARTITION_0) ); - verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionAssignment, expectedPartitionAssignment); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - final Set activeTasks = Set.of( - new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0), - new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_1, PARTITION_0) - ); - final StreamsAssignmentInterface.Assignment expectedTaskAssignment = makeTaskAssignment(activeTasks); - assertEquals(expectedTaskAssignment, onAssignmentCallbackNeededEvent.assignment()); - verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionAssignment); + verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionsToAssign, expectedPartitionsToAssign); + verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionsToAssign); verifyInStateReconciling(membershipManager); - onAssignmentCallbackNeededEvent.future().complete(null); - verify(subscriptionState).enablePartitionsAwaitingCallback(expectedPartitionAssignment); + onTasksAssignedCallbackExecuted.complete(null); + verify(subscriptionState).enablePartitionsAwaitingCallback(expectedPartitionsToAssign); verifyInStateAcknowledging(membershipManager); + verify(streamsAssignmentInterface, never()).requestOnTasksRevokedCallbackInvocation(any()); } @Test public void testReconcilingAndAssignmentCallbackFails() { - final String topicName = "test_topic"; - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, topicName); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); final Set activeTasks = Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); - final StreamsAssignmentInterface.Assignment expectedStreamsAssignment = makeTaskAssignment(activeTasks); + final CompletableFuture onTasksAssignedCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecuted); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - final Collection expectedPartitionAssignment = Set.of(new TopicPartition(topicName, PARTITION_0)); + final Collection expectedPartitionAssignment = Set.of(new TopicPartition(TOPIC_0, PARTITION_0)); verify(subscriptionState).assignFromSubscribedAwaitingCallback(expectedPartitionAssignment, expectedPartitionAssignment); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - assertEquals(expectedStreamsAssignment, onAssignmentCallbackNeededEvent.assignment()); verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionAssignment); verifyInStateReconciling(membershipManager); - onAssignmentCallbackNeededEvent.future().completeExceptionally(new RuntimeException("KABOOM!")); + onTasksAssignedCallbackExecuted.completeExceptionally(new RuntimeException("KABOOM!")); verifyInStateReconciling(membershipManager); verify(subscriptionState, never()).enablePartitionsAwaitingCallback(expectedPartitionAssignment); } @@ -365,6 +368,9 @@ public void testLeaveGroupWhenNotInGroup() { @Test public void testLeaveGroupWhenNotInGroupAndFenced() { + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); fenced(); final CompletableFuture future = membershipManager.leaveGroup(); @@ -380,39 +386,52 @@ public void testLeaveGroupWhenNotInGroupAndFenced() { @Test public void testLeaveGroupWhenInGroupWithAssignment() { - final StreamsAssignmentInterface.Assignment emptyStreamsAssignment = makeTaskAssignment(Collections.emptySet()); - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final Set activeTasks = + Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); + final CompletableFuture onTasksRevokedCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksRevokedCallbackInvocation(activeTasks)) + .thenReturn(onTasksRevokedCallbackExecuted); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); - final CompletableFuture future = membershipManager.leaveGroup(); + final CompletableFuture onGroupLeft = membershipManager.leaveGroup(); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - assertEquals(emptyStreamsAssignment, onAssignmentCallbackNeededEvent.assignment()); + assertFalse(onGroupLeft.isDone()); verify(subscriptionState, never()).unsubscribe(); verifyInStatePrepareLeaving(membershipManager); - final CompletableFuture futureBeforeRevocationCallback = membershipManager.leaveGroup(); - assertEquals(future, futureBeforeRevocationCallback); - onAssignmentCallbackNeededEvent.future().complete(null); + final CompletableFuture onGroupLeftBeforeRevocationCallback = membershipManager.leaveGroup(); + assertEquals(onGroupLeft, onGroupLeftBeforeRevocationCallback); + onTasksRevokedCallbackExecuted.complete(null); verify(subscriptionState).unsubscribe(); - assertFalse(future.isDone()); + assertFalse(onGroupLeft.isDone()); verifyInStateLeaving(membershipManager); - final CompletableFuture futureAfterRevocationCallback = membershipManager.leaveGroup(); - assertEquals(future, futureAfterRevocationCallback); + final CompletableFuture onGroupLeftAfterRevocationCallback = membershipManager.leaveGroup(); + assertEquals(onGroupLeft, onGroupLeftAfterRevocationCallback); membershipManager.transitionToUnsubscribeIfLeaving(); verifyInStateUnsubscribed(membershipManager); } @Test public void testTransitionToUnsubscribeWhenInLeaving() { - final StreamsGroupHeartbeatResponse response = makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0)); setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic"); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + final Set activeTasksSetup = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0) + ); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasksSetup))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); + final CompletableFuture onAllTasksRevokedCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksRevokedCallbackInvocation(activeTasksSetup)) + .thenReturn(onAllTasksRevokedCallbackExecuted); joining(); - reconcile(response); - acknowledging(); - CompletableFuture future = leaving(); + reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); + acknowledging(onTasksAssignedCallbackExecutedSetup); + CompletableFuture future = leaving(onAllTasksRevokedCallbackExecuted); membershipManager.transitionToUnsubscribeIfLeaving(); @@ -424,10 +443,15 @@ public void testTransitionToUnsubscribeWhenInLeaving() { @Test public void testOnPollTimerExpired() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final Set activeTasks = + Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); membershipManager.onPollTimerExpired(); @@ -437,10 +461,15 @@ public void testOnPollTimerExpired() { @Test public void testOnHeartbeatRequestGeneratedWhenInAcknowleding() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final Set activeTasks = + Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); membershipManager.onHeartbeatRequestGenerated(); @@ -449,11 +478,16 @@ public void testOnHeartbeatRequestGeneratedWhenInAcknowleding() { @Test public void testOnHeartbeatRequestGeneratedWhenInAcknowledgingAndNewTargetAssignment() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic0"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final Set activeTasks = + Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_1))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); membershipManager.onHeartbeatRequestGenerated(); @@ -462,8 +496,11 @@ public void testOnHeartbeatRequestGeneratedWhenInAcknowledgingAndNewTargetAssign @Test public void testOnHeartbeatRequestGeneratedWhenInLeaving() { + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); - leaving(); + leaving(onAllTasksLostCallbackExecuted); membershipManager.onHeartbeatRequestGenerated(); @@ -472,6 +509,9 @@ public void testOnHeartbeatRequestGeneratedWhenInLeaving() { @Test public void testOnHeartbeatRequestGeneratedWhenInLeavingAndPollTimerExpired() { + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); membershipManager.onPollTimerExpired(); @@ -489,51 +529,69 @@ public void testOnFencedWhenInJoining() { @Test public void testOnFencedWhenInReconciling() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic0"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final Set activeTasks = + Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - backgroundEventQueue.poll(); testOnFencedWhenInJoiningOrReconcilingOrAcknowledgingOrStable(); } @Test public void testOnFencedWhenInAcknowledging() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic0"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final Set activeTasks = + Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); testOnFencedWhenInJoiningOrReconcilingOrAcknowledgingOrStable(); } @Test public void testOnFencedWhenInStable() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic0"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final Set activeTasks = + Set.of(new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0)); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasks))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); stable(); testOnFencedWhenInJoiningOrReconcilingOrAcknowledgingOrStable(); } private void testOnFencedWhenInJoiningOrReconcilingOrAcknowledgingOrStable() { + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); + membershipManager.onFenced(); verifyInStateFenced(membershipManager); assertEquals(StreamsGroupHeartbeatRequest.JOIN_GROUP_MEMBER_EPOCH, membershipManager.memberEpoch()); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - final StreamsAssignmentInterface.Assignment expectedStreamsAssignment = makeTaskAssignment(Collections.emptySet()); - assertEquals(expectedStreamsAssignment, onAssignmentCallbackNeededEvent.assignment()); - onAssignmentCallbackNeededEvent.future().complete(null); + verify(subscriptionState, never()).assignFromSubscribed(Collections.emptySet()); + onAllTasksLostCallbackExecuted.complete(null); verify(subscriptionState).assignFromSubscribed(Collections.emptySet()); verifyInStateJoining(membershipManager); } @Test public void testOnFencedWhenInPrepareLeaving() { + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); testOnFencedWhenInPrepareLeavingOrLeaving(prepareLeaving()); @@ -541,9 +599,12 @@ public void testOnFencedWhenInPrepareLeaving() { @Test public void testOnFencedWhenInLeaving() { + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); - testOnFencedWhenInPrepareLeavingOrLeaving(leaving()); + testOnFencedWhenInPrepareLeavingOrLeaving(leaving(onAllTasksLostCallbackExecuted)); } private void testOnFencedWhenInPrepareLeavingOrLeaving(final CompletableFuture onGroupLeft) { @@ -558,6 +619,9 @@ private void testOnFencedWhenInPrepareLeavingOrLeaving(final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); testTransitionToFatalWhenInPrepareLeavingOrLeaving(prepareLeaving()); @@ -565,9 +629,12 @@ public void testTransitionToFatalWhenInPrepareLeaving() { @Test public void testTransitionToFatalWhenInLeaving() { + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); - testTransitionToFatalWhenInPrepareLeavingOrLeaving(leaving()); + testTransitionToFatalWhenInPrepareLeavingOrLeaving(leaving(onAllTasksLostCallbackExecuted)); } private void testTransitionToFatalWhenInPrepareLeavingOrLeaving(final CompletableFuture onGroupLeft) { @@ -581,63 +648,89 @@ private void testTransitionToFatalWhenInPrepareLeavingOrLeaving(final Completabl @Test public void testTransitionToFatalWhenInJoining() { + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); - testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(); + testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(onAllTasksLostCallbackExecuted); } @Test public void testTransitionToFatalWhenInReconciling() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic0"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + final Set activeTasksSetup = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0) + ); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasksSetup))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - backgroundEventQueue.poll(); - testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(); + testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(onAllTasksLostCallbackExecuted); } @Test public void testTransitionToFatalWhenInAcknowledging() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic0"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + final Set activeTasksSetup = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0) + ); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasksSetup))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); - testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(); + testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(onAllTasksLostCallbackExecuted); } @Test public void testTransitionToFatalWhenInStable() { - setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, "topic0"); + setupStreamsAssignmentInterfaceWithOneSubtopologyOneSourceTopic(SUB_TOPOLOGY_ID_0, TOPIC_0); + final CompletableFuture onTasksAssignedCallbackExecutedSetup = new CompletableFuture<>(); + final Set activeTasksSetup = Set.of( + new StreamsAssignmentInterface.TaskId(SUB_TOPOLOGY_ID_0, PARTITION_0) + ); + when(streamsAssignmentInterface.requestOnTasksAssignedCallbackInvocation(makeTaskAssignment(activeTasksSetup))) + .thenReturn(onTasksAssignedCallbackExecutedSetup); + final CompletableFuture onAllTasksLostCallbackExecuted = new CompletableFuture<>(); + when(streamsAssignmentInterface.requestOnAllTasksLostCallbackInvocation()) + .thenReturn(onAllTasksLostCallbackExecuted); joining(); reconcile(makeHeartbeatResponse(SUB_TOPOLOGY_ID_0, List.of(PARTITION_0))); - acknowledging(); + acknowledging(onTasksAssignedCallbackExecutedSetup); stable(); - testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(); + testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(onAllTasksLostCallbackExecuted); } - private void testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable() { + private void testTransitionToFatalWhenInJoiningOrReconcilingOrAcknowledgingOrStable(final CompletableFuture future) { membershipManager.transitionToFatal(); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - final StreamsAssignmentInterface.Assignment expectedStreamsAssignment = makeTaskAssignment(Collections.emptySet()); - assertEquals(expectedStreamsAssignment, onAssignmentCallbackNeededEvent.assignment()); - onAssignmentCallbackNeededEvent.future().complete(null); + verify(subscriptionState, never()).assignFromSubscribed(Collections.emptySet()); + future.complete(null); verify(subscriptionState).assignFromSubscribed(Collections.emptySet()); verifyInStateFatal(membershipManager); } @Test - public void testOnTaskAssignmentCallbackCompleted() { + public void testOnTasksAssignedCallbackCompleted() { final CompletableFuture future = new CompletableFuture<>(); - final StreamsOnAssignmentCallbackCompletedEvent event = new StreamsOnAssignmentCallbackCompletedEvent( + final StreamsOnTasksAssignedCallbackCompletedEvent event = new StreamsOnTasksAssignedCallbackCompletedEvent( future, Optional.empty() ); - membershipManager.onTaskAssignmentCallbackCompleted(event); + membershipManager.onTasksAssignedCallbackCompleted(event); assertTrue(future.isDone()); assertFalse(future.isCancelled()); @@ -645,15 +738,30 @@ public void testOnTaskAssignmentCallbackCompleted() { } @Test - public void testOnTaskAssignmentCallbackCompletedWhenCallbackFails() { + public void testOnAllTasksLostCallbackCompleted() { + final CompletableFuture future = new CompletableFuture<>(); + final StreamsOnAllTasksLostCallbackCompletedEvent event = new StreamsOnAllTasksLostCallbackCompletedEvent( + future, + Optional.empty() + ); + + membershipManager.onAllTasksLostCallbackCompleted(event); + + assertTrue(future.isDone()); + assertFalse(future.isCancelled()); + assertFalse(future.isCompletedExceptionally()); + } + + @Test + public void testOnTasksAssignedCallbackCompletedWhenCallbackFails() { final String errorMessage = "KABOOM!"; final CompletableFuture future = new CompletableFuture<>(); - final StreamsOnAssignmentCallbackCompletedEvent event = new StreamsOnAssignmentCallbackCompletedEvent( + final StreamsOnAllTasksLostCallbackCompletedEvent event = new StreamsOnAllTasksLostCallbackCompletedEvent( future, Optional.of(new KafkaException(errorMessage)) ); - membershipManager.onTaskAssignmentCallbackCompleted(event); + membershipManager.onAllTasksLostCallbackCompleted(event); assertTrue(future.isDone()); assertFalse(future.isCancelled()); @@ -781,43 +889,6 @@ private void setupStreamsAssignmentInterfaceWithTwoSubtopologies(final String su ); } - private void setupStreamsAssignmentInterfaceWithTwoConcatenedSubtopologies(final String subtopologyId1, - final String topicName1, - final String subtopologyId2, - final String topicName2) { - when(streamsAssignmentInterface.subtopologyMap()).thenReturn( - mkMap( - mkEntry( - subtopologyId1, - new StreamsAssignmentInterface.Subtopology( - Set.of(topicName1), - Collections.emptySet(), - Collections.emptyMap(), - Collections.emptyMap(), - Collections.emptyList() - ) - ), - mkEntry( - subtopologyId2, - new StreamsAssignmentInterface.Subtopology( - Set.of(topicName2), - Collections.emptySet(), - mkMap(mkEntry( - topicName2, - new StreamsAssignmentInterface.TopicInfo( - Optional.empty(), - Optional.empty(), - Collections.emptyMap() - ) - )), - Collections.emptyMap(), - Collections.emptyList() - ) - ) - ) - ); - } - private StreamsGroupHeartbeatResponse makeHeartbeatResponse(final String subtopologyId, final List partitions) { return makeHeartbeatResponse(Collections.singletonList( @@ -874,23 +945,26 @@ private void reconcile(final StreamsGroupHeartbeatResponse response) { } private void acknowledging() { - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); + final StreamsOnTasksAssignedCallbackNeededEvent onAssignmentCallbackNeededEvent = + (StreamsOnTasksAssignedCallbackNeededEvent) backgroundEventQueue.poll(); onAssignmentCallbackNeededEvent.future().complete(null); verifyInStateAcknowledging(membershipManager); } + private void acknowledging(final CompletableFuture future) { + future.complete(null); + verifyInStateAcknowledging(membershipManager); + } + private CompletableFuture prepareLeaving() { final CompletableFuture onGroupLeft = membershipManager.leaveGroup(); verifyInStatePrepareLeaving(membershipManager); return onGroupLeft; } - private CompletableFuture leaving() { + private CompletableFuture leaving(final CompletableFuture onAllTasksRevokedCallbackExecuted) { final CompletableFuture future = prepareLeaving(); - final StreamsOnAssignmentCallbackNeededEvent onAssignmentCallbackNeededEvent = - (StreamsOnAssignmentCallbackNeededEvent) backgroundEventQueue.poll(); - onAssignmentCallbackNeededEvent.future().complete(null); + onAllTasksRevokedCallbackExecuted.complete(null); verifyInStateLeaving(membershipManager); return future; } diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java index 7e65a8661b0a2..454724db046df 100644 --- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java @@ -30,6 +30,8 @@ import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate; import org.apache.kafka.clients.consumer.internals.OffsetsRequestManager; import org.apache.kafka.clients.consumer.internals.RequestManagers; +import org.apache.kafka.clients.consumer.internals.StreamsGroupHeartbeatRequestManager; +import org.apache.kafka.clients.consumer.internals.StreamsMembershipManager; import org.apache.kafka.clients.consumer.internals.SubscriptionState; import org.apache.kafka.clients.consumer.internals.TopicMetadataRequestManager; import org.apache.kafka.common.Cluster; @@ -80,6 +82,8 @@ public class ApplicationEventProcessorTest { private final CommitRequestManager commitRequestManager = mock(CommitRequestManager.class); private final ConsumerHeartbeatRequestManager heartbeatRequestManager = mock(ConsumerHeartbeatRequestManager.class); private final ConsumerMembershipManager membershipManager = mock(ConsumerMembershipManager.class); + private final StreamsGroupHeartbeatRequestManager streamsGroupHeartbeatRequestManager = mock(StreamsGroupHeartbeatRequestManager.class); + private final StreamsMembershipManager streamsMembershipManager = mock(StreamsMembershipManager.class); private final OffsetsRequestManager offsetsRequestManager = mock(OffsetsRequestManager.class); private SubscriptionState subscriptionState = mock(SubscriptionState.class); private final ConsumerMetadata metadata = mock(ConsumerMetadata.class); @@ -95,6 +99,7 @@ private void setupProcessor(boolean withGroupId) { withGroupId ? Optional.of(commitRequestManager) : Optional.empty(), withGroupId ? Optional.of(heartbeatRequestManager) : Optional.empty(), withGroupId ? Optional.of(membershipManager) : Optional.empty(), + Optional.empty(), Optional.empty() ); @@ -106,6 +111,28 @@ private void setupProcessor(boolean withGroupId) { ); } + private void setupProcessorWithStreamsMembershipManager() { + RequestManagers requestManagers = new RequestManagers( + new LogContext(), + offsetsRequestManager, + mock(TopicMetadataRequestManager.class), + mock(FetchRequestManager.class), + Optional.of(mock(CoordinatorRequestManager.class)), + Optional.of(commitRequestManager), + Optional.empty(), + Optional.empty(), + Optional.of(streamsGroupHeartbeatRequestManager), + Optional.of(streamsMembershipManager) + ); + + processor = new ApplicationEventProcessor( + new LogContext(), + requestManagers, + metadata, + subscriptionState + ); + } + @Test public void testPrepClosingCommitEvents() { setupProcessor(true); @@ -124,6 +151,15 @@ public void testProcessUnsubscribeEventWithGroupId() { verify(membershipManager).leaveGroup(); } + @Test + public void testProcessUnsubscribeEventWithStreamsMembershipManager() { + setupProcessorWithStreamsMembershipManager(); + when(heartbeatRequestManager.membershipManager()).thenReturn(membershipManager); + when(streamsMembershipManager.leaveGroup()).thenReturn(CompletableFuture.completedFuture(null)); + processor.process(new UnsubscribeEvent(0)); + verify(streamsMembershipManager).leaveGroup(); + } + @Test public void testProcessUnsubscribeEventWithoutGroupId() { setupProcessor(false); @@ -466,6 +502,53 @@ private static Stream offsetsGenerator() { ); } + @Test + public void testSubscriptionChangeEventWithStreamsMembershipManager() { + SubscriptionChangeEvent event = new SubscriptionChangeEvent(); + + setupProcessorWithStreamsMembershipManager(); + processor.process(event); + verify(streamsMembershipManager).onSubscriptionUpdated(); + // verify member state doesn't transition to JOINING. + verify(streamsMembershipManager, never()).onConsumerPoll(); + } + + @Test + public void testOnTasksRevokedCallbackExecuted() { + StreamsOnTasksRevokedCallbackCompletedEvent event = new StreamsOnTasksRevokedCallbackCompletedEvent( + new CompletableFuture<>(), + Optional.empty() + ); + + setupProcessorWithStreamsMembershipManager(); + processor.process(event); + verify(streamsMembershipManager).onTasksRevokedCallbackCompleted(event); + } + + @Test + public void testOnTasksAssignedCallbackExecuted() { + StreamsOnTasksAssignedCallbackCompletedEvent event = new StreamsOnTasksAssignedCallbackCompletedEvent( + new CompletableFuture<>(), + Optional.empty() + ); + + setupProcessorWithStreamsMembershipManager(); + processor.process(event); + verify(streamsMembershipManager).onTasksAssignedCallbackCompleted(event); + } + + @Test + public void testOnAllTasksLostCallbackExecuted() { + StreamsOnAllTasksLostCallbackCompletedEvent event = new StreamsOnAllTasksLostCallbackCompletedEvent( + new CompletableFuture<>(), + Optional.empty() + ); + + setupProcessorWithStreamsMembershipManager(); + processor.process(event); + verify(streamsMembershipManager).onAllTasksLostCallbackCompleted(event); + } + private List mockCommitResults() { return Collections.singletonList(mock(NetworkClientDelegate.UnsentRequest.class)); } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java index 37428bd234b4d..372254f352af0 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java @@ -25,7 +25,6 @@ import org.apache.kafka.clients.consumer.InvalidOffsetException; import org.apache.kafka.clients.consumer.OffsetResetStrategy; import org.apache.kafka.clients.consumer.internals.StreamsAssignmentInterface; -import org.apache.kafka.clients.consumer.internals.StreamsAssignmentInterface.Assignment; import org.apache.kafka.clients.consumer.internals.StreamsAssignmentInterface.Subtopology; import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.KafkaFuture; @@ -717,6 +716,11 @@ public StreamThread(final Time time, this.streamsUncaughtExceptionHandler = streamsUncaughtExceptionHandler; this.cacheResizer = cacheResizer; this.streamsAssignmentInterface = streamsAssignmentInterface; + if (streamsAssignmentInterface != null) { + streamsAssignmentInterface.setOnTasksRevokedCallback(this::onTasksRevoked); + streamsAssignmentInterface.setOnTasksAssignedCallback(this::onTasksAssigned); + streamsAssignmentInterface.setOnAllTasksLostCallback(this::onAllTasksLost); + } this.streamsMetadataState = streamsMetadataState; // The following sensors are created here but their references are not stored in this object, since within @@ -1073,6 +1077,8 @@ void runOnceWithoutProcessingThreads() { final long startMs = time.milliseconds(); now = startMs; + maybeHandleAssignmentFromStreamsRebalanceProtocol(); + final long pollLatency; taskManager.resumePollingForPartitionsWithAvailableSpace(); pollLatency = pollPhase(); @@ -1403,23 +1409,59 @@ public void maybeHandleAssignmentFromStreamsRebalanceProtocol() { ); // Process assignment from Streams Rebalance Protocol - final Assignment newAssignment = streamsAssignmentInterface.targetAssignment.get(); - if (!newAssignment.equals(streamsAssignmentInterface.reconciledAssignment.get())) { - - final Map> activeTasksWithPartitions = - pairWithTopicPartitions(newAssignment.activeTasks.stream()); - final Map> standbyTasksWithPartitions = - pairWithTopicPartitions(Stream.concat(newAssignment.standbyTasks.stream(), newAssignment.warmupTasks.stream())); + streamsAssignmentInterface.processStreamsRebalanceEvents(); + } + } - log.info("Processing new assignment {} from Streams Rebalance Protocol", newAssignment); + private Optional onTasksRevoked(final Set activeTasksToRevoke) { + try { + final Map> activeTasksToRevokeWithPartitions = + pairWithTopicPartitions(activeTasksToRevoke.stream()); + final Set partitionsToRevoke = activeTasksToRevokeWithPartitions.values().stream() + .flatMap(Collection::stream) + .collect(Collectors.toSet()); - taskManager.handleAssignment( - activeTasksWithPartitions, - standbyTasksWithPartitions - ); - streamsAssignmentInterface.reconciledAssignment.set(newAssignment); + final long start = time.milliseconds(); + try { + log.info("Revoking active tasks {}.", activeTasksToRevoke); + taskManager.handleRevocation(partitionsToRevoke); + } finally { + log.info("partition revocation took {} ms.", time.milliseconds() - start); } + if (state() != State.PENDING_SHUTDOWN) { + setState(State.PARTITIONS_REVOKED); + } + } catch (final Exception exception) { + return Optional.of(exception); + } + return Optional.empty(); + } + + private Optional onTasksAssigned(final StreamsAssignmentInterface.Assignment assignment) { + try { + final Map> activeTasksWithPartitions = + pairWithTopicPartitions(assignment.activeTasks.stream()); + final Map> standbyTasksWithPartitions = + pairWithTopicPartitions(Stream.concat(assignment.standbyTasks.stream(), assignment.warmupTasks.stream())); + + log.info("Processing new assignment {} from Streams Rebalance Protocol", assignment); + + taskManager.handleAssignment(activeTasksWithPartitions, standbyTasksWithPartitions); + setState(State.PARTITIONS_ASSIGNED); + taskManager.handleRebalanceComplete(); + } catch (final Exception exception) { + return Optional.of(exception); + } + return Optional.empty(); + } + + private Optional onAllTasksLost() { + try { + taskManager.handleLostAll(); + } catch (final Exception exception) { + return Optional.of(exception); } + return Optional.empty(); } static Map getTopicPartitionInfo(final Map> partitionsByHost) { diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java index 49428724b1a1d..a2f504cab0f19 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java @@ -52,8 +52,6 @@ public class StreamsRebalanceListener implements ConsumerRebalanceListener { @Override public void onPartitionsAssigned(final Collection partitions) { - streamThread.maybeHandleAssignmentFromStreamsRebalanceProtocol(); - // NB: all task management is already handled by: // org.apache.kafka.streams.processor.internals.StreamsPartitionAssignor.onAssignment if (assignmentErrorCode.get() == AssignorError.INCOMPLETE_SOURCE_TOPIC_METADATA.code()) {