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.integration;
  
  import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
  import org.apache.kafka.connect.storage.StringConverter;
  import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
  
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Tag;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.TestInfo;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Properties;
  import java.util.concurrent.TimeUnit;
  import java.util.stream.IntStream;
  
  import static org.apache.kafka.connect.integration.TestableSourceConnector.TOPIC_CONFIG;
  import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
  import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
  import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
  import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
  import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
  import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
  import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
  import static org.apache.kafka.connect.runtime.WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_CONFIG;
  import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.COMPATIBLE;
  import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.CONNECT_PROTOCOL_CONFIG;
  import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG;
  import static org.apache.kafka.test.TestUtils.waitForCondition;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertNotEquals;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  /**
   * Integration tests for incremental cooperative rebalancing between Connect workers
   */
  @Tag("integration")
  public class RebalanceSourceConnectorsIntegrationTest {
  
      private static final Logger log = LoggerFactory.getLogger(RebalanceSourceConnectorsIntegrationTest.class);
  
      private static final int NUM_TOPIC_PARTITIONS = 3;
      private static final long CONNECTOR_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
      private static final long WORKER_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(60);
      private static final int NUM_WORKERS = 3;
      private static final int NUM_TASKS = 4;
      private static final String CONNECTOR_NAME = "seq-source1";
      private static final String TOPIC_NAME = "sequential-topic";
  
      private EmbeddedConnectCluster connect;
      
  
      @BeforeEach
      public void setup(TestInfo testInfo) {
          log.info("Starting test {}", testInfo.getDisplayName());
          // setup Connect worker properties
          Map<String, String> workerProps = new HashMap<>();
          workerProps.put(CONNECT_PROTOCOL_CONFIG, COMPATIBLE.toString());
          workerProps.put(OFFSET_COMMIT_INTERVAL_MS_CONFIG, String.valueOf(TimeUnit.SECONDS.toMillis(30)));
          workerProps.put(SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG, String.valueOf(TimeUnit.SECONDS.toMillis(30)));
  
          // setup Kafka broker properties
          Properties brokerProps = new Properties();
          brokerProps.put("auto.create.topics.enable", "false");
  
          // build a Connect cluster backed by a Kafka KRaft cluster
          connect = new EmbeddedConnectCluster.Builder()
                  .name("connect-cluster")
                  .numWorkers(NUM_WORKERS)
                  .numBrokers(1)
                  .workerProps(workerProps)
                  .brokerProps(brokerProps)
                  .build();
  
          // start the clusters
          connect.start();
      }
  
      @AfterEach
      public void close(TestInfo testInfo) {
          log.info("Finished test {}", testInfo.getDisplayName());
          // stop the Connect cluster and its backing Kafka cluster.
          connect.stop();
      }
  
      @Test
      public void testStartTwoConnectors() throws Exception {
          // create test topic
          connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
  
          // setup up props for the source connector
          Map<String, String> props = defaultSourceConnectorProps(TOPIC_NAME);
  
          // start a source connector
          connect.configureConnector(CONNECTOR_NAME, props);
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          // start a source connector
          connect.configureConnector("another-source", props);
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning("another-source", 4,
                  "Connector tasks did not start in time.");
      }
  
      @Test
      public void testReconfigConnector() throws Exception {
          ConnectorHandle connectorHandle = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME);
  
          // create test topic
          String anotherTopic = "another-topic";
          connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
          connect.kafka().createTopic(anotherTopic, NUM_TOPIC_PARTITIONS);
  
          // setup up props for the source connector
          Map<String, String> props = defaultSourceConnectorProps(TOPIC_NAME);
  
          // start a source connector
          connect.configureConnector(CONNECTOR_NAME, props);
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          int numRecordsProduced = 100;
          long recordTransferDurationMs = TimeUnit.SECONDS.toMillis(30);
  
          // consume all records from the source topic or fail, to ensure that they were correctly produced
          int recordNum = connect.kafka().consume(numRecordsProduced, recordTransferDurationMs, TOPIC_NAME).count();
          assertTrue(recordNum >= numRecordsProduced,
                  "Not enough records produced by source connector. Expected at least: " + numRecordsProduced + " + but got " + recordNum);
  
          // expect that we're going to restart the connector and its tasks
          StartAndStopLatch restartLatch = connectorHandle.expectedStarts(1);
  
          // Reconfigure the source connector by changing the Kafka topic used as output
          props.put(TOPIC_CONFIG, anotherTopic);
          connect.configureConnector(CONNECTOR_NAME, props);
  
          // Wait for the connector *and tasks* to be restarted
          assertTrue(restartLatch.await(CONNECTOR_SETUP_DURATION_MS, TimeUnit.MILLISECONDS),
                  "Failed to alter connector configuration and see connector and tasks restart "
                          + "within " + CONNECTOR_SETUP_DURATION_MS + "ms");
  
          // And wait for the Connect to show the connectors and tasks are running
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          // consume all records from the source topic or fail, to ensure that they were correctly produced
          recordNum = connect.kafka().consume(numRecordsProduced, recordTransferDurationMs, anotherTopic).count();
          assertTrue(recordNum >= numRecordsProduced,
                  "Not enough records produced by source connector. Expected at least: " + numRecordsProduced + " + but got " + recordNum);
      }
  
      @Test
      public void testDeleteConnector() throws Exception {
          // create test topic
          connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
  
          // setup up props for the source connector
          Map<String, String> props = defaultSourceConnectorProps(TOPIC_NAME);
  
          // start several source connectors
          IntStream.range(0, 4).forEachOrdered(i -> connect.configureConnector(CONNECTOR_NAME + i, props));
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          // delete connector
          connect.deleteConnector(CONNECTOR_NAME + 3);
  
          connect.assertions().assertConnectorDoesNotExist(CONNECTOR_NAME + 3,
                  "Connector wasn't deleted in time.");
  
          waitForCondition(this::assertConnectorAndTasksAreUniqueAndBalanced,
                  WORKER_SETUP_DURATION_MS, "Connect and tasks are imbalanced between the workers.");
      }
  
      @Test
      public void testAddingWorker() throws Exception {
          // create test topic
          connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
  
          // setup up props for the source connector
          Map<String, String> props = defaultSourceConnectorProps(TOPIC_NAME);
  
          // start a source connector
          IntStream.range(0, 4).forEachOrdered(i -> connect.configureConnector(CONNECTOR_NAME + i, props));
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          connect.addWorker();
  
          connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS + 1,
                  "Connect workers did not start in time.");
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          waitForCondition(this::assertConnectorAndTasksAreUniqueAndBalanced,
                  WORKER_SETUP_DURATION_MS, "Connect and tasks are imbalanced between the workers.");
      }
  
      @Test
      public void testRemovingWorker() throws Exception {
          // create test topic
          connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
  
          // setup up props for the source connector
          Map<String, String> props = defaultSourceConnectorProps(TOPIC_NAME);
  
          // start a source connector
          IntStream.range(0, 4).forEachOrdered(i -> connect.configureConnector(CONNECTOR_NAME + i, props));
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          connect.removeWorker();
  
          connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS - 1,
                  "Connect workers did not start in time.");
  
          waitForCondition(this::assertConnectorAndTasksAreUniqueAndBalanced,
                  WORKER_SETUP_DURATION_MS, "Connect and tasks are imbalanced between the workers.");
      }
  
      @Test
      public void testMultipleWorkersRejoining() throws Exception {
          // create test topic
          connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
  
          // setup up props for the source connector
          Map<String, String> props = defaultSourceConnectorProps(TOPIC_NAME);
  
          // start a source connector
          IntStream.range(0, 4).forEachOrdered(i -> connect.configureConnector(CONNECTOR_NAME + i, props));
  
          connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
                  "Connector tasks did not start in time.");
  
          waitForCondition(this::assertConnectorAndTasksAreUniqueAndBalanced,
                  WORKER_SETUP_DURATION_MS, "Connect and tasks are imbalanced between the workers.");
  
          Thread.sleep(TimeUnit.SECONDS.toMillis(10));
  
          connect.removeWorker();
          connect.removeWorker();
  
          connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS - 2,
                  "Connect workers did not stop in time.");
  
          Thread.sleep(TimeUnit.SECONDS.toMillis(10));
  
          connect.addWorker();
          connect.addWorker();
  
          connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS,
                  "Connect workers did not start in time.");
  
          Thread.sleep(TimeUnit.SECONDS.toMillis(10));
  
          for (int i = 0; i < 4; ++i) {
              connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + i, NUM_TASKS, "Connector tasks did not start in time.");
          }
  
          waitForCondition(this::assertConnectorAndTasksAreUniqueAndBalanced,
                  WORKER_SETUP_DURATION_MS, "Connect and tasks are imbalanced between the workers.");
      }
  
      private Map<String, String> defaultSourceConnectorProps(String topic) {
          // setup up props for the source connector
          Map<String, String> props = new HashMap<>();
          props.put(CONNECTOR_CLASS_CONFIG, TestableSourceConnector.class.getSimpleName());
          props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
          props.put(TOPIC_CONFIG, topic);
          props.put("throughput", String.valueOf(10));
          props.put("messages.per.poll", String.valueOf(10));
          props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
          props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
          props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
          props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
          return props;
      }
  
      private boolean assertConnectorAndTasksAreUniqueAndBalanced() {
          try {
              Map<String, Collection<String>> connectors = new HashMap<>();
              Map<String, Collection<String>> tasks = new HashMap<>();
              for (String connector : connect.connectors()) {
                  ConnectorStateInfo info = connect.connectorStatus(connector);
                  connectors.computeIfAbsent(info.connector().workerId(), k -> new ArrayList<>())
                          .add(connector);
                  info.tasks().forEach(
                      t -> tasks.computeIfAbsent(t.workerId(), k -> new ArrayList<>())
                             .add(connector + "-" + t.id()));
              }
  
              int maxConnectors = connectors.values().stream().mapToInt(Collection::size).max().orElse(0);
              int minConnectors = connectors.values().stream().mapToInt(Collection::size).min().orElse(0);
              int maxTasks = tasks.values().stream().mapToInt(Collection::size).max().orElse(0);
              int minTasks = tasks.values().stream().mapToInt(Collection::size).min().orElse(0);
  
              log.debug("Connector balance: {}", formatAssignment(connectors));
              log.debug("Task balance: {}", formatAssignment(tasks));
  
              assertNotEquals(0, maxConnectors, "Found no connectors running!");
              assertNotEquals(0, maxTasks, "Found no tasks running!");
              assertEquals(connectors.size(),
                      connectors.values().stream().distinct().count(),
                      "Connector assignments are not unique: " + connectors);
              assertEquals(tasks.size(),
                      tasks.values().stream().distinct().count(),
                      "Task assignments are not unique: " + tasks);
              assertTrue(maxConnectors - minConnectors < 2, "Connectors are imbalanced: " + formatAssignment(connectors));
              assertTrue(maxTasks - minTasks < 2, "Tasks are imbalanced: " + formatAssignment(tasks));
              return true;
          } catch (Exception e) {
              log.error("Could not check connector state info.", e);
              return false;
          }
      }
  
      private static String formatAssignment(Map<String, Collection<String>> assignment) {
          StringBuilder result = new StringBuilder();
          for (String worker : assignment.keySet().stream().sorted().toList()) {
              result.append(String.format("\n%s=%s", worker, assignment.getOrDefault(worker,
                      Collections.emptyList())));
          }
          return result.toString();
      }
  
  }
  