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.common.protocol.types.ArrayOf;
  import org.apache.kafka.common.protocol.types.Field;
  import org.apache.kafka.common.protocol.types.Schema;
  import org.apache.kafka.common.protocol.types.SchemaException;
  import org.apache.kafka.common.protocol.types.Struct;
  import org.apache.kafka.common.protocol.types.Type;
  import org.apache.kafka.connect.util.ConnectorTaskId;
  
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.LinkedHashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  
  import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocol;
  import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection;
  import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.EAGER;
  
  /**
   * This class implements the protocol for Kafka Connect workers in a group. It includes the format of worker state used when
   * joining the group and distributing assignments, and the format of assignments of connectors and tasks to workers.
   */
  public class ConnectProtocol {
      public static final String VERSION_KEY_NAME = "version";
      public static final String URL_KEY_NAME = "url";
      public static final String CONFIG_OFFSET_KEY_NAME = "config-offset";
      public static final String CONNECTOR_KEY_NAME = "connector";
      public static final String LEADER_KEY_NAME = "leader";
      public static final String LEADER_URL_KEY_NAME = "leader-url";
      public static final String ERROR_KEY_NAME = "error";
      public static final String TASKS_KEY_NAME = "tasks";
      public static final String ASSIGNMENT_KEY_NAME = "assignment";
      public static final int CONNECTOR_TASK = -1;
  
      public static final short CONNECT_PROTOCOL_V0 = 0;
      public static final Schema CONNECT_PROTOCOL_HEADER_SCHEMA = new Schema(
              new Field(VERSION_KEY_NAME, Type.INT16));
  
      /**
       * Connect Protocol Header V0:
       * <pre>
       *   Version            => Int16
       * </pre>
       */
      private static final Struct CONNECT_PROTOCOL_HEADER_V0 = new Struct(CONNECT_PROTOCOL_HEADER_SCHEMA)
              .set(VERSION_KEY_NAME, CONNECT_PROTOCOL_V0);
  
      /**
       * Config State V0:
       * <pre>
       *   Url                => [String]
       *   ConfigOffset       => Int64
       * </pre>
       */
      public static final Schema CONFIG_STATE_V0 = new Schema(
              new Field(URL_KEY_NAME, Type.STRING),
              new Field(CONFIG_OFFSET_KEY_NAME, Type.INT64));
  
      /**
       * Connector Assignment V0:
       * <pre>
       *   Connector          => [String]
       *   Tasks              => [Int32]
       * </pre>
       *
       * <p>Assignments for each worker are a set of connectors and tasks. These are categorized by
       * connector ID. A sentinel task ID (CONNECTOR_TASK) is used to indicate the connector itself
       * (i.e. that the assignment includes responsibility for running the Connector instance in
       * addition to any tasks it generates).</p>
       */
      public static final Schema CONNECTOR_ASSIGNMENT_V0 = new Schema(
              new Field(CONNECTOR_KEY_NAME, Type.STRING),
              new Field(TASKS_KEY_NAME, new ArrayOf(Type.INT32)));
  
      /**
       * Assignment V0:
       * <pre>
       *   Error              => Int16
       *   Leader             => [String]
       *   LeaderUrl          => [String]
       *   ConfigOffset       => Int64
       *   Assignment         => [Connector Assignment]
       * </pre>
       */
      public static final Schema ASSIGNMENT_V0 = new Schema(
              new Field(ERROR_KEY_NAME, Type.INT16),
              new Field(LEADER_KEY_NAME, Type.STRING),
              new Field(LEADER_URL_KEY_NAME, Type.STRING),
              new Field(CONFIG_OFFSET_KEY_NAME, Type.INT64),
              new Field(ASSIGNMENT_KEY_NAME, new ArrayOf(CONNECTOR_ASSIGNMENT_V0)));
  
      /**
       * The fields are serialized in sequence as follows:
       * Subscription V0:
       * <pre>
       *   Version            => Int16
       *   Url                => [String]
       *   ConfigOffset       => Int64
       * </pre>
       *
       * @param workerState the current state of the worker metadata
       * @return the serialized state of the worker metadata
       */
      public static ByteBuffer serializeMetadata(WorkerState workerState) {
          Struct struct = new Struct(CONFIG_STATE_V0);
          struct.set(URL_KEY_NAME, workerState.url());
          struct.set(CONFIG_OFFSET_KEY_NAME, workerState.offset());
          ByteBuffer buffer = ByteBuffer.allocate(CONNECT_PROTOCOL_HEADER_V0.sizeOf() + CONFIG_STATE_V0.sizeOf(struct));
          CONNECT_PROTOCOL_HEADER_V0.writeTo(buffer);
          CONFIG_STATE_V0.write(buffer, struct);
          buffer.flip();
          return buffer;
      }
  
      /**
       * Returns the collection of Connect protocols that are supported by this version along
       * with their serialized metadata. The protocols are ordered by preference.
       *
       * @param workerState the current state of the worker metadata
       * @return the collection of Connect protocol metadata
       */
      public static JoinGroupRequestProtocolCollection metadataRequest(WorkerState workerState) {
          return new JoinGroupRequestProtocolCollection(Collections.singleton(
                  new JoinGroupRequestProtocol()
                          .setName(EAGER.protocol())
                          .setMetadata(ConnectProtocol.serializeMetadata(workerState).array()))
                  .iterator());
      }
  
      /**
       * Given a byte buffer that contains protocol metadata return the deserialized form of the
       * metadata.
       *
       * @param buffer A buffer containing the protocols metadata
       * @return the deserialized metadata
       * @throws SchemaException on incompatible Connect protocol version
       */
      public static WorkerState deserializeMetadata(ByteBuffer buffer) {
          Struct header = CONNECT_PROTOCOL_HEADER_SCHEMA.read(buffer);
          Short version = header.getShort(VERSION_KEY_NAME);
          checkVersionCompatibility(version);
          Struct struct = CONFIG_STATE_V0.read(buffer);
          long configOffset = struct.getLong(CONFIG_OFFSET_KEY_NAME);
          String url = struct.getString(URL_KEY_NAME);
          return new WorkerState(url, configOffset);
      }
  
      /**
       * The fields are serialized in sequence as follows:
       * Complete Assignment V0:
       * <pre>
       *   Version            => Int16
       *   Error              => Int16
       *   Leader             => [String]
       *   LeaderUrl          => [String]
       *   ConfigOffset       => Int64
       *   Assignment         => [Connector Assignment]
       * </pre>
       */
      public static ByteBuffer serializeAssignment(Assignment assignment) {
          Struct struct = new Struct(ASSIGNMENT_V0);
          struct.set(ERROR_KEY_NAME, assignment.error());
          struct.set(LEADER_KEY_NAME, assignment.leader());
          struct.set(LEADER_URL_KEY_NAME, assignment.leaderUrl());
          struct.set(CONFIG_OFFSET_KEY_NAME, assignment.offset());
          List<Struct> taskAssignments = new ArrayList<>();
          for (Map.Entry<String, Collection<Integer>> connectorEntry : assignment.asMap().entrySet()) {
              Struct taskAssignment = new Struct(CONNECTOR_ASSIGNMENT_V0);
              taskAssignment.set(CONNECTOR_KEY_NAME, connectorEntry.getKey());
              Collection<Integer> tasks = connectorEntry.getValue();
              taskAssignment.set(TASKS_KEY_NAME, tasks.toArray());
              taskAssignments.add(taskAssignment);
          }
          struct.set(ASSIGNMENT_KEY_NAME, taskAssignments.toArray());
  
          ByteBuffer buffer = ByteBuffer.allocate(CONNECT_PROTOCOL_HEADER_V0.sizeOf() + ASSIGNMENT_V0.sizeOf(struct));
          CONNECT_PROTOCOL_HEADER_V0.writeTo(buffer);
          ASSIGNMENT_V0.write(buffer, struct);
          buffer.flip();
          return buffer;
      }
  
      /**
       * Given a byte buffer that contains an assignment as defined by this protocol, return the
       * deserialized form of the assignment.
       *
       * @param buffer the buffer containing a serialized assignment
       * @return the deserialized assignment
       * @throws SchemaException on incompatible Connect protocol version
       */
      public static Assignment deserializeAssignment(ByteBuffer buffer) {
          Struct header = CONNECT_PROTOCOL_HEADER_SCHEMA.read(buffer);
          Short version = header.getShort(VERSION_KEY_NAME);
          checkVersionCompatibility(version);
          Struct struct = ASSIGNMENT_V0.read(buffer);
          short error = struct.getShort(ERROR_KEY_NAME);
          String leader = struct.getString(LEADER_KEY_NAME);
          String leaderUrl = struct.getString(LEADER_URL_KEY_NAME);
          long offset = struct.getLong(CONFIG_OFFSET_KEY_NAME);
          List<String> connectorIds = new ArrayList<>();
          List<ConnectorTaskId> taskIds = new ArrayList<>();
          for (Object structObj : struct.getArray(ASSIGNMENT_KEY_NAME)) {
              Struct assignment = (Struct) structObj;
              String connector = assignment.getString(CONNECTOR_KEY_NAME);
              for (Object taskIdObj : assignment.getArray(TASKS_KEY_NAME)) {
                  Integer taskId = (Integer) taskIdObj;
                  if (taskId == CONNECTOR_TASK)
                      connectorIds.add(connector);
                  else
                      taskIds.add(new ConnectorTaskId(connector, taskId));
              }
          }
          return new Assignment(error, leader, leaderUrl, offset, connectorIds, taskIds);
      }
  
      /**
       * A class that captures the deserialized form of a worker's metadata.
       */
      public static class WorkerState {
          private final String url;
          private final long offset;
  
          public WorkerState(String url, long offset) {
              this.url = url;
              this.offset = offset;
          }
  
          public String url() {
              return url;
          }
  
          /**
           * The most up-to-date (maximum) configuration offset known to this worker.
           *
           * @return the configuration offset
           */
          public long offset() {
              return offset;
          }
  
          @Override
          public String toString() {
              return "WorkerState{" +
                      "url='" + url + '\'' +
                      ", offset=" + offset +
                      '}';
          }
      }
  
      /**
       * The basic assignment of connectors and tasks introduced with V0 version of the Connect protocol.
       */
      public static class Assignment {
          public static final short NO_ERROR = 0;
          // Configuration offsets mismatched in a way that the leader could not resolve. Workers should read to the end
          // of the config log and try to re-join
          public static final short CONFIG_MISMATCH = 1;
  
          private final short error;
          private final String leader;
          private final String leaderUrl;
          private final long offset;
          private final Collection<String> connectorIds;
          private final Collection<ConnectorTaskId> taskIds;
  
          /**
           * Create an assignment indicating responsibility for the given connector instances and task Ids.
           *
           * @param error error code for this assignment; {@link ConnectProtocol.Assignment#NO_ERROR}
           *              indicates no error during assignment
           * @param leader Connect group's leader Id; may be null only on the empty assignment
           * @param leaderUrl Connect group's leader URL; may be null only on the empty assignment
           * @param configOffset the most up-to-date configuration offset according to this assignment
           * @param connectorIds list of connectors that the worker should instantiate and run; may not be null
           * @param taskIds list of task IDs that the worker should instantiate and run; may not be null
           */
          public Assignment(short error, String leader, String leaderUrl, long configOffset,
                            Collection<String> connectorIds, Collection<ConnectorTaskId> taskIds) {
              this.error = error;
              this.leader = leader;
              this.leaderUrl = leaderUrl;
              this.offset = configOffset;
              this.connectorIds = Objects.requireNonNull(connectorIds,
                      "Assigned connector IDs may be empty but not null");
              this.taskIds = Objects.requireNonNull(taskIds,
                      "Assigned task IDs may be empty but not null");
          }
  
          /**
           * Return the error code of this assignment; 0 signals successful assignment ({@link ConnectProtocol.Assignment#NO_ERROR}).
           *
           * @return the error code of the assignment
           */
          public short error() {
              return error;
          }
  
          /**
           * Return the ID of the leader Connect worker in this assignment.
           *
           * @return the ID of the leader
           */
          public String leader() {
              return leader;
          }
  
          /**
           * Return the URL to which the leader accepts requests from other members of the group.
           *
           * @return the leader URL
           */
          public String leaderUrl() {
              return leaderUrl;
          }
  
          /**
           * Check if this assignment failed.
           *
           * @return true if this assignment failed; false otherwise
           */
          public boolean failed() {
              return error != NO_ERROR;
          }
  
          /**
           * Return the most up-to-date offset in the configuration topic according to this assignment
           *
           * @return the up-to-date configuration topic offset
           */
          public long offset() {
              return offset;
          }
  
          /**
           * The connectors included in this assignment.
           *
           * @return the connectors
           */
          public Collection<String> connectors() {
              return connectorIds;
          }
  
          /**
           * The tasks included in this assignment.
           *
           * @return the tasks
           */
          public Collection<ConnectorTaskId> tasks() {
              return taskIds;
          }
  
          @Override
          public String toString() {
              return "Assignment{" +
                      "error=" + error +
                      ", leader='" + leader + '\'' +
                      ", leaderUrl='" + leaderUrl + '\'' +
                      ", offset=" + offset +
                      ", connectorIds=" + connectorIds +
                      ", taskIds=" + taskIds +
                      '}';
          }
  
          protected Map<String, Collection<Integer>> asMap() {
              // Using LinkedHashMap preserves the ordering, which is helpful for tests and debugging
              Map<String, Collection<Integer>> taskMap = new LinkedHashMap<>();
              for (String connectorId : new HashSet<>(connectorIds)) {
                  taskMap.computeIfAbsent(connectorId, key -> new ArrayList<>()).add(CONNECTOR_TASK);
              }
              for (ConnectorTaskId taskId : taskIds) {
                  String connectorId = taskId.connector();
                  taskMap.computeIfAbsent(connectorId, key -> new ArrayList<>()).add(taskId.task());
              }
              return taskMap;
          }
      }
  
      private static void checkVersionCompatibility(short version) {
          // check for invalid versions
          if (version < CONNECT_PROTOCOL_V0)
              throw new SchemaException("Unsupported subscription version: " + version);
  
          // otherwise, assume versions can be parsed as V0
      }
  }
  