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.connect.runtime.distributed;
  
  import org.apache.kafka.clients.GroupRebalanceConfig;
  import org.apache.kafka.clients.consumer.CloseOptions;
  import org.apache.kafka.clients.consumer.internals.AbstractCoordinator;
  import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
  import org.apache.kafka.common.metrics.Measurable;
  import org.apache.kafka.common.metrics.Metrics;
  import org.apache.kafka.common.requests.JoinGroupRequest;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.common.utils.Timer;
  import org.apache.kafka.connect.storage.ClusterConfigState;
  import org.apache.kafka.connect.storage.ConfigBackingStore;
  import org.apache.kafka.connect.util.ConnectorTaskId;
  
  import org.slf4j.Logger;
  
  import java.io.Closeable;
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.HashMap;
  import java.util.LinkedHashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Set;
  import java.util.function.Supplier;
  
  import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection;
  import static org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
  import static org.apache.kafka.common.utils.Utils.UncheckedCloseable;
  import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.EAGER;
  
  
  /**
   * This class manages the coordination process with the Kafka group coordinator on the broker for managing assignments
   * to workers.
   */
  public class WorkerCoordinator extends AbstractCoordinator implements Closeable {
      private final Logger log;
      private final String restUrl;
      private final ConfigBackingStore configStorage;
      private volatile ExtendedAssignment assignmentSnapshot;
      private ClusterConfigState configSnapshot;
      private final WorkerRebalanceListener listener;
      private final ConnectProtocolCompatibility protocolCompatibility;
      private LeaderState leaderState;
  
      private boolean rejoinRequested;
      private volatile ConnectProtocolCompatibility currentConnectProtocol;
      private volatile int lastCompletedGenerationId;
      private final ConnectAssignor eagerAssignor;
      private final ConnectAssignor incrementalAssignor;
      private final int coordinatorDiscoveryTimeoutMs;
  
      /**
       * Initialize the coordination manager.
       */
      public WorkerCoordinator(GroupRebalanceConfig config,
                               LogContext logContext,
                               ConsumerNetworkClient client,
                               Metrics metrics,
                               String metricGrpPrefix,
                               Time time,
                               String restUrl,
                               ConfigBackingStore configStorage,
                               WorkerRebalanceListener listener,
                               ConnectProtocolCompatibility protocolCompatibility,
                               int maxDelay) {
          super(config,
                logContext,
                client,
                metrics,
                metricGrpPrefix,
                time);
          this.log = logContext.logger(WorkerCoordinator.class);
          this.restUrl = restUrl;
          this.configStorage = configStorage;
          this.assignmentSnapshot = null;
          new WorkerCoordinatorMetrics(metrics, metricGrpPrefix);
          this.listener = listener;
          this.rejoinRequested = false;
          this.protocolCompatibility = protocolCompatibility;
          this.incrementalAssignor = new IncrementalCooperativeAssignor(logContext, time, maxDelay);
          this.eagerAssignor = new EagerAssignor(logContext);
          this.currentConnectProtocol = protocolCompatibility;
          this.coordinatorDiscoveryTimeoutMs = config.heartbeatIntervalMs;
          this.lastCompletedGenerationId = Generation.NO_GENERATION.generationId;
      }
  
      @Override
      public void requestRejoin(final String reason) {
          log.debug("Request joining group due to: {}", reason);
          rejoinRequested = true;
      }
  
      @Override
      public String protocolType() {
          return "connect";
      }
  
      // expose for tests
      @Override
      protected synchronized boolean ensureCoordinatorReady(final Timer timer) {
          return super.ensureCoordinatorReady(timer);
      }
  
      public void poll(long timeout, Supplier<UncheckedCloseable> onPoll) {
          // poll for io until the timeout expires
          final long start = time.milliseconds();
          long now = start;
          long remaining;
  
          do {
              if (coordinatorUnknown()) {
                  log.debug("Broker coordinator is marked unknown. Attempting discovery with a timeout of {}ms",
                          coordinatorDiscoveryTimeoutMs);
                  boolean coordinatorReady;
                  try (UncheckedCloseable polling = onPoll.get()) {
                      coordinatorReady = ensureCoordinatorReady(time.timer(coordinatorDiscoveryTimeoutMs));
                  }
                  if (coordinatorReady) {
                      log.debug("Broker coordinator is ready");
                  } else {
                      log.debug("Can not connect to broker coordinator");
                      final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
                      if (localAssignmentSnapshot != null && !localAssignmentSnapshot.failed()) {
                          log.info("Broker coordinator was unreachable for {}ms. Revoking previous assignment {} to " +
                                  "avoid running tasks while not being a member the group", coordinatorDiscoveryTimeoutMs, localAssignmentSnapshot);
                          listener.onRevoked(localAssignmentSnapshot.leader(), localAssignmentSnapshot.connectors(), localAssignmentSnapshot.tasks());
                          assignmentSnapshot = null;
                      }
                  }
                  now = time.milliseconds();
              }
  
              if (rejoinNeededOrPending()) {
                  try (UncheckedCloseable polling = onPoll.get()) {
                      ensureActiveGroup();
                  }
                  now = time.milliseconds();
              }
  
              pollHeartbeat(now);
  
              long elapsed = now - start;
              remaining = timeout - elapsed;
  
              // Note that because the network client is shared with the background heartbeat thread,
              // we do not want to block in poll longer than the time to the next heartbeat.
              long pollTimeout = Math.min(Math.max(0, remaining), timeToNextHeartbeat(now));
              try (UncheckedCloseable polling = onPoll.get()) {
                  client.poll(time.timer(pollTimeout));
              }
  
              now = time.milliseconds();
              elapsed = now - start;
              remaining = timeout - elapsed;
          } while (remaining > 0);
      }
  
      @Override
      public JoinGroupRequestProtocolCollection metadata() {
          configSnapshot = configStorage.snapshot();
          final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
          ExtendedWorkerState workerState = new ExtendedWorkerState(restUrl, configSnapshot.offset(), localAssignmentSnapshot);
          return switch (protocolCompatibility) {
              case EAGER -> ConnectProtocol.metadataRequest(workerState);
              case COMPATIBLE -> IncrementalCooperativeConnectProtocol.metadataRequest(workerState, false);
              case SESSIONED -> IncrementalCooperativeConnectProtocol.metadataRequest(workerState, true);
          };
      }
  
      @Override
      protected void onJoinComplete(int generation, String memberId, String protocol, ByteBuffer memberAssignment) {
          ExtendedAssignment newAssignment = IncrementalCooperativeConnectProtocol.deserializeAssignment(memberAssignment);
          log.debug("Deserialized new assignment: {}", newAssignment);
          currentConnectProtocol = ConnectProtocolCompatibility.fromProtocol(protocol);
          // At this point we always consider ourselves to be a member of the cluster, even if there was an assignment
          // error (the leader couldn't make the assignment) or we are behind the config and cannot yet work on our assigned
          // tasks. It's the responsibility of the code driving this process to decide how to react (e.g. trying to get
          // up to date, try to rejoin again, leaving the group and backing off, etc.).
          rejoinRequested = false;
          if (currentConnectProtocol != EAGER) {
              if (!newAssignment.revokedConnectors().isEmpty() || !newAssignment.revokedTasks().isEmpty()) {
                  listener.onRevoked(newAssignment.leader(), newAssignment.revokedConnectors(), newAssignment.revokedTasks());
              }
  
              final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
              if (localAssignmentSnapshot != null) {
                  localAssignmentSnapshot.connectors().removeAll(newAssignment.revokedConnectors());
                  localAssignmentSnapshot.tasks().removeAll(newAssignment.revokedTasks());
                  log.debug("After revocations snapshot of assignment: {}", localAssignmentSnapshot);
                  newAssignment.connectors().addAll(localAssignmentSnapshot.connectors());
                  newAssignment.tasks().addAll(localAssignmentSnapshot.tasks());
              }
              log.debug("Augmented new assignment: {}", newAssignment);
          }
          assignmentSnapshot = newAssignment;
          lastCompletedGenerationId = generation;
          listener.onAssigned(newAssignment, generation);
      }
  
      @Override
      protected Map<String, ByteBuffer> onLeaderElected(String leaderId,
                                                        String protocol,
                                                        List<JoinGroupResponseMember> allMemberMetadata,
                                                        boolean skipAssignment) {
          if (skipAssignment)
              throw new IllegalStateException("Can't skip assignment because Connect does not support static membership.");
  
          return ConnectProtocolCompatibility.fromProtocol(protocol) == EAGER
                 ? eagerAssignor.performAssignment(leaderId, protocol, allMemberMetadata, this)
                 : incrementalAssignor.performAssignment(leaderId, protocol, allMemberMetadata, this);
      }
  
      @Override
      protected boolean onJoinPrepare(Timer timer, int generation, String memberId) {
          log.info("Rebalance started");
          leaderState(null);
          final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
          if (currentConnectProtocol == EAGER) {
              log.debug("Revoking previous assignment {}", localAssignmentSnapshot);
              if (localAssignmentSnapshot != null && !localAssignmentSnapshot.failed())
                  listener.onRevoked(localAssignmentSnapshot.leader(), localAssignmentSnapshot.connectors(), localAssignmentSnapshot.tasks());
          } else {
              log.debug("Cooperative rebalance triggered. Keeping assignment {} until it's "
                        + "explicitly revoked.", localAssignmentSnapshot);
          }
          return true;
      }
  
      @Override
      protected boolean rejoinNeededOrPending() {
          final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
          return super.rejoinNeededOrPending() || (localAssignmentSnapshot == null || localAssignmentSnapshot.failed()) || rejoinRequested;
      }
  
      @Override
      public String memberId() {
          Generation generation = generationIfStable();
          if (generation != null)
              return generation.memberId;
          return JoinGroupRequest.UNKNOWN_MEMBER_ID;
      }
  
      @Override
      protected void handlePollTimeoutExpiry() {
          listener.onPollTimeoutExpiry();
          maybeLeaveGroup(CloseOptions.GroupMembershipOperation.DEFAULT, "worker poll timeout has expired.");
      }
  
      /**
       * Return the current generation. The generation refers to this worker's knowledge with
       * respect to which generation is the latest one and, therefore, this information is local.
       *
       * @return the generation ID or -1 if no generation is defined
       */
      public int generationId() {
          return super.generation().generationId;
      }
  
      /**
       * Return id that corresponds to the group generation that was active when the last join was successful
       *
       * @return the generation ID of the last group that was joined successfully by this member or -1 if no generation
       * was stable at that point
       */
      public int lastCompletedGenerationId() {
          return lastCompletedGenerationId;
      }
  
      public void revokeAssignment(ExtendedAssignment assignment) {
          listener.onRevoked(assignment.leader(), assignment.connectors(), assignment.tasks());
      }
  
      private boolean isLeader() {
          final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
          return localAssignmentSnapshot != null && memberId().equals(localAssignmentSnapshot.leader());
      }
  
      public String ownerUrl(String connector) {
          if (rejoinNeededOrPending() || !isLeader())
              return null;
          return leaderState().ownerUrl(connector);
      }
  
      public String ownerUrl(ConnectorTaskId task) {
          if (rejoinNeededOrPending() || !isLeader())
              return null;
          return leaderState().ownerUrl(task);
      }
  
      /**
       * Get an up-to-date snapshot of the cluster configuration.
       *
       * @return the state of the cluster configuration; the result is not locally cached
       */
      public ClusterConfigState configFreshSnapshot() {
          return configStorage.snapshot();
      }
  
      /**
       * Get a snapshot of the cluster configuration.
       *
       * @return the state of the cluster configuration
       */
      public ClusterConfigState configSnapshot() {
          return configSnapshot;
      }
  
      /**
       * Set the state of the cluster configuration to this worker coordinator.
       *
       * @param update the updated state of the cluster configuration
       */
      public void configSnapshot(ClusterConfigState update) {
          configSnapshot = update;
      }
  
      /**
       * Get the leader state stored in this worker coordinator.
       *
       * @return the leader state
       */
      private LeaderState leaderState() {
          return leaderState;
      }
  
      /**
       * Store the leader state to this worker coordinator.
       *
       * @param update the updated leader state
       */
      public void leaderState(LeaderState update) {
          leaderState = update;
      }
  
      /**
       * Get the version of the connect protocol that is currently active in the group of workers.
       *
       * @return the current connect protocol version
       */
      public short currentProtocolVersion() {
          return currentConnectProtocol.protocolVersion();
      }
  
      private class WorkerCoordinatorMetrics {
          public final String metricGrpName;
  
          public WorkerCoordinatorMetrics(Metrics metrics, String metricGrpPrefix) {
              this.metricGrpName = metricGrpPrefix + "-coordinator-metrics";
  
              Measurable numConnectors = (config, now) -> {
                  final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
                  if (localAssignmentSnapshot == null) {
                      return 0.0;
                  }
                  return localAssignmentSnapshot.connectors().size();
              };
  
              Measurable numTasks = (config, now) -> {
                  final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
                  if (localAssignmentSnapshot == null) {
                      return 0.0;
                  }
                  return localAssignmentSnapshot.tasks().size();
              };
  
              metrics.addMetric(metrics.metricName("assigned-connectors",
                                this.metricGrpName,
                                "The number of connector instances currently assigned to this worker"), numConnectors);
              metrics.addMetric(metrics.metricName("assigned-tasks",
                                this.metricGrpName,
                                "The number of tasks currently assigned to this worker"), numTasks);
          }
      }
  
      public static <K, V> Map<V, K> invertAssignment(Map<K, Collection<V>> assignment) {
          Map<V, K> inverted = new HashMap<>();
          for (Map.Entry<K, Collection<V>> assignmentEntry : assignment.entrySet()) {
              K key = assignmentEntry.getKey();
              for (V value : assignmentEntry.getValue())
                  inverted.put(value, key);
          }
          return inverted;
      }
  
      public static class LeaderState {
          private final Map<String, ExtendedWorkerState> allMembers;
          private final Map<String, String> connectorOwners;
          private final Map<ConnectorTaskId, String> taskOwners;
  
          public LeaderState(Map<String, ExtendedWorkerState> allMembers,
                             Map<String, Collection<String>> connectorAssignment,
                             Map<String, Collection<ConnectorTaskId>> taskAssignment) {
              this.allMembers = allMembers;
              this.connectorOwners = invertAssignment(connectorAssignment);
              this.taskOwners = invertAssignment(taskAssignment);
          }
  
          private String ownerUrl(ConnectorTaskId id) {
              String ownerId = taskOwners.get(id);
              if (ownerId == null)
                  return null;
              return allMembers.get(ownerId).url();
          }
  
          private String ownerUrl(String connector) {
              String ownerId = connectorOwners.get(connector);
              if (ownerId == null)
                  return null;
              return allMembers.get(ownerId).url();
          }
  
          @Override
          public boolean equals(Object o) {
              if (this == o) return true;
              if (!(o instanceof LeaderState that)) return false;
              return Objects.equals(allMembers, that.allMembers)
                      && Objects.equals(connectorOwners, that.connectorOwners)
                      && Objects.equals(taskOwners, that.taskOwners);
          }
  
          @Override
          public int hashCode() {
              return Objects.hash(allMembers, connectorOwners, taskOwners);
          }
  
          @Override
          public String toString() {
              return "LeaderState{"
                      + "allMembers=" + allMembers
                      + ", connectorOwners=" + connectorOwners
                      + ", taskOwners=" + taskOwners
                      + '}';
          }
      }
  
      public static class ConnectorsAndTasks {
          public static final ConnectorsAndTasks EMPTY =
                  new ConnectorsAndTasks(Collections.emptyList(), Collections.emptyList());
  
          private final Collection<String> connectors;
          private final Collection<ConnectorTaskId> tasks;
  
          private ConnectorsAndTasks(Collection<String> connectors, Collection<ConnectorTaskId> tasks) {
              this.connectors = connectors;
              this.tasks = tasks;
          }
  
          public static class Builder {
              private Set<String> withConnectors = new LinkedHashSet<>();
              private Set<ConnectorTaskId> withTasks = new LinkedHashSet<>();
  
              public Builder() {
              }
  
              public ConnectorsAndTasks.Builder with(Collection<String> connectors,
                                                     Collection<ConnectorTaskId> tasks) {
                  withConnectors = new LinkedHashSet<>(connectors);
                  withTasks = new LinkedHashSet<>(tasks);
                  return this;
              }
  
              public ConnectorsAndTasks.Builder addConnectors(Collection<String> connectors) {
                  this.withConnectors.addAll(connectors);
                  return this;
              }
  
              public ConnectorsAndTasks.Builder addTasks(Collection<ConnectorTaskId> tasks) {
                  this.withTasks.addAll(tasks);
                  return this;
              }
  
              public ConnectorsAndTasks.Builder addAll(ConnectorsAndTasks connectorsAndTasks) {
                  return this
                          .addConnectors(connectorsAndTasks.connectors())
                          .addTasks(connectorsAndTasks.tasks());
              }
  
              public ConnectorsAndTasks build() {
                  return new ConnectorsAndTasks(withConnectors, withTasks);
              }
          }
  
          public Collection<String> connectors() {
              return connectors;
          }
  
          public Collection<ConnectorTaskId> tasks() {
              return tasks;
          }
  
          public int size() {
              return connectors.size() + tasks.size();
          }
  
          public boolean isEmpty() {
              return connectors.isEmpty() && tasks.isEmpty();
          }
  
          @Override
          public String toString() {
              return "{ connectorIds=" + connectors + ", taskIds=" + tasks + '}';
          }
      }
  
      public static class WorkerLoad {
          private final String worker;
          private final Collection<String> connectors;
          private final Collection<ConnectorTaskId> tasks;
  
          private WorkerLoad(
                  String worker,
                  Collection<String> connectors,
                  Collection<ConnectorTaskId> tasks
          ) {
              this.worker = worker;
              this.connectors = connectors;
              this.tasks = tasks;
          }
  
          public static class Builder {
              private final String withWorker;
              private Collection<String> withConnectors;
              private Collection<ConnectorTaskId> withTasks;
  
              public Builder(String worker) {
                  this.withWorker = Objects.requireNonNull(worker, "worker cannot be null");
              }
  
              public WorkerLoad.Builder withCopies(Collection<String> connectors,
                                                   Collection<ConnectorTaskId> tasks) {
                  withConnectors = new ArrayList<>(
                          Objects.requireNonNull(connectors, "connectors may be empty but not null"));
                  withTasks = new ArrayList<>(
                          Objects.requireNonNull(tasks, "tasks may be empty but not null"));
                  return this;
              }
  
              public WorkerLoad.Builder with(Collection<String> connectors,
                                             Collection<ConnectorTaskId> tasks) {
                  withConnectors = Objects.requireNonNull(connectors,
                          "connectors may be empty but not null");
                  withTasks = Objects.requireNonNull(tasks, "tasks may be empty but not null");
                  return this;
              }
  
              public WorkerLoad build() {
                  return new WorkerLoad(
                          withWorker,
                          withConnectors != null ? withConnectors : new ArrayList<>(),
                          withTasks != null ? withTasks : new ArrayList<>());
              }
          }
  
          public String worker() {
              return worker;
          }
  
          public Collection<String> connectors() {
              return connectors;
          }
  
          public Collection<ConnectorTaskId> tasks() {
              return tasks;
          }
  
          public int connectorsSize() {
              return connectors.size();
          }
  
          public int tasksSize() {
              return tasks.size();
          }
  
          public void assign(String connector) {
              connectors.add(connector);
          }
  
          public void assign(ConnectorTaskId task) {
              tasks.add(task);
          }
  
          public int size() {
              return connectors.size() + tasks.size();
          }
  
          public boolean isEmpty() {
              return connectors.isEmpty() && tasks.isEmpty();
          }
  
          public static Comparator<WorkerLoad> connectorComparator() {
              return (left, right) -> {
                  int res = left.connectors.size() - right.connectors.size();
                  return res != 0 ? res : left.worker == null
                                          ? right.worker == null ? 0 : -1
                                          : left.worker.compareTo(right.worker);
              };
          }
  
          public static Comparator<WorkerLoad> taskComparator() {
              return (left, right) -> {
                  int res = left.tasks.size() - right.tasks.size();
                  return res != 0 ? res : left.worker == null
                                          ? right.worker == null ? 0 : -1
                                          : left.worker.compareTo(right.worker);
              };
          }
  
          @Override
          public String toString() {
              return "{ worker=" + worker + ", connectorIds=" + connectors + ", taskIds=" + tasks + '}';
          }
  
          @Override
          public boolean equals(Object o) {
              if (this == o) {
                  return true;
              }
              if (!(o instanceof WorkerLoad that)) {
                  return false;
              }
              return worker.equals(that.worker) &&
                      connectors.equals(that.connectors) &&
                      tasks.equals(that.tasks);
          }
  
          @Override
          public int hashCode() {
              return Objects.hash(worker, connectors, tasks);
          }
      }
  
  }
  