From 2cec8899af19d7243787e26a582a5c859b6acd9e Mon Sep 17 00:00:00 2001
From: Jiashuo <js982986555@live.com>
Date: Mon, 6 Sep 2021 11:06:47 +0800
Subject: [PATCH] feat: support nfs server rate limit (#901)

---
 include/dsn/http/http_server.h |  3 +--
 src/nfs/nfs_client_impl.cpp    |  2 +-
 src/nfs/nfs_server_impl.cpp    | 37 ++++++++++++++++++++++++++++++++++
 src/nfs/nfs_server_impl.h      |  8 ++++++++
 4 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/include/dsn/http/http_server.h b/include/dsn/http/http_server.h
index fcfb317713..ff5faf7ca6 100644
--- a/include/dsn/http/http_server.h
+++ b/include/dsn/http/http_server.h
@@ -28,8 +28,7 @@ namespace dsn {
 DSN_DECLARE_bool(enable_http_server);
 
 /// The rpc code for all the HTTP RPCs.
-/// Since http is used only for system monitoring, it is restricted to lowest priority.
-DEFINE_TASK_CODE_RPC(RPC_HTTP_SERVICE, TASK_PRIORITY_LOW, THREAD_POOL_DEFAULT);
+DEFINE_TASK_CODE_RPC(RPC_HTTP_SERVICE, TASK_PRIORITY_COMMON, THREAD_POOL_DEFAULT);
 
 enum http_method
 {
diff --git a/src/nfs/nfs_client_impl.cpp b/src/nfs/nfs_client_impl.cpp
index c9cd17eca7..edcec4e92c 100644
--- a/src/nfs/nfs_client_impl.cpp
+++ b/src/nfs/nfs_client_impl.cpp
@@ -75,7 +75,7 @@ DSN_DEFINE_int32("nfs",
                  "maximum retry count when copy failed");
 DSN_DEFINE_int32("nfs",
                  rpc_timeout_ms,
-                 10000,
+                 1e5, // 100s
                  "rpc timeout in milliseconds for nfs copy, "
                  "0 means use default timeout of rpc engine");
 
diff --git a/src/nfs/nfs_server_impl.cpp b/src/nfs/nfs_server_impl.cpp
index be04c2ecd3..d444fcb653 100644
--- a/src/nfs/nfs_server_impl.cpp
+++ b/src/nfs/nfs_server_impl.cpp
@@ -42,6 +42,9 @@
 namespace dsn {
 namespace service {
 
+DSN_DEFINE_int32("nfs", max_send_rate_megabytes, 100, "max rate of send to remote node(MB/s)");
+DSN_TAG_VARIABLE(max_send_rate_megabytes, FT_MUTABLE);
+
 DSN_DECLARE_int32(file_close_timer_interval_ms_on_server);
 DSN_DECLARE_int32(file_close_expire_time_ms);
 
@@ -62,6 +65,9 @@ nfs_service_impl::nfs_service_impl() : ::dsn::serverlet<nfs_service_impl>("nfs")
         "recent_copy_fail_count",
         COUNTER_TYPE_VOLATILE_NUMBER,
         "nfs server copy fail count count in the recent period");
+
+    _send_token_bucket = std::make_unique<folly::DynamicTokenBucket>();
+    register_cli_commands();
 }
 
 void nfs_service_impl::on_copy(const ::dsn::service::copy_request &request,
@@ -131,6 +137,8 @@ void nfs_service_impl::on_copy(const ::dsn::service::copy_request &request,
 
 void nfs_service_impl::internal_read_callback(error_code err, size_t sz, callback_para &cp)
 {
+    _send_token_bucket->consumeWithBorrowAndWait(
+        sz, FLAGS_max_send_rate_megabytes << 20, 1.5 * (FLAGS_max_send_rate_megabytes << 20));
     {
         zauto_lock l(_handles_map_lock);
         auto it = _handles_map.find(cp.file_path);
@@ -241,5 +249,34 @@ void nfs_service_impl::close_file() // release out-of-date file handle
             it++;
     }
 }
+
+// TODO(jiashuo1): just for compatibility, ready to delete it later
+void nfs_service_impl::register_cli_commands()
+{
+    static std::once_flag flag;
+    std::call_once(flag, [&]() {
+        _nfs_max_send_rate_megabytes_cmd = dsn::command_manager::instance().register_command(
+            {"nfs.max_send_rate_megabytes"},
+            "nfs.max_send_rate_megabytes [num]",
+            "control the max rate(MB/s) to send file to remote node",
+            [](const std::vector<std::string> &args) {
+                std::string result("OK");
+
+                if (args.empty()) {
+                    return std::to_string(FLAGS_max_send_rate_megabytes);
+                }
+
+                int32_t max_send_rate_megabytes = 0;
+                if (!dsn::buf2int32(args[0], max_send_rate_megabytes) ||
+                    max_send_rate_megabytes <= 0) {
+                    return std::string("ERR: invalid arguments");
+                }
+
+                FLAGS_max_send_rate_megabytes = max_send_rate_megabytes;
+                return result;
+            });
+    });
+}
+
 } // namespace service
 } // namespace dsn
diff --git a/src/nfs/nfs_server_impl.h b/src/nfs/nfs_server_impl.h
index 9abd89b516..49edf96760 100644
--- a/src/nfs/nfs_server_impl.h
+++ b/src/nfs/nfs_server_impl.h
@@ -30,6 +30,7 @@
 #include <iostream>
 #include <dsn/cpp/serverlet.h>
 #include <dsn/utility/flags.h>
+#include <dsn/tool-api/command_manager.h>
 
 #include "nfs_code_definition.h"
 #include "nfs_types.h"
@@ -50,10 +51,13 @@ class nfs_service_impl : public ::dsn::serverlet<nfs_service_impl>
             RPC_NFS_GET_FILE_SIZE, "get_file_size", &nfs_service_impl::on_get_file_size);
     }
 
+    void register_cli_commands();
+
     void close_service()
     {
         unregister_rpc_handler(RPC_NFS_COPY);
         unregister_rpc_handler(RPC_NFS_GET_FILE_SIZE);
+        UNREGISTER_VALID_HANDLER(_nfs_max_send_rate_megabytes_cmd);
     }
 
 protected:
@@ -122,9 +126,13 @@ class nfs_service_impl : public ::dsn::serverlet<nfs_service_impl>
 
     ::dsn::task_ptr _file_close_timer;
 
+    std::unique_ptr<folly::DynamicTokenBucket> _send_token_bucket; // rate limiter of send to remote
+
     perf_counter_wrapper _recent_copy_data_size;
     perf_counter_wrapper _recent_copy_fail_count;
 
+    dsn_handle_t _nfs_max_send_rate_megabytes_cmd;
+
     dsn::task_tracker _tracker;
 };
 } // namespace service