feat: add shell command for query_disk_info (#498)

src/shell/command_output.h

@@ -0,0 +1,29 @@
+// Copyright (c) 2017, Xiaomi, Inc.  All rights reserved.
+// This source code is licensed under the Apache License Version 2.0, which
+// can be found in the LICENSE file in the root directory of this source tree.
+
+#include "shell/commands.h"
+#include <dsn/utility/smart_pointers.h>
+
+class command_output
+{
+public:
+    explicit command_output(const std::string &file_name) : _file_name(file_name)
+    {
+        if (!file_name.empty()) {
+            _file_stream = dsn::make_unique<std::ofstream>(_file_name);
+        }
+    }
+    std::ostream *stream() const
+    {
+        if (_file_stream && !_file_stream->is_open()) {
+            fmt::print(stderr, "open output file {} failed!\n", _file_name);
+            return nullptr;
+        }
+        return _file_stream ? _file_stream.get() : &std::cout;
+    }
+
+private:
+    std::string _file_name;
+    std::unique_ptr<std::ofstream> _file_stream;
+};

src/shell/commands.h

@@ -239,3 +239,9 @@ bool remove_dup(command_executor *e, shell_context *sc, arguments args);
 bool start_dup(command_executor *e, shell_context *sc, arguments args);
 
 bool pause_dup(command_executor *e, shell_context *sc, arguments args);
+
+// == disk rebalance (see 'commands/disk_rebalance.cpp') == //
+
+bool query_disk_capacity(command_executor *e, shell_context *sc, arguments args);
+
+bool query_disk_replica(command_executor *e, shell_context *sc, arguments args);

src/shell/commands/disk_rebalance.cpp

@@ -0,0 +1,258 @@
+// Copyright (c) 2017, Xiaomi, Inc.  All rights reserved.
+// This source code is licensed under the Apache License Version 2.0, which
+// can be found in the LICENSE file in the root directory of this source tree.
+
+#include "shell/commands.h"
+#include "shell/argh.h"
+#include "shell/command_output.h"
+
+#include <math.h>
+#include <fmt/ostream.h>
+#include <dsn/utility/errors.h>
+#include <dsn/utility/output_utils.h>
+#include <dsn/utility/string_conv.h>
+#include <dsn/dist/replication/duplication_common.h>
+
+bool validate_cmd(const argh::parser &cmd,
+                  const std::set<std::string> &params,
+                  const std::set<std::string> &flags)
+{
+    if (cmd.size() > 1) {
+        fmt::print(stderr, "too many params!\n");
+        return false;
+    }
+
+    for (const auto &param : cmd.params()) {
+        if (params.find(param.first) == params.end()) {
+            fmt::print(stderr, "unknown param {} = {}\n", param.first, param.second);
+            return false;
+        }
+    }
+
+    for (const auto &flag : cmd.flags()) {
+        if (params.find(flag) != params.end()) {
+            fmt::print(stderr, "missing value of {}\n", flag);
+            return false;
+        }
+
+        if (flags.find(flag) == flags.end()) {
+            fmt::print(stderr, "unknown flag {}\n", flag);
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool query_disk_info(
+    shell_context *sc,
+    const argh::parser &cmd,
+    const std::string &node_address,
+    const std::string &app_name,
+    /*out*/ std::map<dsn::rpc_address, dsn::error_with<query_disk_info_response>> &err_resps)
+{
+    std::map<dsn::rpc_address, dsn::replication::node_status::type> nodes;
+    auto error = sc->ddl_client->list_nodes(::dsn::replication::node_status::NS_INVALID, nodes);
+    if (error != dsn::ERR_OK) {
+        fmt::print(stderr, "list nodes failed, error={}\n", error.to_string());
+        return false;
+    }
+
+    std::vector<dsn::rpc_address> targets;
+    for (const auto &node : nodes) {
+        if (node_address.empty() || node_address == node.first.to_std_string()) {
+            targets.push_back(node.first);
+            if (!node_address.empty())
+                break;
+        }
+    }
+
+    if (targets.empty()) {
+        fmt::print(stderr, "invalid target replica server address!\n");
+        return false;
+    }
+    sc->ddl_client->query_disk_info(targets, app_name, err_resps);
+    return true;
+}
+
+bool query_disk_capacity(command_executor *e, shell_context *sc, arguments args)
+{
+    // disk_capacity [-n|--node replica_server(ip:port)] [-o|--out file_name][-j|-json][-d|--detail]
+    const std::set<std::string> &params = {"n", "node", "o", "out"};
+    const std::set<std::string> &flags = {"j", "json", "d", "detail"};
+    argh::parser cmd(args.argc, args.argv, argh::parser::PREFER_PARAM_FOR_UNREG_OPTION);
+    if (!validate_cmd(cmd, params, flags)) {
+        return false;
+    }
+
+    bool format_to_json = cmd[{"-j", "--json"}];
+    bool query_detail_info = cmd[{"-d", "--detail"}];
+    std::string node_address = cmd({"-n", "--node"}).str();
+    std::string file_name = cmd({"-o", "--out"}).str();
+
+    command_output out(file_name);
+    if (!out.stream()) {
+        return false;
+    }
+
+    std::map<dsn::rpc_address, dsn::error_with<query_disk_info_response>> err_resps;
+    // passing empty app_name(app_name = "") means query all app disk info.
+    if (!query_disk_info(sc, cmd, node_address, "", err_resps)) {
+        return false;
+    }
+
+    dsn::utils::table_printer node_printer;
+    node_printer.add_title("node");
+    node_printer.add_column("total_capacity(MB)");
+    node_printer.add_column("avalable_capacity(MB)");
+    node_printer.add_column("avalable_ratio(%)");
+    node_printer.add_column("capacity_balance");
+
+    dsn::utils::multi_table_printer multi_printer;
+    for (const auto &err_resp : err_resps) {
+        dsn::error_s err = err_resp.second.get_error();
+        if (err.is_ok()) {
+            err = dsn::error_s::make(err_resp.second.get_value().err);
+        }
+        if (!err.is_ok()) {
+            fmt::print(stderr,
+                       "disk of node[{}] info skiped because request failed, error={}\n",
+                       err_resp.first.to_std_string(),
+                       err.description());
+            continue;
+        }
+
+        const auto &resp = err_resp.second.get_value();
+        int total_capacity_ratio =
+            resp.total_capacity_mb == 0
+                ? 0
+                : std::round(resp.total_available_mb * 100.0 / resp.total_capacity_mb);
+
+        int variance = 0;
+        for (const auto &disk_info : resp.disk_infos) {
+            int disk_available_ratio =
+                disk_info.disk_capacity_mb == 0
+                    ? 0
+                    : std::round(disk_info.disk_available_mb * 100.0 / disk_info.disk_capacity_mb);
+            variance += std::pow((disk_available_ratio - total_capacity_ratio), 2);
+        }
+
+        int capacity_balance = std::sqrt(variance);
+
+        if (query_detail_info) {
+            dsn::utils::table_printer disk_printer(err_resp.first.to_std_string());
+            disk_printer.add_title("disk");
+            disk_printer.add_column("total_capacity(MB)");
+            disk_printer.add_column("avalable_capacity(MB)");
+            disk_printer.add_column("avalable_ratio(%)");
+            disk_printer.add_column("capacity_balance");
+
+            for (const auto &disk_info : resp.disk_infos) {
+                int disk_available_ratio = disk_info.disk_capacity_mb == 0
+                                               ? 0
+                                               : std::round(disk_info.disk_available_mb * 100.0 /
+                                                            disk_info.disk_capacity_mb);
+                int disk_density = disk_available_ratio - total_capacity_ratio;
+                disk_printer.add_row(disk_info.tag);
+                disk_printer.append_data(disk_info.disk_capacity_mb);
+                disk_printer.append_data(disk_info.disk_available_mb);
+                disk_printer.append_data(disk_available_ratio);
+                disk_printer.append_data(disk_density);
+            }
+            disk_printer.add_row("total");
+            disk_printer.append_data(resp.total_capacity_mb);
+            disk_printer.append_data(resp.total_available_mb);
+            disk_printer.append_data(total_capacity_ratio);
+            disk_printer.append_data(capacity_balance);
+
+            multi_printer.add(std::move(disk_printer));
+        } else {
+            node_printer.add_row(err_resp.first.to_std_string());
+            node_printer.append_data(resp.total_capacity_mb);
+            node_printer.append_data(resp.total_available_mb);
+            node_printer.append_data(total_capacity_ratio);
+            node_printer.append_data(capacity_balance);
+        }
+    }
+    if (query_detail_info) {
+        multi_printer.output(*out.stream(),
+                             format_to_json ? tp_output_format::kJsonPretty
+                                            : tp_output_format::kTabular);
+    } else {
+        node_printer.output(*out.stream(),
+                            format_to_json ? tp_output_format::kJsonPretty
+                                           : tp_output_format::kTabular);
+    }
+
+    return true;
+}
+
+bool query_disk_replica(command_executor *e, shell_context *sc, arguments args)
+{
+    // disk_capacity [-n|--node replica_server(ip:port)][-a|-app app_name][-o|--out
+    // file_name][-j|--json]
+    const std::set<std::string> &params = {"n", "node", "a", "app", "o", "out"};
+    const std::set<std::string> &flags = {"j", "json"};
+    argh::parser cmd(args.argc, args.argv, argh::parser::PREFER_PARAM_FOR_UNREG_OPTION);
+    if (!validate_cmd(cmd, params, flags)) {
+        return false;
+    }
+
+    bool format_to_json = cmd[{"-j", "--json"}];
+    std::string node_address = cmd({"-n", "--node"}).str();
+    std::string app_name = cmd({"-a", "--app"}).str();
+    std::string file_name = cmd({"-o", "--out"}).str();
+
+    command_output out(file_name);
+    if (!out.stream()) {
+        return false;
+    }
+
+    std::map<dsn::rpc_address, dsn::error_with<query_disk_info_response>> err_resps;
+    if (!query_disk_info(sc, cmd, node_address, app_name, err_resps)) {
+        return false;
+    }
+
+    dsn::utils::multi_table_printer multi_printer;
+    for (const auto &err_resp : err_resps) {
+        dsn::error_s err = err_resp.second.get_error();
+        if (err.is_ok()) {
+            err = dsn::error_s::make(err_resp.second.get_value().err);
+        }
+        if (!err.is_ok()) {
+            fmt::print(stderr,
+                       "disk of node[{}] info skiped because request failed, error={}\n",
+                       err_resp.first.to_std_string(),
+                       err.description());
+            continue;
+        }
+        dsn::utils::table_printer disk_printer(err_resp.first.to_std_string());
+        disk_printer.add_title("disk");
+        disk_printer.add_column("primary_count");
+        disk_printer.add_column("secondary_count");
+        disk_printer.add_column("replica_count");
+
+        const auto &resp = err_resp.second.get_value();
+        for (const auto &disk_info : resp.disk_infos) {
+            int primary_count = 0;
+            int secondary_count = 0;
+            for (const auto &replica_count : disk_info.holding_primary_replica_counts) {
+                primary_count += replica_count.second;
+            }
+
+            for (const auto &replica_count : disk_info.holding_secondary_replica_counts) {
+                secondary_count += replica_count.second;
+            }
+            disk_printer.add_row(disk_info.tag);
+            disk_printer.append_data(primary_count);
+            disk_printer.append_data(secondary_count);
+            disk_printer.append_data(primary_count + secondary_count);
+
+            multi_printer.add(std::move(disk_printer));
+        }
+    }
+    multi_printer.output(
+        *out.stream(), format_to_json ? tp_output_format::kJsonPretty : tp_output_format::kTabular);
+
+    return true;
+}

src/shell/main.cpp

@@ -437,6 +437,14 @@ static command_executor commands[] = {
     {"remove_dup", "remove duplication", "<app_name> <dup_id>", remove_dup},
     {"start_dup", "start duplication", "<app_name> <dup_id>", start_dup},
     {"pause_dup", "pause duplication", "<app_name> <dup_id>", pause_dup},
+    {"disk_capacity",
+     "query disk capacity info",
+     "[-n|--node replica_server(ip:port)][-o|--out file_name][-j|-json][-d|--detail]",
+     query_disk_capacity},
+    {"disk_replica",
+     "query disk replica count info",
+     "[-n|--node replica_server(ip:port)][-a|-app app_name][-o|--out file_name][-j|--json]",
+     query_disk_replica},
     {
         "exit", "exit shell", "", exit_shell,
     },