Migrate Executor off private Kafka APIs

Signed-off-by: Kyle Liberti <kliberti.us@gmail.com>

Author: kyguy

cruise-control-metrics-reporter/src/test/java/com/linkedin/kafka/cruisecontrol/metricsreporter/utils/CCContainerizedKraftCluster.java

@@ -396,25 +396,6 @@ public void shutDownBroker(int brokerId) {
       timeout,
       String.format("Broker %s did not shutdown properly.", brokerId)
     );
-
-    // Wait until describeLogDirs fails
-    waitUntil(
-      () -> {
-        try {
-          _adminClient.describeLogDirs(Collections.singletonList(brokerId))
-            .allDescriptions()
-            .get(5, TimeUnit.SECONDS);
-          return false;
-        } catch (InterruptedException ie) {
-          throw new RuntimeException(ie);
-        } catch (Exception e) {
-          return true;
-        }
-      },
-      result -> result,
-      timeout,
-      String.format("Broker %s did not fully shut down (logDirs RPC still succeeds).", brokerId)
-    );
   }
 
   public static class BrokerWaitStrategy extends AbstractWaitStrategy {

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/common/MetadataAdminClient.java

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2017 LinkedIn Corp. Licensed under the BSD 2-Clause License (the "License"). See License in the project root for license information.
+ */
+
+package com.linkedin.kafka.cruisecontrol.common;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.DescribeClusterResult;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartitionInfo;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Client for fetching Kafka cluster metadata using Kafka Admin APIs.
+ *
+ * This replaces the use of MetadataClient which relies on internal Kafka APIs.
+ */
+public class MetadataAdminClient {
+  private static final Logger LOG = LoggerFactory.getLogger(MetadataAdminClient.class);
+
+  private final Admin _adminClient;
+
+  /**
+   * Creates a new MetadataAdminClient.
+   *
+   * @param adminClient The AdminClient to use for fetching cluster metadata.
+   */
+  public MetadataAdminClient(Admin adminClient) {
+    _adminClient = adminClient;
+  }
+
+  /**
+   * Close adminClient
+   */
+  public void close() {
+    if (_adminClient != null) {
+      try {
+        _adminClient.close();
+      } catch (Exception e) {
+        LOG.warn("Failed to close AdminClient", e);
+      }
+    }
+  }
+
+  /**
+   * Get cluster metadata.
+   *
+   * @return Cluster containing the cluster metadata.
+   */
+  public Cluster cluster() {
+    try {
+      DescribeClusterResult describeResult = _adminClient.describeCluster();
+      Collection<Node> nodes = describeResult.nodes().get();
+      String clusterId = describeResult.clusterId().get();
+
+      Set<String> topicNames = _adminClient.listTopics()
+        .names().get();
+
+      Map<String, TopicDescription> topicDescriptions = _adminClient.describeTopics(topicNames)
+        .allTopicNames().get();
+
+      List<PartitionInfo> partitionInfos = new ArrayList<>();
+      for (TopicDescription desc : topicDescriptions.values()) {
+        for (TopicPartitionInfo partInfo : desc.partitions()) {
+          Node leader = partInfo.leader();
+          Node[] replicas = partInfo.replicas().toArray(Node[]::new);
+          Node[] isr = partInfo.isr().toArray(Node[]::new);
+          partitionInfos.add(new PartitionInfo(desc.name(), partInfo.partition(), leader, replicas, isr));
+        }
+      }
+
+      return new Cluster(clusterId, nodes, partitionInfos, Collections.emptySet(), Collections.emptySet());
+    } catch (ExecutionException | InterruptedException e) {
+      LOG.error("Failed to fetch cluster metadata", e);
+      throw new RuntimeException("Failed to fetch cluster metadata", e);
+    }
+  }
+}

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/common/MetadataClient.java

@@ -26,7 +26,13 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-
+/**
+ * Legacy client for fetching Kafka cluster metadata.
+ *
+ * @deprecated This class is deprecated and will be removed in a future release in favor of {@link MetadataAdminClient}
+ * which uses public Kafka APIs.
+ */
+@Deprecated
 public class MetadataClient {
   private static final Logger LOG = LoggerFactory.getLogger(MetadataClient.class);
   private static final LogContext LOG_CONTEXT = new LogContext();

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionUtils.java

@@ -29,6 +29,7 @@
 import java.util.concurrent.TimeoutException;
 import javax.annotation.Nullable;
 import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.admin.AlterPartitionReassignmentsResult;
 import org.apache.kafka.clients.admin.DescribeReplicaLogDirsResult;
@@ -347,7 +348,7 @@ public static void populateMinIsrState(Cluster cluster,
    * @param adminClient The adminClient to ask for ongoing partition reassignments.
    * @return The set of {@link TopicPartition partitions} that are being reassigned.
    */
-  public static Set<TopicPartition> partitionsBeingReassigned(AdminClient adminClient)
+  public static Set<TopicPartition> partitionsBeingReassigned(Admin adminClient)
       throws InterruptedException, ExecutionException, TimeoutException {
     return ongoingPartitionReassignments(adminClient).keySet();
   }
@@ -362,7 +363,7 @@ public static Set<TopicPartition> partitionsBeingReassigned(AdminClient adminCli
    * @param adminClient The adminClient to ask for ongoing partition reassignments.
    * @return The map of {@link PartitionReassignment reassignment} by {@link TopicPartition partitions}.
    */
-  public static Map<TopicPartition, PartitionReassignment> ongoingPartitionReassignments(AdminClient adminClient)
+  public static Map<TopicPartition, PartitionReassignment> ongoingPartitionReassignments(Admin adminClient)
       throws InterruptedException, ExecutionException, TimeoutException {
     Map<TopicPartition, PartitionReassignment> partitionReassignments = null;
     int attempts = 0;
@@ -404,7 +405,7 @@ private static Optional<NewPartitionReassignment> reassignmentValue(List<Integer
    * @param tasks Preferred leader election tasks to execute.
    * @return The {@link ElectLeadersResult result} of preferred leader election request -- cannot be {@code null}.
    */
-  public static ElectLeadersResult submitPreferredLeaderElection(AdminClient adminClient, List<ExecutionTask> tasks) {
+  public static ElectLeadersResult submitPreferredLeaderElection(Admin adminClient, List<ExecutionTask> tasks) {
     if (validateNotNull(tasks, "Tasks to execute cannot be null.").isEmpty()) {
       throw new IllegalArgumentException("Tasks to execute cannot be empty.");
     }
@@ -442,7 +443,7 @@ public static ElectLeadersResult submitPreferredLeaderElection(AdminClient admin
    * @param tasks Inter-broker replica reassignment tasks to execute.
    * @return The {@link AlterPartitionReassignmentsResult result} of reassignment request -- cannot be {@code null}.
    */
-  public static AlterPartitionReassignmentsResult submitReplicaReassignmentTasks(AdminClient adminClient, List<ExecutionTask> tasks) {
+  public static AlterPartitionReassignmentsResult submitReplicaReassignmentTasks(Admin adminClient, List<ExecutionTask> tasks) {
     if (validateNotNull(tasks, "Tasks to execute cannot be null.").isEmpty()) {
       throw new IllegalArgumentException("Tasks to execute cannot be empty.");
     }

cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/Executor.java

@@ -9,10 +9,10 @@
 import com.codahale.metrics.Timer;
 import com.google.common.util.concurrent.AtomicDouble;
 import com.linkedin.kafka.cruisecontrol.KafkaCruiseControlUtils;
+import com.linkedin.kafka.cruisecontrol.common.MetadataAdminClient;
 import com.linkedin.kafka.cruisecontrol.common.TopicMinIsrCache;
 import com.linkedin.kafka.cruisecontrol.config.KafkaCruiseControlConfig;
 import com.linkedin.kafka.cruisecontrol.common.KafkaCruiseControlThreadFactory;
-import com.linkedin.kafka.cruisecontrol.common.MetadataClient;
 import com.linkedin.kafka.cruisecontrol.config.constants.ExecutorConfig;
 import com.linkedin.kafka.cruisecontrol.detector.AnomalyDetectorManager;
 import com.linkedin.kafka.cruisecontrol.exception.OngoingExecutionException;
@@ -42,7 +42,6 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
-import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.admin.AlterPartitionReassignmentsResult;
 import org.apache.kafka.clients.admin.DescribeConfigsResult;
@@ -52,8 +51,6 @@
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.config.ConfigResource;
-import org.apache.kafka.common.internals.ClusterResourceListeners;
-import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -87,7 +84,7 @@ public class Executor {
   private static final long EXECUTION_PROGRESS_CHECK_INTERVAL_ADJUSTING_MS = 1000;
   // The execution progress is controlled by the ExecutionTaskManager.
   private final ExecutionTaskManager _executionTaskManager;
-  private final MetadataClient _metadataClient;
+  private final MetadataAdminClient _metadataClient;
   private volatile long _executionProgressCheckIntervalMs;
   private final long _defaultExecutionProgressCheckIntervalMs;
   private Long _requestedExecutionProgressCheckIntervalMs;
@@ -168,7 +165,7 @@ public Executor(KafkaCruiseControlConfig config,
   Executor(KafkaCruiseControlConfig config,
            Time time,
            MetricRegistry dropwizardMetricRegistry,
-           MetadataClient metadataClient,
+           MetadataAdminClient metadataClient,
            ExecutorNotifier executorNotifier,
            AnomalyDetectorManager anomalyDetectorManager) {
     _numExecutionStopped = new AtomicInteger(0);
@@ -189,15 +186,7 @@ public Executor(KafkaCruiseControlConfig config,
     _executionTaskManager = new ExecutionTaskManager(_adminClient, dropwizardMetricRegistry, time, config);
     // Register gauge sensors.
     registerGaugeSensors(dropwizardMetricRegistry);
-    _metadataClient = metadataClient != null ? metadataClient
-                                             : new MetadataClient(config,
-                                                                  new Metadata(ExecutionUtils.METADATA_REFRESH_BACKOFF,
-                                                                               ExecutionUtils.METADATA_REFRESH_BACKOFF_MAX,
-                                                                               ExecutionUtils.METADATA_EXPIRY_MS,
-                                                                               new LogContext(),
-                                                                               new ClusterResourceListeners()),
-                                                                  -1L,
-                                                                  time);
+    _metadataClient = metadataClient != null ? metadataClient : new MetadataAdminClient(_adminClient);
     _defaultExecutionProgressCheckIntervalMs = config.getLong(ExecutorConfig.EXECUTION_PROGRESS_CHECK_INTERVAL_MS_CONFIG);
     _executionProgressCheckIntervalMs = _defaultExecutionProgressCheckIntervalMs;
     _leaderMovementTimeoutMs = config.getLong(ExecutorConfig.LEADER_MOVEMENT_TIMEOUT_MS_CONFIG);
@@ -875,7 +864,7 @@ private synchronized void initProposalExecution(Collection<ExecutionProposal> pr
                                                                recentlyRemovedBrokers(), isTriggeredByUserRequest);
     _executionTaskManager.setExecutionModeForTaskTracker(_isKafkaAssignerMode);
     // Get a snapshot of (1) cluster and (2) minIsr with time by topic name.
-    StrategyOptions strategyOptions = new StrategyOptions.Builder(_metadataClient.refreshMetadata().cluster())
+    StrategyOptions strategyOptions = new StrategyOptions.Builder(_metadataClient.cluster())
         .minIsrWithTimeByTopic(_topicMinIsrCache.minIsrWithTimeByTopic()).build();
     _executionTaskManager.addExecutionProposals(proposals, brokersToSkipConcurrencyCheck, strategyOptions, replicaMovementStrategy);
     _concurrencyAdjuster.initAdjustment(loadMonitor,
@@ -1799,7 +1788,7 @@ private Cluster getClusterForExecutionProgressCheck() {
       if (LOG.isDebugEnabled()) {
         LOG.debug("Tasks in execution: {}", inExecutionTasks());
       }
-      return _metadataClient.refreshMetadata().cluster();
+      return _metadataClient.cluster();
     }
 
     /**

cruise-control/src/test/java/com/linkedin/kafka/cruisecontrol/executor/ExecutorTest.java

@@ -6,7 +6,7 @@
 
 import com.codahale.metrics.MetricRegistry;
 import com.linkedin.kafka.cruisecontrol.KafkaCruiseControlUtils;
-import com.linkedin.kafka.cruisecontrol.common.MetadataClient;
+import com.linkedin.kafka.cruisecontrol.common.MetadataAdminClient;
 import com.linkedin.kafka.cruisecontrol.common.TestConstants;
 import com.linkedin.kafka.cruisecontrol.config.BrokerCapacityConfigFileResolver;
 import com.linkedin.kafka.cruisecontrol.config.KafkaCruiseControlConfig;
@@ -103,6 +103,8 @@ public int clusterSize() {
   @Before
   public void setUp() {
     Properties adminClientProps = new Properties();
+    adminClientProps.setProperty(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "5000");
+    adminClientProps.setProperty(AdminClientConfig.METADATA_MAX_AGE_CONFIG, "0");
     setSecurityConfigs(adminClientProps, "admin");
 
     _cluster = new CCContainerizedKraftCluster(clusterSize(), buildBrokerConfigs(), adminClientProps);
@@ -405,7 +407,7 @@ private static boolean verifyFutureError(Future<?> future, Class<? extends Throw
   @Test
   public void testSetRequestedExecutionProgressCheckIntervalMs() {
     KafkaCruiseControlConfig config = new KafkaCruiseControlConfig(getExecutorProperties());
-    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataClient.class),
+    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataAdminClient.class),
                                      null, EasyMock.mock(AnomalyDetectorManager.class));
     long minExecutionProgressCheckIntervalMs = config.getLong(ExecutorConfig.MIN_EXECUTION_PROGRESS_CHECK_INTERVAL_MS_CONFIG);
 
@@ -419,7 +421,7 @@ public void testSetRequestedExecutionProgressCheckIntervalMs() {
   @Test
   public void testSetExecutionProgressCheckIntervalMsWithRequestedValue() {
     KafkaCruiseControlConfig config = new KafkaCruiseControlConfig(getExecutorProperties());
-    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataClient.class),
+    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataAdminClient.class),
                                      null, EasyMock.mock(AnomalyDetectorManager.class));
     long defaultExecutionProgressCheckIntervalMs = config.getLong(ExecutorConfig.EXECUTION_PROGRESS_CHECK_INTERVAL_MS_CONFIG);
     long minExecutionProgressCheckIntervalMs = config.getLong(ExecutorConfig.MIN_EXECUTION_PROGRESS_CHECK_INTERVAL_MS_CONFIG);
@@ -440,7 +442,7 @@ public void testSetExecutionProgressCheckIntervalMsWithRequestedValue() {
   @Test
   public void testSetExecutionProgressCheckIntervalMsWithNoRequestedValue() {
     KafkaCruiseControlConfig config = new KafkaCruiseControlConfig(getExecutorProperties());
-    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataClient.class),
+    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataAdminClient.class),
                                      null, EasyMock.mock(AnomalyDetectorManager.class));
     long defaultExecutionProgressCheckIntervalMs = config.getLong(ExecutorConfig.EXECUTION_PROGRESS_CHECK_INTERVAL_MS_CONFIG);
     long minExecutionProgressCheckIntervalMs = config.getLong(ExecutorConfig.MIN_EXECUTION_PROGRESS_CHECK_INTERVAL_MS_CONFIG);
@@ -461,7 +463,7 @@ public void testSetExecutionProgressCheckIntervalMsWithNoRequestedValue() {
   @Test
   public void testResetExecutionProgressCheckIntervalMs() {
     KafkaCruiseControlConfig config = new KafkaCruiseControlConfig(getExecutorProperties());
-    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataClient.class),
+    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataAdminClient.class),
         null, EasyMock.mock(AnomalyDetectorManager.class));
     long defaultExecutionProgressCheckIntervalMs = config.getLong(ExecutorConfig.EXECUTION_PROGRESS_CHECK_INTERVAL_MS_CONFIG);
     executor.resetExecutionProgressCheckIntervalMs();
@@ -477,8 +479,8 @@ public void testResetExecutionProgressCheckIntervalMs() {
   public void testExecutionKnobs() {
     KafkaCruiseControlConfig config = new KafkaCruiseControlConfig(getExecutorProperties());
     assertThrows(IllegalStateException.class,
-                 () -> new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataClient.class), null, null));
-    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataClient.class),
+                 () -> new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataAdminClient.class), null, null));
+    Executor executor = new Executor(config, null, new MetricRegistry(), EasyMock.mock(MetadataAdminClient.class),
                                      null, EasyMock.mock(AnomalyDetectorManager.class));
 
     // Verify correctness of add/drop recently removed/demoted brokers.
@@ -502,7 +504,7 @@ public void testTimeoutAndExecutionStop() throws InterruptedException, OngoingEx
 
     KafkaCruiseControlConfig configs = new KafkaCruiseControlConfig(getExecutorProperties());
     Time time = new MockTime();
-    MetadataClient mockMetadataClient = EasyMock.mock(MetadataClient.class);
+    MetadataAdminClient mockMetadataClient = EasyMock.mock(MetadataAdminClient.class);
     // Fake the metadata to never change so the leader movement will timeout.
     Node node0 = new Node(BROKER_ID_0, "host0", 100);
     Node node1 = new Node(BROKER_ID_1, "host1", 100);
@@ -512,9 +514,7 @@ public void testTimeoutAndExecutionStop() throws InterruptedException, OngoingEx
     PartitionInfo partitionInfo = new PartitionInfo(TP1.topic(), TP1.partition(), node1, replicas, replicas);
     Cluster cluster = new Cluster("id", Arrays.asList(node0, node1), Collections.singleton(partitionInfo),
                                   Collections.emptySet(), Collections.emptySet());
-    MetadataClient.ClusterAndGeneration clusterAndGeneration = new MetadataClient.ClusterAndGeneration(cluster, 0);
-    EasyMock.expect(mockMetadataClient.refreshMetadata()).andReturn(clusterAndGeneration).anyTimes();
-    EasyMock.expect(mockMetadataClient.cluster()).andReturn(clusterAndGeneration.cluster()).anyTimes();
+    EasyMock.expect(mockMetadataClient.cluster()).andReturn(cluster).anyTimes();
     LoadMonitor mockLoadMonitor = getMockLoadMonitor();
     AnomalyDetectorManager mockAnomalyDetectorManager = getMockAnomalyDetector(RANDOM_UUID, false);
     UserTaskManager.UserTaskInfo mockUserTaskInfo = getMockUserTaskInfo();

gradle.properties

@@ -6,4 +6,4 @@ kafkaVersion=4.0.0
 nettyVersion=4.1.118.Final
 jettyVersion=9.4.56.v20240826
 vertxVersion=4.5.8
-testcontainersVersion=1.21.3
+testcontainersVersion=1.21.4