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.util.clusters;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.Iterator;
  import java.util.LinkedHashSet;
  import java.util.Map;
  import java.util.Properties;
  import java.util.Set;
  import java.util.UUID;
  import java.util.concurrent.atomic.AtomicInteger;
  
  import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
  import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
  import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
  import static org.apache.kafka.connect.runtime.WorkerConfig.BOOTSTRAP_SERVERS_CONFIG;
  import static org.apache.kafka.connect.runtime.WorkerConfig.PLUGIN_DISCOVERY_CONFIG;
  import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG;
  import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.CONFIG_TOPIC_CONFIG;
  import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG;
  import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG;
  import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG;
  import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG;
  import static org.apache.kafka.connect.runtime.rest.RestServerConfig.LISTENERS_CONFIG;
  
  /**
   * Start an embedded Connect cluster that can be used for integration tests. Internally, this class also spins up a
   * backing Kafka KRaft cluster for the Connect cluster leveraging {@link org.apache.kafka.common.test.KafkaClusterTestKit}. Methods
   * on the same {@code EmbeddedConnectCluster} are not guaranteed to be thread-safe. This class also provides various
   * utility methods to perform actions on the Connect cluster such as connector creation, config validation, connector
   * restarts, pause / resume, connector deletion etc.
   */
  public class EmbeddedConnectCluster extends EmbeddedConnect {
  
      private static final Logger log = LoggerFactory.getLogger(EmbeddedConnectCluster.class);
  
      public static final int DEFAULT_NUM_WORKERS = 1;
      private static final String REST_HOST_NAME = "localhost";
  
      private static final String DEFAULT_WORKER_NAME_PREFIX = "connect-worker-";
  
      private final Set<WorkerHandle> connectCluster;
      private final Map<String, String> workerProps;
      private final String connectClusterName;
      private final int numInitialWorkers;
      private final String workerNamePrefix;
      private final AtomicInteger nextWorkerId = new AtomicInteger(0);
  
      private EmbeddedConnectCluster(
              int numBrokers,
              Properties brokerProps,
              boolean maskExitProcedures,
              Map<String, String> clientProps,
              Map<String, String> workerProps,
              String name,
              int numWorkers
      ) {
          super(numBrokers, brokerProps, maskExitProcedures, clientProps);
          this.workerProps = workerProps;
          this.connectClusterName = name;
          this.connectCluster = new LinkedHashSet<>();
          this.numInitialWorkers = numWorkers;
          // leaving non-configurable for now
          this.workerNamePrefix = DEFAULT_WORKER_NAME_PREFIX;
      }
  
      /**
       * Provision and start an additional worker to the Connect cluster.
       *
       * @return the worker handle of the worker that was provisioned
       */
      public WorkerHandle addWorker() {
          WorkerHandle worker = WorkerHandle.start(workerNamePrefix + nextWorkerId.getAndIncrement(), workerProps);
          connectCluster.add(worker);
          log.info("Started worker {}", worker);
          return worker;
      }
  
      /**
       * Decommission one of the workers from this Connect cluster. Which worker is removed is
       * implementation dependent and selection is not guaranteed to be consistent. Use this method
       * when you don't care which worker stops.
       *
       * @see #removeWorker(WorkerHandle)
       */
      public void removeWorker() {
          WorkerHandle toRemove = null;
          for (Iterator<WorkerHandle> it = connectCluster.iterator(); it.hasNext(); toRemove = it.next()) {
          }
          if (toRemove != null) {
              removeWorker(toRemove);
          }
      }
  
      /**
       * Decommission a specific worker from this Connect cluster.
       *
       * @param worker the handle of the worker to remove from the cluster
       * @throws IllegalStateException if the Connect cluster has no workers
       */
      public void removeWorker(WorkerHandle worker) {
          if (connectCluster.isEmpty()) {
              throw new IllegalStateException("Cannot remove worker. Cluster is empty");
          }
          stopWorker(worker);
          connectCluster.remove(worker);
      }
  
      /**
       * Determine whether the Connect cluster has any workers running.
       *
       * @return true if any worker is running, or false otherwise
       */
      public boolean anyWorkersHealthy() {
          return workers().stream().anyMatch(this::isHealthy);
      }
  
      /**
       * Determine whether the Connect cluster has all workers running.
       *
       * @return true if all workers are running, or false otherwise
       */
      public boolean allWorkersHealthy() {
          return workers().stream().allMatch(this::isHealthy);
      }
  
      @Override
      public void startConnect() {
          log.info("Starting Connect cluster '{}' with {} workers", connectClusterName, numInitialWorkers);
  
          workerProps.put(BOOTSTRAP_SERVERS_CONFIG, kafka().bootstrapServers());
          // use a random available port
          workerProps.put(LISTENERS_CONFIG, "HTTP://" + REST_HOST_NAME + ":0");
  
          String internalTopicsReplFactor = String.valueOf(numBrokers);
          workerProps.putIfAbsent(GROUP_ID_CONFIG, "connect-integration-test-" + connectClusterName);
          workerProps.putIfAbsent(OFFSET_STORAGE_TOPIC_CONFIG, "connect-offset-topic-" + connectClusterName);
          workerProps.putIfAbsent(OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, internalTopicsReplFactor);
          workerProps.putIfAbsent(CONFIG_TOPIC_CONFIG, "connect-config-topic-" + connectClusterName);
          workerProps.putIfAbsent(CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, internalTopicsReplFactor);
          workerProps.putIfAbsent(STATUS_STORAGE_TOPIC_CONFIG, "connect-status-topic-" + connectClusterName);
          workerProps.putIfAbsent(STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, internalTopicsReplFactor);
          workerProps.putIfAbsent(KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.storage.StringConverter");
          workerProps.putIfAbsent(VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.storage.StringConverter");
          workerProps.putIfAbsent(PLUGIN_DISCOVERY_CONFIG, "hybrid_fail");
  
          for (int i = 0; i < numInitialWorkers; i++) {
              addWorker();
          }
      }
  
      @Override
      public String toString() {
          return String.format("EmbeddedConnectCluster(name= %s, numBrokers= %d, numInitialWorkers= %d, workerProps= %s)",
              connectClusterName,
              numBrokers,
              numInitialWorkers,
              workerProps);
      }
  
      public String getName() {
          return connectClusterName;
      }
  
      /**
       * Get the provisioned workers.
       *
       * @return the set of handles of the provisioned workers
       */
      public Set<WorkerHandle> workers() {
          return new LinkedHashSet<>(connectCluster);
      }
  
      public static class Builder extends EmbeddedConnectBuilder<EmbeddedConnectCluster, Builder> {
          private String name = UUID.randomUUID().toString();
          private int numWorkers = DEFAULT_NUM_WORKERS;
  
          public Builder name(String name) {
              this.name = name;
              return this;
          }
  
          public Builder numWorkers(int numWorkers) {
              this.numWorkers = numWorkers;
              return this;
          }
  
          @Override
          protected EmbeddedConnectCluster build(
                  int numBrokers,
                  Properties brokerProps,
                  boolean maskExitProcedures,
                  Map<String, String> clientProps,
                  Map<String, String> workerProps
          ) {
              return new EmbeddedConnectCluster(
                      numBrokers,
                      brokerProps,
                      maskExitProcedures,
                      clientProps,
                      workerProps,
                      name,
                      numWorkers
              );
          }
      }
  
  }
  