Addressing feedback

Signed-off-by: Kyle Liberti <[email protected]>

Author: kyguy

cruise-control-metrics-reporter/src/main/java/com/linkedin/kafka/cruisecontrol/metricsreporter/CruiseControlMetricsReporter.java

@@ -566,7 +566,8 @@ private void setIfAbsent(Properties props, String key, String value) {
    *         exceptions related to reflection (e.g., {@link NoSuchMethodException}), invocation issues,
    *         execution exceptions, timeouts, and interruptions.
    */
-  /* test */ static TopicDescription getTopicDescription(AdminClient adminClient, String ccMetricsTopic) throws KafkaTopicDescriptionException {
+  /* test */ public static TopicDescription getTopicDescription(AdminClient adminClient, String ccMetricsTopic)
+    throws KafkaTopicDescriptionException {
     try {
       // For compatibility with Kafka 4.0 and beyond we must use new API methods.
       Method topicDescriptionMethod = topicNameValuesMethod();

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

@@ -4,7 +4,6 @@
 
 package com.linkedin.kafka.cruisecontrol.metricsreporter;
 
-import com.linkedin.kafka.cruisecontrol.metricsreporter.exception.KafkaTopicDescriptionException;
 import com.linkedin.kafka.cruisecontrol.metricsreporter.metric.CruiseControlMetric;
 import com.linkedin.kafka.cruisecontrol.metricsreporter.metric.MetricSerde;
 import com.linkedin.kafka.cruisecontrol.metricsreporter.utils.CCContainerizedKraftCluster;
@@ -20,11 +19,8 @@
 import java.util.Set;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Predicate;
 import java.util.regex.Pattern;
 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.TopicDescription;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -36,7 +32,6 @@
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
-import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.apache.kafka.common.serialization.StringDeserializer;
 import org.junit.After;
 import org.junit.Before;
@@ -45,7 +40,6 @@
 
 import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporter.DEFAULT_BOOTSTRAP_SERVERS_HOST;
 import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporter.DEFAULT_BOOTSTRAP_SERVERS_PORT;
-import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporter.getTopicDescription;
 import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporterConfig.CRUISE_CONTROL_METRICS_REPORTER_INTERVAL_MS_CONFIG;
 import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporterConfig.CRUISE_CONTROL_METRICS_TOPIC_AUTO_CREATE_CONFIG;
 import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporterConfig.CRUISE_CONTROL_METRICS_TOPIC_CONFIG;
@@ -206,41 +200,9 @@ public void testReportingMetrics() {
     assertEquals("Expected " + expectedMetricTypes + ", but saw " + metricTypes, expectedMetricTypes, metricTypes);
   }
 
-  private TopicDescription waitForTopicMetadata(Admin adminClient,
-                                                Duration timeout,
-                                                Predicate<TopicDescription> condition)
-    throws InterruptedException, TimeoutException {
-
-    long deadline = System.currentTimeMillis() + timeout.toMillis();
-
-    while (System.currentTimeMillis() < deadline) {
-      try {
-        TopicDescription topicDescription = getTopicDescription((AdminClient) adminClient, TOPIC);
-
-        if (condition.test(topicDescription)) {
-          return topicDescription;
-        }
-      } catch (KafkaTopicDescriptionException e) {
-        if (!(e.getCause() instanceof UnknownTopicOrPartitionException)) {
-          throw new RuntimeException("Failed to describe topic: " + TOPIC, e);
-        }
-        // else ignore and retry
-      }
-
-      Thread.sleep(500);
-    }
-
-    throw new TimeoutException("Timeout waiting for topic metadata condition to be met: " + TOPIC);
-  }
-
   @Test
   public void testUpdatingMetricsTopicConfig() throws InterruptedException, TimeoutException {
-    Properties props = new Properties();
-    setSecurityConfigs(props, "admin");
-    props.setProperty(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers());
-    Admin adminClient = Admin.create(props);
-
-    TopicDescription topicDescription = waitForTopicMetadata(adminClient, Duration.ofSeconds(30), td -> true);
+    TopicDescription topicDescription = _cluster.waitForTopicMetadata(TOPIC, Duration.ofSeconds(30), td -> true);
     assertEquals(1, topicDescription.partitions().size());
 
     KafkaContainer broker = _cluster.getBrokers().get(0);
@@ -261,7 +223,7 @@ public void testUpdatingMetricsTopicConfig() throws InterruptedException, Timeou
 
     // Wait for topic metadata configuration change to propagate
     int oldPartitionCount = topicDescription.partitions().size();
-    TopicDescription newTopicDescription = waitForTopicMetadata(adminClient, Duration.ofSeconds(30),
+    TopicDescription newTopicDescription = _cluster.waitForTopicMetadata(TOPIC, Duration.ofSeconds(30),
       td -> td.partitions().size() != oldPartitionCount);
 
     assertEquals(2, newTopicDescription.partitions().size());

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

@@ -7,11 +7,14 @@
 import com.github.dockerjava.api.command.InspectContainerResponse;
 import com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporter;
 import com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporterConfig;
+import com.linkedin.kafka.cruisecontrol.metricsreporter.exception.KafkaTopicDescriptionException;
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.admin.DescribeClusterResult;
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.TopicDescription;
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.Network;
 import org.testcontainers.containers.wait.strategy.AbstractWaitStrategy;
@@ -33,9 +36,12 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Predicate;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporter.getTopicDescription;
 import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporterConfig.CRUISE_CONTROL_METRICS_REPORTER_INTERVAL_MS_CONFIG;
 import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporterConfig.CRUISE_CONTROL_METRICS_TOPIC_CONFIG;
 import static com.linkedin.kafka.cruisecontrol.metricsreporter.CruiseControlMetricsReporterConfig.DEFAULT_CRUISE_CONTROL_METRICS_TOPIC;
@@ -79,6 +85,7 @@ public class CCContainerizedKraftCluster implements Startable {
   private final List<KafkaContainer> _brokers;
   private final String _bootstrapServers;
   private final List<String> _brokerAddressList;
+  private final Admin _adminClient;
 
   public CCContainerizedKraftCluster(int numOfBrokers, List<Map<Object, Object>> brokerConfigs, Properties adminClientProps) {
     if (numOfBrokers <= 0) {
@@ -105,6 +112,7 @@ public CCContainerizedKraftCluster(int numOfBrokers, List<Map<Object, Object>> b
     Properties adminClientPropsWithBootstrapAddress = new Properties();
     adminClientPropsWithBootstrapAddress.putAll(adminClientProps);
     adminClientPropsWithBootstrapAddress.setProperty(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, _bootstrapServers);
+    this._adminClient = Admin.create(adminClientPropsWithBootstrapAddress);
 
     this._brokers =
       IntStream
@@ -154,7 +162,7 @@ protected void containerIsStarting(InspectContainerResponse containerInfo) {
             //.withLogConsumer(outputFrame -> System.out.print(networkAlias + " | " + outputFrame.getUtf8String()))
             // Uncomment the following line when debugging SSL connection problems.
             //.withEnv("KAFKA_OPTS", "-Djavax.net.debug=ssl,handshake")
-            .waitingFor(new BrokerWaitStrategy(brokerNum, metricsTopic, adminClientPropsWithBootstrapAddress)
+            .waitingFor(new BrokerWaitStrategy(brokerNum, metricsTopic, _adminClient)
               .withStartupTimeout(Duration.ofMinutes(1))
             );
           kafkaContainer.setPortBindings(List.of(containerHostPort + ":" + CONTAINER_EXTERNAL_LISTENER_PORT));
@@ -300,49 +308,133 @@ public void start() {
 
   @Override
   public void stop() {
+    this._adminClient.close();
     this._brokers.stream().parallel().forEach(GenericContainer::close);
   }
 
+  private static <T> T waitUntil(Supplier<T> supplier, Predicate<T> condition, Duration timeout, String timeoutMessage) {
+    long deadline = System.currentTimeMillis() + timeout.toMillis();
+
+    while (System.currentTimeMillis() < deadline) {
+      T value;
+      try {
+        value = supplier.get();
+      } catch (Exception e) {
+        throw new RuntimeException(e);
+      }
+
+      if (condition.test(value)) {
+        return value;
+      }
+
+      try {
+        Thread.sleep(500);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException(e);
+      }
+    }
+
+    throw new RuntimeException(timeoutMessage);
+  }
+
+  /**
+   * Waits until the metadata for a Kafka topic meets the specified condition,
+   * or until the given timeout period elapses.
+   *
+   * @param topicName the name of the topic whose metadata should be monitored
+   * @param timeout   the maximum duration to wait for the condition to be satisfied
+   * @param condition a {@link Predicate} that tests the {@link TopicDescription} for readiness
+   * @return the {@link TopicDescription} once the condition evaluates to {@code true}
+   */
+  public TopicDescription waitForTopicMetadata(String topicName, Duration timeout, Predicate<TopicDescription> condition) {
+    return waitUntil(
+      () -> {
+        try {
+          return getTopicDescription((AdminClient) _adminClient, topicName);
+        } catch (KafkaTopicDescriptionException e) {
+          if (e.getCause() instanceof UnknownTopicOrPartitionException) {
+            return null;
+          }
+          throw new RuntimeException("Failed to describe topic: " + topicName, e);
+        }
+      },
+      td -> td != null && condition.test(td),
+      timeout,
+      String.format("Timeout waiting for topic %s metadata to be ready.", topicName)
+    );
+  }
+
+  /**
+   * Shuts down the Kafka broker with the given broker ID and waits for it
+   * to be removed from the cluster metadata.
+   *
+   * @param brokerId the ID of the broker to wait for shutdown.
+   */
+  public void shutDownBroker(int brokerId) {
+    Duration timeout = Duration.ofSeconds(30);
+    _brokers.get(brokerId).stop();
+
+    // Wait until brokerId is no longer in the cluster metadata.
+    waitUntil(
+      () -> {
+        try {
+          return _adminClient.describeCluster().nodes().get().stream()
+            .map(node -> String.valueOf(node.id()))
+            .collect(Collectors.toSet());
+        } catch (InterruptedException | ExecutionException e) {
+          throw new RuntimeException(e);
+        }
+      },
+      brokerIds -> !brokerIds.contains(String.valueOf(brokerId)),
+      timeout,
+      String.format("Broker %s was not removed from cluster metadata after shutdown.", brokerId)
+    );
+  }
+
   public static class BrokerWaitStrategy extends AbstractWaitStrategy {
     private final int _brokerId;
     private final String _metricsTopic;
-    private final Properties _adminClientProps;
+    private final Admin _adminClient;
 
-    public BrokerWaitStrategy(int brokerId, String metricsTopic, Properties adminClientProps) {
+    public BrokerWaitStrategy(int brokerId, String metricsTopic, Admin adminClient) {
       this._brokerId = brokerId;
       this._metricsTopic = metricsTopic;
-      this._adminClientProps = adminClientProps;
+      this._adminClient = adminClient;
     }
 
     @Override
     protected void waitUntilReady() {
-      long deadline = System.currentTimeMillis() + startupTimeout.toMillis();
-
-      try (Admin adminClient = Admin.create(_adminClientProps)) {
-        while (System.currentTimeMillis() < deadline) {
+      waitUntil(
+        () -> {
+          // Returns true if broker is online and metrics topic exists
           try {
-            DescribeClusterResult cluster = adminClient.describeCluster();
-            boolean brokerOnline = cluster.nodes().get().stream().anyMatch(node -> node.id() == _brokerId);
+            boolean brokerOnline = _adminClient.describeCluster().nodes().get()
+              .stream()
+              .anyMatch(node -> node.id() == _brokerId);
 
             if (!brokerOnline) {
-              Thread.sleep(500);
-              continue;
+              return false;
             }
 
-            adminClient.describeTopics(Collections.singleton(_metricsTopic)).allTopicNames().get(5, TimeUnit.SECONDS);
-            return;
+            return _adminClient.describeTopics(Collections.singleton(_metricsTopic))
+              .allTopicNames()
+              .get(5, TimeUnit.SECONDS)
+              .containsKey(_metricsTopic);
 
           } catch (InterruptedException e) {
             // Restore interrupt status.
             Thread.currentThread().interrupt();
-            throw new RuntimeException("Interrupted while waiting for broker to become ready", e);
+            throw new RuntimeException(e);
           } catch (ExecutionException | TimeoutException e) {
             // Kafka broker is not ready yet, ignore and retry.
+            return false;
           }
-        }
-
-        throw new RuntimeException(String.format("Broker %d did not become ready within timeout of %d ms", _brokerId, startupTimeout.toMillis()));
-      }
+        },
+        ready -> ready,
+        startupTimeout,
+        String.format("Broker %d did not become ready within timeout of %d ms", _brokerId, startupTimeout.toMillis())
+      );
     }
   }
 }