From 67012ed15c38852a56c25d1babbd22473f6eae06 Mon Sep 17 00:00:00 2001
From: Shuo <js982986555@live.com>
Date: Thu, 19 Mar 2020 23:19:08 +0800
Subject: [PATCH] feat: add query_disk_info api for shell command (#416)

---
 .../dist/replication/replication_ddl_client.h | 41 +++++++++++++++++++
 .../ddl_lib/replication_ddl_client.cpp        | 17 ++++++++
 2 files changed, 58 insertions(+)

diff --git a/include/dsn/dist/replication/replication_ddl_client.h b/include/dsn/dist/replication/replication_ddl_client.h
index fff8d58a7b..c6dc3d3efb 100644
--- a/include/dsn/dist/replication/replication_ddl_client.h
+++ b/include/dsn/dist/replication/replication_ddl_client.h
@@ -169,6 +169,11 @@ class replication_ddl_client
 
     dsn::error_code ddd_diagnose(gpid pid, std::vector<ddd_partition_info> &ddd_partitions);
 
+    void query_disk_info(
+        const std::vector<dsn::rpc_address> &targets,
+        const std::string &app_name,
+        /*out*/ std::map<dsn::rpc_address, error_with<query_disk_info_response>> &resps);
+
 private:
     bool static valid_app_char(int c);
 
@@ -222,9 +227,45 @@ class replication_ddl_client
         return error_with<TResponse>(std::move(rpc.response()));
     }
 
+    /// Send request to multi replica server synchronously.
+    template <typename TRpcHolder, typename TResponse = typename TRpcHolder::response_type>
+    void call_rpcs_async(std::map<dsn::rpc_address, TRpcHolder> &rpcs,
+                         std::map<dsn::rpc_address, error_with<TResponse>> &resps,
+                         int reply_thread_hash = 0,
+                         bool enable_retry = true)
+    {
+        dsn::task_tracker tracker;
+        error_code err = ERR_UNKNOWN;
+        for (auto &rpc : rpcs) {
+            rpc.second.call(
+                rpc.first, &tracker, [&err, &resps, &rpcs, &rpc](error_code code) mutable {
+                    err = code;
+                    if (err == dsn::ERR_OK) {
+                        resps.emplace(rpc.first, std::move(rpc.second.response()));
+                        rpcs.erase(rpc.first);
+                    } else {
+                        resps.emplace(
+                            rpc.first,
+                            std::move(error_s::make(err, "unable to send rpc to server")));
+                    }
+                });
+        }
+        tracker.wait_outstanding_tasks();
+
+        if (enable_retry && rpcs.size() > 0) {
+            std::map<dsn::rpc_address, dsn::error_with<TResponse>> retry_resps;
+            call_rpcs_async(rpcs, retry_resps, reply_thread_hash, false);
+            for (auto &resp : retry_resps) {
+                resps.emplace(resp.first, std::move(resp.second));
+            }
+        }
+    }
+
 private:
     dsn::rpc_address _meta_server;
     dsn::task_tracker _tracker;
+
+    typedef rpc_holder<query_disk_info_request, query_disk_info_response> query_disk_info_rpc;
 };
 } // namespace replication
 } // namespace dsn
diff --git a/src/dist/replication/ddl_lib/replication_ddl_client.cpp b/src/dist/replication/ddl_lib/replication_ddl_client.cpp
index f3bcdac646..9b45e4b49e 100644
--- a/src/dist/replication/ddl_lib/replication_ddl_client.cpp
+++ b/src/dist/replication/ddl_lib/replication_ddl_client.cpp
@@ -1517,5 +1517,22 @@ replication_ddl_client::ddd_diagnose(gpid pid, std::vector<ddd_partition_info> &
 
     return dsn::ERR_OK;
 }
+
+void replication_ddl_client::query_disk_info(
+    const std::vector<dsn::rpc_address> &targets,
+    const std::string &app_name,
+    /*out*/ std::map<dsn::rpc_address, error_with<query_disk_info_response>> &resps)
+{
+    std::map<dsn::rpc_address, query_disk_info_rpc> query_disk_info_rpcs;
+    for (const auto &target : targets) {
+        auto request = make_unique<query_disk_info_request>();
+        request->node = target;
+        request->app_name = app_name;
+        query_disk_info_rpcs.emplace(target,
+                                     query_disk_info_rpc(std::move(request), RPC_QUERY_DISK_INFO));
+    }
+    call_rpcs_async(query_disk_info_rpcs, resps);
+}
+
 } // namespace replication
 } // namespace dsn