KafkaReadyCommand.java

/*
 * Copyright 2017 Confluent Inc.
 * <p>
 * Licensed 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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 io.confluent.admin.utils.cli;

import net.sourceforge.argparse4j.ArgumentParsers;
import net.sourceforge.argparse4j.inf.ArgumentParser;
import net.sourceforge.argparse4j.inf.ArgumentParserException;
import net.sourceforge.argparse4j.inf.MutuallyExclusiveGroup;
import net.sourceforge.argparse4j.inf.Namespace;

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.common.utils.Utils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import io.confluent.admin.utils.ClusterStatus;

import static net.sourceforge.argparse4j.impl.Arguments.store;

/**
 * This command checks if the kafka cluster has the expected number of brokers and is ready to
 * accept
 * requests.
 * where:
 * config                 : path to properties with client config.
 * min-expected-brokers   : minimum brokers to wait for.
 * timeout                : timeout in ms for all operations. This includes looking up metadata in
 * Zookeeper or fetching metadata for the brokers.
 * (bootstrap-brokers
 * or
 * zookeeper-connect)     : Either a bootstrap broker list or zookeeper connect string
 * security-protocol      : Security protocol to use to connect to the broker.
 */
public class KafkaReadyCommand {

  private static final Logger log = LogManager.getLogger(KafkaReadyCommand.class);
  public static final String KAFKA_READY = "kafka-ready";
  private static final String CONFIG_PROVIDERS_PREFIX = "config.providers";

  // When set to "true", config.providers entries are stripped from the worker config
  // before running the kafka-ready check. This prevents ClassNotFoundException when
  // config provider plugin JARs are on the worker's plugin.path but not on the
  // CUB_CLASSPATH used by kafka-ready. Default: disabled (original behavior).
  static final String SKIP_CONFIG_PROVIDERS_ENV = "CUB_KAFKA_READY_SKIP_CONFIG_PROVIDERS";

  private static ArgumentParser createArgsParser() {
    ArgumentParser kafkaReady = ArgumentParsers
        .newArgumentParser(KAFKA_READY)
        .defaultHelp(true)
        .description("Check if Kafka is ready.");

    kafkaReady.addArgument("min-expected-brokers")
        .action(store())
        .required(true)
        .type(Integer.class)
        .metavar("MIN_EXPECTED_BROKERS")
        .help("Minimum number of brokers to wait for.");

    kafkaReady.addArgument("timeout")
        .action(store())
        .required(true)
        .type(Integer.class)
        .metavar("TIMEOUT_IN_MS")
        .help("Time (in ms) to wait for service to be ready.");

    kafkaReady.addArgument("--config", "-c")
        .action(store())
        .type(String.class)
        .metavar("CONFIG")
        .help("List of bootstrap brokers.");

    MutuallyExclusiveGroup kafkaOrZK = kafkaReady.addMutuallyExclusiveGroup();
    kafkaOrZK.addArgument("--bootstrap-servers", "-b")
        .action(store())
        .type(String.class)
        .metavar("BOOTSTRAP_SERVERS")
        .help("List of bootstrap brokers.");

    kafkaOrZK.addArgument("--zookeeper-connect", "-z")
        .action(store())
        .type(String.class)
        .metavar("ZOOKEEPER_CONNECT")
        .help("Zookeeper connect string.");

    kafkaReady.addArgument("--security-protocol", "-s")
        .action(store())
        .type(String.class)
        .metavar("SECURITY_PROTOCOL")
        .setDefault("PLAINTEXT")
        .help("Which endpoint to connect to ? ");

    return kafkaReady;
  }

  public static void main(String[] args) {
    ArgumentParser parser = createArgsParser();
    boolean success = false;
    try {
      Namespace res = parser.parseArgs(args);
      log.debug("Arguments {}. ", res);

      Map<String, String> workerProps = new HashMap<>();

      if (res.getString("config") == null
          && !(res.getString("security_protocol").equals("PLAINTEXT"))) {
        log.error("config is required for all protocols except PLAINTEXT");
        success = false;
      } else {
        if (res.getString("config") != null) {
          workerProps = Utils.propsToStringMap(Utils.loadProps(res.getString("config")));
        }
        if (res.getString("bootstrap_servers") != null) {
          workerProps.put(
              CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
              res.getString("bootstrap_servers")
          );
        } else {
          log.error("Bootstrap servers should be provided through config or bootstrap_servers");
          throw new RuntimeException(
              "Bootstrap servers should be provided through config or bootstrap_servers"
          );
        }

        if (isSkipConfigProvidersEnabled()) {
          workerProps = stripConfigProviders(workerProps);
        }

        success = ClusterStatus.isKafkaReady(
            workerProps,
            res.getInt("min_expected_brokers"),
            res.getInt("timeout")
        );
      }

    } catch (ArgumentParserException e) {
      if (args.length == 0) {
        parser.printHelp();
        success = true;
      } else {
        parser.handleError(e);
        success = false;
      }
    } catch (Exception e) {
      log.error("Error while running kafka-ready.", e);
      success = false;
    }

    if (success) {
      System.exit(0);
    } else {
      System.exit(1);
    }
  }

  /**
   * Returns a copy of the properties map with config.providers entries removed.
   * If no config.providers entries are found, returns the original map unmodified.
   */
  static Map<String, String> stripConfigProviders(Map<String, String> props) {
    boolean hasProviderKeys = false;
    for (String key : props.keySet()) {
      if (key.startsWith(CONFIG_PROVIDERS_PREFIX)) {
        hasProviderKeys = true;
        break;
      }
    }
    if (!hasProviderKeys) {
      return props;
    }

    Map<String, String> result = new HashMap<>(props);
    int count = 0;
    Iterator<String> it = result.keySet().iterator();
    while (it.hasNext()) {
      String key = it.next();
      if (key.startsWith(CONFIG_PROVIDERS_PREFIX)) {
        it.remove();
        count++;
      }
    }
    log.info("Stripped {} config.providers properties from kafka-ready config ({}=true).",
        count, SKIP_CONFIG_PROVIDERS_ENV);
    return result;
  }

  static boolean isSkipConfigProvidersEnabled() {
    return "true".equalsIgnoreCase(System.getenv(SKIP_CONFIG_PROVIDERS_ENV));
  }
}