org.apache.kafka.tools.LogDirsCommand Maven / Gradle / Ivy

Go to download
/*
 * 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.tools;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import joptsimple.OptionSpec;
import joptsimple.OptionSpecBuilder;
import org.apache.kafka.clients.admin.Admin;
import org.apache.kafka.clients.admin.AdminClientConfig;
import org.apache.kafka.clients.admin.DescribeLogDirsResult;
import org.apache.kafka.clients.admin.LogDirDescription;
import org.apache.kafka.clients.admin.ReplicaInfo;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.utils.Exit;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.server.util.CommandDefaultOptions;
import org.apache.kafka.server.util.CommandLineUtils;

import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class LogDirsCommand {

    public static void main(String... args) {
        Exit.exit(mainNoExit(args));
    }

    static int mainNoExit(String... args) {
        try {
            execute(args);
            return 0;
        } catch (TerseException e) {
            System.err.println(e.getMessage());
            return 1;
        } catch (Throwable e) {
            System.err.println(e.getMessage());
            System.err.println(Utils.stackTrace(e));
            return 1;
        }
    }

    private static void execute(String... args) throws Exception {
        LogDirsCommandOptions options = new LogDirsCommandOptions(args);
        try (Admin adminClient = createAdminClient(options)) {
            execute(options, adminClient);
        }
    }

    static void execute(LogDirsCommandOptions options, Admin adminClient) throws Exception {
        Set topics = options.topics();
        Set clusterBrokers = adminClient.describeCluster().nodes().get().stream().map(Node::id).collect(Collectors.toSet());
        Set inputBrokers = options.brokers();
        Set existingBrokers = inputBrokers.isEmpty() ? new HashSet<>(clusterBrokers) : new HashSet<>(inputBrokers);
        existingBrokers.retainAll(clusterBrokers);
        Set nonExistingBrokers = new HashSet<>(inputBrokers);
        nonExistingBrokers.removeAll(clusterBrokers);

        if (!nonExistingBrokers.isEmpty()) {
            throw new TerseException(
                    String.format(
                            "ERROR: The given brokers do not exist from --broker-list: %s. Current existent brokers: %s",
                            commaDelimitedStringFromIntegerSet(nonExistingBrokers),
                            commaDelimitedStringFromIntegerSet(clusterBrokers)));
        } else {
            System.out.println("Querying brokers for log directories information");
            DescribeLogDirsResult describeLogDirsResult = adminClient.describeLogDirs(existingBrokers);
            Map> logDirInfosByBroker = describeLogDirsResult.allDescriptions().get();

            System.out.printf(
                    "Received log directory information from brokers %s%n",
                    commaDelimitedStringFromIntegerSet(existingBrokers));
            System.out.println(formatAsJson(logDirInfosByBroker, topics));
        }
    }

    private static String commaDelimitedStringFromIntegerSet(Set set) {
        return set.stream().map(String::valueOf).collect(Collectors.joining(","));
    }

    private static List> fromReplicasInfoToPrintableRepresentation(Map replicasInfo) {
        return replicasInfo.entrySet().stream().map(entry -> {
            TopicPartition topicPartition = entry.getKey();
            return new HashMap() {{
                    put("partition", topicPartition.toString());
                    put("size", entry.getValue().size());
                    put("offsetLag", entry.getValue().offsetLag());
                    put("isFuture", entry.getValue().isFuture());
                }};
        }).collect(Collectors.toList());
    }

    private static List> fromLogDirInfosToPrintableRepresentation(Map logDirInfos, Set topicSet) {
        return logDirInfos.entrySet().stream().map(entry -> {
            String logDir = entry.getKey();
            return new HashMap() {{
                    put("logDir", logDir);
                    put("error", entry.getValue().error() != null ? entry.getValue().error().getClass().getName() : null);
                    put("partitions", fromReplicasInfoToPrintableRepresentation(
                            entry.getValue().replicaInfos().entrySet().stream().filter(entry -> {
                                TopicPartition topicPartition = entry.getKey();
                                return topicSet.isEmpty() || topicSet.contains(topicPartition.topic());
                            }).collect(Collectors.toMap(Entry::getKey, Entry::getValue))
                    ));
                }};
        }).collect(Collectors.toList());
    }

    private static String formatAsJson(Map> logDirInfosByBroker, Set topicSet) throws JsonProcessingException {
        return new ObjectMapper().writeValueAsString(new HashMap() {{
                put("version", 1);
                put("brokers", logDirInfosByBroker.entrySet().stream().map(entry -> {
                    int broker = entry.getKey();
                    Map logDirInfos = entry.getValue();
                    return new HashMap() {{
                            put("broker", broker);
                            put("logDirs", fromLogDirInfosToPrintableRepresentation(logDirInfos, topicSet));
                        }};
                }).collect(Collectors.toList()));
            }});
    }

    private static Admin createAdminClient(LogDirsCommandOptions options) throws IOException {
        Properties props = new Properties();
        if (options.hasCommandConfig()) {
            props.putAll(Utils.loadProps(options.commandConfig()));
        }
        props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, options.bootstrapServers());
        props.putIfAbsent(AdminClientConfig.CLIENT_ID_CONFIG, "log-dirs-tool");
        return Admin.create(props);
    }

    // Visible for testing
    static class LogDirsCommandOptions extends CommandDefaultOptions {
        private final OptionSpec bootstrapServerOpt;
        private final OptionSpec commandConfigOpt;
        private final OptionSpecBuilder describeOpt;
        private final OptionSpec topicListOpt;
        private final OptionSpec brokerListOpt;

        public LogDirsCommandOptions(String... args) {
            super(args);

            bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED: the server(s) to use for bootstrapping")
                    .withRequiredArg()
                    .describedAs("The server(s) to use for bootstrapping")
                    .ofType(String.class);
            commandConfigOpt = parser.accepts("command-config", "Property file containing configs to be passed to Admin Client.")
                    .withRequiredArg()
                    .describedAs("Admin client property file")
                    .ofType(String.class);
            describeOpt = parser.accepts("describe", "Describe the specified log directories on the specified brokers.");
            topicListOpt = parser.accepts("topic-list", "The list of topics to be queried in the form \"topic1,topic2,topic3\". " +
                            "All topics will be queried if no topic list is specified")
                    .withRequiredArg()
                    .describedAs("Topic list")
                    .defaultsTo("")
                    .ofType(String.class);
            brokerListOpt = parser.accepts("broker-list", "The list of brokers to be queried in the form \"0,1,2\". " +
                            "All brokers in the cluster will be queried if no broker list is specified")
                    .withRequiredArg()
                    .describedAs("Broker list")
                    .ofType(String.class)
                    .defaultsTo("");

            options = parser.parse(args);

            CommandLineUtils.maybePrintHelpOrVersion(this, "This tool helps to query log directory usage on the specified brokers.");

            CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt, describeOpt);
        }

        private Stream splitAtCommasAndFilterOutEmpty(OptionSpec option) {
            return Arrays.stream(options.valueOf(option).split(",")).filter(x -> !x.isEmpty());
        }

        private String bootstrapServers() {
            return options.valueOf(bootstrapServerOpt);
        }

        private boolean hasCommandConfig() {
            return options.has(commandConfigOpt);
        }

        private String commandConfig() {
            return options.valueOf(commandConfigOpt);
        }

        private Set topics() {
            return splitAtCommasAndFilterOutEmpty(topicListOpt).collect(Collectors.toSet());
        }

        private Set brokers() {
            return splitAtCommasAndFilterOutEmpty(brokerListOpt).map(Integer::valueOf).collect(Collectors.toSet());
        }
    }
}