content: /*
   * 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;
  
  import org.apache.kafka.clients.consumer.CloseOptions;
  import org.apache.kafka.clients.consumer.ConsumerConfig;
  import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
  import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
  import org.apache.kafka.clients.consumer.internals.events.CompletableBackgroundEvent;
  import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
  import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
  import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
  import org.apache.kafka.clients.consumer.internals.metrics.ConsumerRebalanceMetricsManager;
  import org.apache.kafka.clients.consumer.internals.metrics.RebalanceMetricsManager;
  import org.apache.kafka.common.KafkaException;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
  import org.apache.kafka.common.metrics.Metrics;
  import org.apache.kafka.common.protocol.Errors;
  import org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest;
  import org.apache.kafka.common.requests.ConsumerGroupHeartbeatResponse;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.common.utils.Time;
  
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.SortedSet;
  import java.util.TreeSet;
  import java.util.concurrent.CompletableFuture;
  import java.util.stream.Collectors;
  
  import static org.apache.kafka.clients.consumer.CloseOptions.GroupMembershipOperation.DEFAULT;
  import static org.apache.kafka.clients.consumer.CloseOptions.GroupMembershipOperation.LEAVE_GROUP;
  import static org.apache.kafka.clients.consumer.CloseOptions.GroupMembershipOperation.REMAIN_IN_GROUP;
  import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED;
  import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_LOST;
  import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_REVOKED;
  
  /**
   * Group manager for a single consumer that has a group id defined in the config
   * {@link ConsumerConfig#GROUP_ID_CONFIG}, to use the Kafka-based offset management capability,
   * and the consumer group protocol to get automatically assigned partitions when calling the
   * subscribe API.
   *
   * <p/>
   *
   * While the subscribe API hasn't been called (or if the consumer called unsubscribe), this manager
   * will only be responsible for keeping the member in the {@link MemberState#UNSUBSCRIBED} state,
   * where it can commit offsets to the group identified by the {@link #groupId()}, without joining
   * the group.
   *
   * <p/>
   *
   * If the consumer subscribe API is called, this manager will use the {@link #groupId()} to join the
   * consumer group, and based on the consumer group protocol heartbeats, will handle the full
   * lifecycle of the member as it joins the group, reconciles assignments, handles fencing and
   * fatal errors, and leaves the group.
   *
   * <p/>
   *
   * Reconciliation process:<p/>
   * The member accepts all assignments received from the broker, resolves topic names from
   * metadata, reconciles the resolved assignments, and keeps the unresolved to be reconciled when
   * discovered with a metadata update. Reconciliations of resolved assignments are executed
   * sequentially and acknowledged to the server as they complete. The reconciliation process
   * involves multiple async operations, so the member will continue to heartbeat while these
   * operations complete, to make sure that the member stays in the group while reconciling.
   *
   * <p/>
   *
   * Reconciliation steps:
   * <ol>
   *     <li>Resolve topic names for all topic IDs received in the target assignment. Topic names
   *     found in metadata are then ready to be reconciled. Topic IDs not found are kept as
   *     unresolved, and the member request metadata updates until it resolves them (or the broker
   *     removes it from the target assignment.</li>
   *     <li>Commit offsets if auto-commit is enabled.</li>
   *     <li>Invoke the user-defined onPartitionsRevoked listener.</li>
   *     <li>Invoke the user-defined onPartitionsAssigned listener.</li>
   *     <li>When the above steps complete, the member acknowledges the reconciled assignment,
   *     which is the subset of the target that was resolved from metadata and actually reconciled.
   *     The ack is performed by sending a heartbeat request back to the broker, including the
   *     reconciled assignment.</li>
   * </ol>
   *
   * Note that user-defined callbacks are triggered from this manager that runs in the
   * BackgroundThread, but executed in the Application Thread, where a failure will be returned to
   * the user if the callbacks fail. This manager is only concerned about the callbacks completion to
   * know that it can proceed with the reconciliation.
   */
  public class ConsumerMembershipManager extends AbstractMembershipManager<ConsumerGroupHeartbeatResponse> {
  
      /**
       * Group instance ID to be used by the member, provided when creating the current membership manager.
       */
      protected final Optional<String> groupInstanceId;
  
      private final Optional<String> rackId;
  
      /**
       * Rebalance timeout. To be used as time limit for the commit request issued
       * when a new assignment is received, that is retried until it succeeds, fails with a
       * non-retriable error, it the time limit expires.
       */
      private final int rebalanceTimeoutMs;
  
      /**
       * Name of the server-side assignor this member has configured to use. It will be sent
       * out to the server on the {@link ConsumerGroupHeartbeatRequest}. If not defined, the server
       * will select the assignor implementation to use.
       */
      private final Optional<String> serverAssignor;
  
      /**
       * Manager to perform commit requests needed before rebalance (if auto-commit is enabled)
       */
      private final CommitRequestManager commitRequestManager;
  
      /**
       * Serves as the conduit by which we can report events to the application thread. This is needed as we send
       * {@link ConsumerRebalanceListenerCallbackNeededEvent callbacks} and, if needed,
       * {@link ErrorEvent errors} to the application thread.
       */
      private final BackgroundEventHandler backgroundEventHandler;
  
      public ConsumerMembershipManager(String groupId,
                                       Optional<String> groupInstanceId,
                                       Optional<String> rackId,
                                       int rebalanceTimeoutMs,
                                       Optional<String> serverAssignor,
                                       SubscriptionState subscriptions,
                                       CommitRequestManager commitRequestManager,
                                       ConsumerMetadata metadata,
                                       LogContext logContext,
                                       BackgroundEventHandler backgroundEventHandler,
                                       Time time,
                                       Metrics metrics,
                                       boolean autoCommitEnabled) {
          this(groupId,
              groupInstanceId,
              rackId,
              rebalanceTimeoutMs,
              serverAssignor,
              subscriptions,
              commitRequestManager,
              metadata,
              logContext,
              backgroundEventHandler,
              time,
              new ConsumerRebalanceMetricsManager(metrics),
              autoCommitEnabled);
      }
  
      // Visible for testing
      ConsumerMembershipManager(String groupId,
                                Optional<String> groupInstanceId,
                                Optional<String> rackId,
                                int rebalanceTimeoutMs,
                                Optional<String> serverAssignor,
                                SubscriptionState subscriptions,
                                CommitRequestManager commitRequestManager,
                                ConsumerMetadata metadata,
                                LogContext logContext,
                                BackgroundEventHandler backgroundEventHandler,
                                Time time,
                                RebalanceMetricsManager metricsManager,
                                boolean autoCommitEnabled) {
          super(groupId,
              subscriptions,
              metadata,
              logContext.logger(ConsumerMembershipManager.class),
              time,
              metricsManager,
              autoCommitEnabled);
          this.groupInstanceId = groupInstanceId;
          this.rackId = rackId;
          this.rebalanceTimeoutMs = rebalanceTimeoutMs;
          this.serverAssignor = serverAssignor;
          this.commitRequestManager = commitRequestManager;
          this.backgroundEventHandler = backgroundEventHandler;
      }
  
      /**
       * @return Instance ID used by the member when joining the group. If non-empty, it will indicate that
       * this is a static member.
       */
      public Optional<String> groupInstanceId() {
          return groupInstanceId;
      }
  
      public Optional<String> rackId() {
          return rackId;
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public void onHeartbeatSuccess(ConsumerGroupHeartbeatResponse response) {
          ConsumerGroupHeartbeatResponseData responseData = response.data();
          if (responseData.errorCode() != Errors.NONE.code()) {
              String errorMessage = String.format(
                      "Unexpected error in Heartbeat response. Expected no error, but received: %s",
                      Errors.forCode(responseData.errorCode())
              );
              throw new IllegalArgumentException(errorMessage);
          }
          MemberState state = state();
          if (state == MemberState.LEAVING) {
              log.debug("Ignoring heartbeat response received from broker. Member {} with epoch {} is " +
                      "already leaving the group.", memberId, memberEpoch);
              return;
          }
          if (state == MemberState.UNSUBSCRIBED && responseData.memberEpoch() < 0 && maybeCompleteLeaveInProgress()) {
              log.debug("Member {} with epoch {} received a successful response to the heartbeat " +
                      "to leave the group and completed the leave operation. ", memberId, memberEpoch);
              return;
          }
          if (isNotInGroup()) {
              log.debug("Ignoring heartbeat response received from broker. Member {} is in {} state" +
                      " so it's not a member of the group. ", memberId, state);
              return;
          }
          if (responseData.memberEpoch() < 0) {
              log.debug("Ignoring heartbeat response received from broker. Member {} with epoch {} " +
                      "is in {} state and the member epoch is invalid: {}. ", memberId, memberEpoch, state,
                      responseData.memberEpoch());
              maybeCompleteLeaveInProgress();
              return;
          }
  
          updateMemberEpoch(responseData.memberEpoch());
  
          ConsumerGroupHeartbeatResponseData.Assignment assignment = responseData.assignment();
  
          if (assignment != null) {
              if (!state.canHandleNewAssignment()) {
                  // New assignment received but member is in a state where it cannot take new
                  // assignments (ex. preparing to leave the group)
                  log.debug("Ignoring new assignment {} received from server because member is in {} state.",
                          assignment, state);
                  return;
              }
  
              Map<Uuid, SortedSet<Integer>> newAssignment = new HashMap<>();
              assignment.topicPartitions().forEach(topicPartition ->
                  newAssignment.put(topicPartition.topicId(), new TreeSet<>(topicPartition.partitions())));
              processAssignmentReceived(newAssignment);
          }
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      protected CompletableFuture<Void> signalReconciliationStarted() {
          // Issue a commit request that will be retried until it succeeds, fails with a
          // non-retriable error, or the time limit expires. Retry on stale member epoch error, in a
          // best effort to commit the offsets in the case where the epoch might have changed while
          // the current reconciliation is in process. Note this is using the rebalance timeout as
          // it is the limit enforced by the broker to complete the reconciliation process.
          return commitRequestManager.maybeAutoCommitSyncBeforeRebalance(getDeadlineMsForTimeout(rebalanceTimeoutMs));
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      protected void signalReconciliationCompleting() {
          // Reschedule the auto commit starting from now that the member has a new assignment.
          commitRequestManager.resetAutoCommitTimer();
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      protected CompletableFuture<Void> signalMemberLeavingGroup() {
          return invokeOnPartitionsRevokedOrLostToReleaseAssignment();
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      protected CompletableFuture<Void> signalPartitionsLost(Set<TopicPartition> partitionsLost) {
          return invokeOnPartitionsLostCallback(partitionsLost);
      }
  
      /**
       * Release member assignment by calling the user defined callbacks for onPartitionsRevoked or
       * onPartitionsLost.
       * <ul>
       *     <li>If the member is part of the group (epoch > 0), this will invoke onPartitionsRevoked.
       *     This will be the case when releasing assignment because the member is intentionally
       *     leaving the group (after a call to unsubscribe)</li>
       *
       *     <li>If the member is not part of the group (epoch <=0), this will invoke onPartitionsLost.
       *     This will be the case when releasing assignment after being fenced .</li>
       * </ul>
       *
       * @return Future that will complete when the callback execution completes.
       */
      private CompletableFuture<Void> invokeOnPartitionsRevokedOrLostToReleaseAssignment() {
          SortedSet<TopicPartition> droppedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
          droppedPartitions.addAll(subscriptions.assignedPartitions());
  
          log.info("Member {} is triggering callbacks to release assignment {} and leave group",
                  memberId, droppedPartitions);
  
          CompletableFuture<Void> callbackResult;
          if (droppedPartitions.isEmpty()) {
              // No assignment to release.
              callbackResult = CompletableFuture.completedFuture(null);
          } else {
              // Release assignment.
              if (memberEpoch > 0) {
                  // Member is part of the group. Invoke onPartitionsRevoked.
                  callbackResult = revokePartitions(droppedPartitions);
              } else {
                  // Member is not part of the group anymore. Invoke onPartitionsLost.
                  callbackResult = invokeOnPartitionsLostCallback(droppedPartitions);
              }
          }
          return callbackResult;
      }
  
      /**
       * @return Server-side assignor implementation configured for the member, that will be sent
       * out to the server to be used. If empty, then the server will select the assignor.
       */
      public Optional<String> serverAssignor() {
          return this.serverAssignor;
      }
  
      private CompletableFuture<Void> invokeOnPartitionsRevokedCallback(Set<TopicPartition> partitionsRevoked) {
          // This should not trigger the callback if partitionsRevoked is empty, to keep the
          // current behaviour.
          Optional<ConsumerRebalanceListener> listener = subscriptions.rebalanceListener();
          if (!partitionsRevoked.isEmpty() && listener.isPresent()) {
              return enqueueConsumerRebalanceListenerCallback(ON_PARTITIONS_REVOKED, partitionsRevoked);
          } else {
              return CompletableFuture.completedFuture(null);
          }
      }
  
      private CompletableFuture<Void> invokeOnPartitionsAssignedCallback(Set<TopicPartition> partitionsAssigned) {
          // This should always trigger the callback, even if partitionsAssigned is empty, to keep
          // the current behaviour.
          Optional<ConsumerRebalanceListener> listener = subscriptions.rebalanceListener();
          if (listener.isPresent()) {
              return enqueueConsumerRebalanceListenerCallback(ON_PARTITIONS_ASSIGNED, partitionsAssigned);
          } else {
              return CompletableFuture.completedFuture(null);
          }
      }
  
      private CompletableFuture<Void> invokeOnPartitionsLostCallback(Set<TopicPartition> partitionsLost) {
          // This should not trigger the callback if partitionsLost is empty, to keep the current
          // behaviour.
          Optional<ConsumerRebalanceListener> listener = subscriptions.rebalanceListener();
          if (!partitionsLost.isEmpty() && listener.isPresent()) {
              return enqueueConsumerRebalanceListenerCallback(ON_PARTITIONS_LOST, partitionsLost);
          } else {
              return CompletableFuture.completedFuture(null);
          }
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public CompletableFuture<Void> signalPartitionsAssigned(Set<TopicPartition> partitionsAssigned) {
          return invokeOnPartitionsAssignedCallback(partitionsAssigned);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public void signalPartitionsBeingRevoked(Set<TopicPartition> partitionsToRevoke) {
          logPausedPartitionsBeingRevoked(partitionsToRevoke);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public CompletableFuture<Void> signalPartitionsRevoked(Set<TopicPartition> partitionsRevoked) {
          return invokeOnPartitionsRevokedCallback(partitionsRevoked);
      }
  
      /**
       * Log partitions being revoked that were already paused, since the pause flag will be
       * effectively lost.
       */
      private void logPausedPartitionsBeingRevoked(Set<TopicPartition> partitionsToRevoke) {
          Set<TopicPartition> revokePausedPartitions = subscriptions.pausedPartitions();
          revokePausedPartitions.retainAll(partitionsToRevoke);
          if (!revokePausedPartitions.isEmpty()) {
              log.info("The pause flag in partitions [{}] will be removed due to revocation.", revokePausedPartitions.stream().map(TopicPartition::toString).collect(Collectors.joining(", ")));
          }
      }
  
      @Override
      public boolean isLeavingGroup() {
          CloseOptions.GroupMembershipOperation leaveGroupOperation = leaveGroupOperation();
          if (REMAIN_IN_GROUP == leaveGroupOperation) {
              return false;
          }
  
          MemberState state = state();
          boolean isLeavingState = state == MemberState.PREPARE_LEAVING || state == MemberState.LEAVING;
  
          // Default operation: both static and dynamic consumers will send a leave heartbeat
          boolean hasLeaveOperation = DEFAULT == leaveGroupOperation ||
              // Leave operation: both static and dynamic consumers will send a leave heartbeat
              LEAVE_GROUP == leaveGroupOperation ||
              // Remain in group: only static consumers will send a leave heartbeat, while dynamic members will not
              groupInstanceId().isPresent();
  
          return isLeavingState && hasLeaveOperation;
      }
  
      /**
       * Enqueue a {@link ConsumerRebalanceListenerCallbackNeededEvent} to trigger the execution of the
       * appropriate {@link ConsumerRebalanceListener} {@link ConsumerRebalanceListenerMethodName method} on the
       * application thread.
       *
       * <p/>
       *
       * Because the reconciliation process (run in the background thread) will be blocked by the application thread
       * until it completes this, we need to provide a {@link CompletableFuture} by which to remember where we left off.
       *
       * @param methodName Callback method that needs to be executed on the application thread
       * @param partitions Partitions to supply to the callback method
       * @return Future that will be chained within the rest of the reconciliation logic
       */
      private CompletableFuture<Void> enqueueConsumerRebalanceListenerCallback(ConsumerRebalanceListenerMethodName methodName,
                                                                               Set<TopicPartition> partitions) {
          SortedSet<TopicPartition> sortedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
          sortedPartitions.addAll(partitions);
  
          CompletableBackgroundEvent<Void> event = new ConsumerRebalanceListenerCallbackNeededEvent(methodName, sortedPartitions);
          backgroundEventHandler.add(event);
          log.debug("The event to trigger the {} method execution was enqueued successfully", methodName.fullyQualifiedMethodName());
          return event.future();
      }
  
      /**
       * Signals that a {@link ConsumerRebalanceListener} callback has completed. This is invoked when the
       * application thread has completed the callback and has submitted a
       * {@link ConsumerRebalanceListenerCallbackCompletedEvent} to the network I/O thread. At this point, we
       * notify the state machine that it's complete so that it can move to the next appropriate step of the
       * rebalance process.
       *
       * @param event Event with details about the callback that was executed
       */
      public void consumerRebalanceListenerCallbackCompleted(ConsumerRebalanceListenerCallbackCompletedEvent event) {
          ConsumerRebalanceListenerMethodName methodName = event.methodName();
          Optional<KafkaException> error = event.error();
          CompletableFuture<Void> future = event.future();
  
          if (error.isPresent()) {
              Exception e = error.get();
              log.warn(
                      "The {} method completed with an error ({}); signaling to continue to the next phase of rebalance",
                      methodName.fullyQualifiedMethodName(),
                      e.getMessage()
              );
  
              future.completeExceptionally(e);
          } else {
              log.debug(
                      "The {} method completed successfully; signaling to continue to the next phase of rebalance",
                      methodName.fullyQualifiedMethodName()
              );
  
              future.complete(null);
          }
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public int joinGroupEpoch() {
          return ConsumerGroupHeartbeatRequest.JOIN_GROUP_MEMBER_EPOCH;
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public int leaveGroupEpoch() {
          boolean isStaticMember = groupInstanceId.isPresent();
          // Currently, the server doesn't have a mechanism for static members to permanently leave the group.
          // Therefore, we use LEAVE_GROUP_MEMBER_EPOCH to force the GroupMetadataManager to fence
          // this member, effectively removing it from the group.
          if (LEAVE_GROUP == leaveGroupOperation) {
              return ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
          }
  
          return isStaticMember ?
              ConsumerGroupHeartbeatRequest.LEAVE_GROUP_STATIC_MEMBER_EPOCH :
              ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
      }
  }
  