[feature](meta-service) Support querying and adjusting rpc qps limit on meta service

cloud/src/meta-service/meta_service_http.cpp

@@ -22,6 +22,7 @@
 #include <brpc/uri.h>
 #include <fmt/format.h>
 #include <gen_cpp/cloud.pb.h>
+#include <glog/logging.h>
 #include <google/protobuf/message.h>
 #include <google/protobuf/service.h>
 #include <google/protobuf/util/json_util.h>
@@ -30,8 +31,14 @@
 #include <rapidjson/prettywriter.h>
 #include <rapidjson/stringbuffer.h>
 
+#include <algorithm>
+#include <array>
+#include <cstdint>
+#include <functional>
 #include <memory>
 #include <optional>
+#include <string>
+#include <string_view>
 #include <type_traits>
 #include <variant>
 #include <vector>
@@ -42,6 +49,7 @@
 #include "meta-service/txn_kv.h"
 #include "meta-service/txn_kv_error.h"
 #include "meta_service.h"
+#include "rate-limiter/rate_limiter.h"
 
 namespace doris::cloud {
 
@@ -333,6 +341,113 @@ static HttpResponse process_alter_iam(MetaServiceImpl* service, brpc::Controller
     return http_json_reply(resp.status());
 }
 
+static HttpResponse process_adjust_rate_limit(MetaServiceImpl* service, brpc::Controller* cntl) {
+    const auto& uri = cntl->http_request().uri();
+    auto qps_limit_str = std::string {http_query(uri, "qps_limit")};
+    auto rpc_name = std::string {http_query(uri, "rpc_name")};
+    auto instance_id = std::string {http_query(uri, "instance_id")};
+
+    auto process_set_qps_limit = [&](std::function<bool(int64_t)> cb) -> HttpResponse {
+        DCHECK(!qps_limit_str.empty());
+        int64_t qps_limit = -1;
+        try {
+            qps_limit = std::stoll(qps_limit_str);
+        } catch (const std::exception& ex) {
+            return http_json_reply(
+                    MetaServiceCode::INVALID_ARGUMENT,
+                    fmt::format("param `qps_limit` is not a legal int64 type:{}", ex.what()));
+        }
+        if (qps_limit < 0) {
+            return http_json_reply(MetaServiceCode::INVALID_ARGUMENT,
+                                   "`qps_limit` should not be less than 0");
+        }
+        if (cb(qps_limit)) {
+            return http_json_reply(MetaServiceCode::OK, "sucess to adjust rate limit");
+        }
+        return http_json_reply(MetaServiceCode::INVALID_ARGUMENT,
+                               fmt::format("failed to adjust rate limit for qps_limit={}, "
+                                           "rpc_name={}, instance_id={}, plz ensure correct "
+                                           "rpc/instance name",
+                                           qps_limit_str, rpc_name, instance_id));
+    };
+
+    auto set_global_qps_limit = [process_set_qps_limit, service]() {
+        return process_set_qps_limit([service](int64_t qps_limit) {
+            return service->rate_limiter()->set_rate_limit(qps_limit);
+        });
+    };
+
+    auto set_rpc_qps_limit = [&]() {
+        return process_set_qps_limit([&](int64_t qps_limit) {
+            return service->rate_limiter()->set_rate_limit(qps_limit, rpc_name);
+        });
+    };
+
+    auto set_instance_qps_limit = [&]() {
+        return process_set_qps_limit([&](int64_t qps_limit) {
+            return service->rate_limiter()->set_instance_rate_limit(qps_limit, instance_id);
+        });
+    };
+
+    auto set_instance_rpc_qps_limit = [&]() {
+        return process_set_qps_limit([&](int64_t qps_limit) {
+            return service->rate_limiter()->set_rate_limit(qps_limit, rpc_name, instance_id);
+        });
+    };
+
+    auto process_invalid_arguments = [&]() -> HttpResponse {
+        return http_json_reply(MetaServiceCode::INVALID_ARGUMENT,
+                               fmt::format("invalid argument: qps_limit(required)={}, "
+                                           "rpc_name(optional)={}, instance_id(optional)={}",
+                                           qps_limit_str, rpc_name, instance_id));
+    };
+
+    // We have 3 optional params and 2^3 combination, and 4 of them are illegal.
+    // We register callbacks for them in porcessors accordings to the level, represented by 3 bits.
+    std::array<std::function<HttpResponse()>, 8> processors;
+    std::fill_n(processors.begin(), 8, std::move(process_invalid_arguments));
+    processors[0b001] = std::move(set_global_qps_limit);
+    processors[0b011] = std::move(set_rpc_qps_limit);
+    processors[0b101] = std::move(set_instance_qps_limit);
+    processors[0b111] = std::move(set_instance_rpc_qps_limit);
+
+    uint8_t level = (0x01 & qps_limit_str.empty()) | ((0x01 & rpc_name.empty()) << 1) |
+                    ((0x01 & instance_id.empty()) << 2);
+
+    DCHECK_LT(level, 8);
+
+    return processors[level]();
+}
+
+static HttpResponse process_query_rate_limit(MetaServiceImpl* service, brpc::Controller* cntl) {
+    auto rate_limiter = service->rate_limiter();
+    rapidjson::Document d;
+    auto get_qps_limit = [&d](std::string_view rpc_name,
+                              std::shared_ptr<RpcRateLimiter> rpc_limiter) {
+        rapidjson::Document node;
+        rapidjson::Document sub;
+        auto get_qps_token_limit = [&](std::string_view instance_id,
+                                       std::shared_ptr<RpcRateLimiter::QpsToken> qps_token) {
+            sub.AddMember(rapidjson::StringRef(instance_id.data(), instance_id.size()),
+                          qps_token->max_qps_limit(), d.GetAllocator());
+        };
+        rpc_limiter->for_each_qps_token(std::move(get_qps_token_limit));
+
+        auto max_qps_limit = std::to_string(rpc_limiter->max_qps_limit());
+        node.AddMember("RPC qps limit",
+                       rapidjson::StringRef(max_qps_limit.data(), max_qps_limit.size()),
+                       d.GetAllocator());
+        node.AddMember("instance specific qps limit", sub, d.GetAllocator());
+        d.AddMember(rapidjson::StringRef(rpc_name.data(), rpc_name.size()), node, d.GetAllocator());
+    };
+    rate_limiter->for_each_rpc_limiter(std::move(get_qps_limit));
+
+    rapidjson::StringBuffer sb;
+    rapidjson::PrettyWriter<rapidjson::StringBuffer> writer(sb);
+    d.Accept(writer);
+    return http_json_reply(MetaServiceCode::OK, sb.GetString());
+}
+
 static HttpResponse process_decode_key(MetaServiceImpl*, brpc::Controller* ctrl) {
     auto& uri = ctrl->http_request().uri();
     std::string_view key = http_query(uri, "key");
@@ -615,13 +730,17 @@ void MetaServiceImpl::http(::google::protobuf::RpcController* controller,
             {"abort_tablet_job", process_abort_tablet_job},
             {"alter_ram_user", process_alter_ram_user},
             {"alter_iam", process_alter_iam},
+            {"adjust_rate_limit", process_adjust_rate_limit},
+            {"list_rate_limit", process_query_rate_limit},
             {"v1/abort_txn", process_abort_txn},
             {"v1/abort_tablet_job", process_abort_tablet_job},
             {"v1/alter_ram_user", process_alter_ram_user},
             {"v1/alter_iam", process_alter_iam},
+            {"v1/adjust_rate_limit", process_adjust_rate_limit},
+            {"v1/list_rate_limit", process_query_rate_limit},
     };
 
-    auto cntl = static_cast<brpc::Controller*>(controller);
+    auto* cntl = static_cast<brpc::Controller*>(controller);
     brpc::ClosureGuard closure_guard(done);
 
     // Prepare input request info

cloud/src/rate-limiter/rate_limiter.cpp

@@ -17,22 +17,27 @@
 
 #include "rate_limiter.h"
 
+#include <bthread/mutex.h>
 #include <butil/strings/string_split.h>
 
+#include <algorithm>
 #include <chrono>
+#include <cstdint>
 #include <memory>
 #include <mutex>
+#include <ranges>
+#include <shared_mutex>
 
 #include "common/bvars.h"
 #include "common/config.h"
 #include "common/configbase.h"
 
 namespace doris::cloud {
 
-void RateLimiter::init(google::protobuf::Service* service) {
-    std::map<std::string, int64_t> rpc_name_to_max_qps_limit;
+std::unordered_map<std::string, int64_t> parse_specific_qps_limit(const std::string& list_str) {
+    std::unordered_map<std::string, int64_t> rpc_name_to_max_qps_limit;
     std::vector<std::string> max_qps_limit_list;
-    butil::SplitString(config::specific_max_qps_limit, ';', &max_qps_limit_list);
+    butil::SplitString(list_str, ';', &max_qps_limit_list);
     for (const auto& v : max_qps_limit_list) {
         auto p = v.find(':');
         if (p != std::string::npos && p != (v.size() - 1)) {
@@ -41,36 +46,92 @@ void RateLimiter::init(google::protobuf::Service* service) {
                 int64_t max_qps_limit = std::stoll(v.substr(p + 1));
                 if (max_qps_limit > 0) {
                     rpc_name_to_max_qps_limit[rpc_name] = max_qps_limit;
-                    LOG(INFO) << "set rpc: " << rpc_name << " max_qps_limit: " << max_qps_limit;
                 }
             } catch (...) {
-                LOG(WARNING) << "failed to set max_qps_limit to rpc: " << rpc_name
+                LOG(WARNING) << "failed to parse max_qps_limit to rpc: " << rpc_name
                              << " config: " << v;
             }
         }
     }
+    return rpc_name_to_max_qps_limit;
+}
+
+template <typename Callable>
+void for_each_rpc_name(google::protobuf::Service* service, Callable cb) {
     auto method_size = service->GetDescriptor()->method_count();
     for (auto i = 0; i < method_size; ++i) {
         std::string rpc_name = service->GetDescriptor()->method(i)->name();
-        int64_t max_qps_limit = config::default_max_qps_limit;
+        cb(rpc_name);
+    }
+}
 
-        auto it = rpc_name_to_max_qps_limit.find(rpc_name);
-        if (it != rpc_name_to_max_qps_limit.end()) {
+void RateLimiter::init(google::protobuf::Service* service) {
+    auto rpc_name_to_specific_limit = parse_specific_qps_limit(config::specific_max_qps_limit);
+    std::unique_lock write_lock(mutex_);
+    for_each_rpc_name(service, [&](const std::string& rpc_name) {
+        auto it = rpc_name_to_specific_limit.find(rpc_name);
+        int64_t max_qps_limit = config::default_max_qps_limit;
+        if (it != rpc_name_to_specific_limit.end()) {
             max_qps_limit = it->second;
         }
         limiters_[rpc_name] = std::make_shared<RpcRateLimiter>(rpc_name, max_qps_limit);
+    });
+    for (const auto& [k, _] : rpc_name_to_specific_limit) {
+        rpc_with_specific_limit_.insert(k);
     }
 }
 
 std::shared_ptr<RpcRateLimiter> RateLimiter::get_rpc_rate_limiter(const std::string& rpc_name) {
-    // no need to be locked, because it is only modified during initialization
     auto it = limiters_.find(rpc_name);
     if (it == limiters_.end()) {
         return nullptr;
     }
     return it->second;
 }
 
+bool RateLimiter::set_rate_limit(int64_t qps_limit) {
+    std::lock_guard lock(mutex_);
+    auto filter = [this](const auto& kv) { return !rpc_with_specific_limit_.contains(kv.first); };
+    for (const auto& [_, v] : limiters_ | std::views::filter(std::move(filter))) {
+        v->set_max_qps_limit(qps_limit);
+    }
+    return true;
+}
+
+bool RateLimiter::set_rate_limit(int64_t qps_limit, const std::string& rpc_name) {
+    if (!limiters_.contains(rpc_name)) {
+        return false;
+    }
+    auto limiter = limiters_.at(rpc_name);
+    std::lock_guard lock(mutex_);
+    limiter->set_max_qps_limit(qps_limit);
+    rpc_with_specific_limit_.insert(rpc_name);
+    return true;
+}
+
+bool RateLimiter::set_rate_limit(int64_t qps_limit, const std::string& rpc_name,
+                                 const std::string& instance_id) {
+    if (!limiters_.contains(rpc_name)) {
+        return false;
+    }
+    auto limiter = limiters_.at(rpc_name);
+    return limiter->set_max_qps_limit(qps_limit, instance_id);
+}
+
+bool RateLimiter::set_instance_rate_limit(int64_t qps_limit, const std::string& instance_id) {
+    return std::ranges::all_of(limiters_, [&](const auto& kv) {
+        return kv.second->set_max_qps_limit(qps_limit, instance_id);
+    });
+    return true;
+}
+
+void RateLimiter::for_each_rpc_limiter(
+        std::function<void(std::string_view, std::shared_ptr<RpcRateLimiter>)> cb) {
+    for (const auto& [rpc_name, rpc_limiter] : limiters_) {
+        cb(rpc_name, rpc_limiter);
+    }
+}
+
 bool RpcRateLimiter::get_qps_token(const std::string& instance_id,
                                    std::function<int()>& get_bvar_qps) {
     if (!config::use_detailed_metrics || instance_id.empty()) {
@@ -93,6 +154,35 @@ bool RpcRateLimiter::get_qps_token(const std::string& instance_id,
     return qps_token->get_token(get_bvar_qps);
 }
 
+void RpcRateLimiter::set_max_qps_limit(int64_t max_qps_limit) {
+    std::lock_guard<bthread::Mutex> l(mutex_);
+    max_qps_limit_ = max_qps_limit;
+    auto filter = [this](const auto& kv) {
+        return !instance_with_specific_limit_.contains(kv.first);
+    };
+    for (auto& [k, v] : qps_limiter_ | std::views::filter(std::move(filter))) {
+        v->set_max_qps_limit(max_qps_limit);
+    }
+}
+
+bool RpcRateLimiter::set_max_qps_limit(int64_t max_qps_limit, const std::string& instance_id) {
+    std::lock_guard<bthread::Mutex> l(mutex_);
+    if (!qps_limiter_.contains(instance_id)) {
+        qps_limiter_[instance_id] = std::make_shared<QpsToken>(max_qps_limit);
+    } else {
+        qps_limiter_.at(instance_id)->set_max_qps_limit(max_qps_limit);
+    }
+    instance_with_specific_limit_.insert(instance_id);
+    return true;
+}
+
+void RpcRateLimiter::for_each_qps_token(
+        std::function<void(std::string_view, std::shared_ptr<QpsToken>)> cb) {
+    for (const auto& [instance_id, qps_token] : qps_limiter_) {
+        cb(instance_id, qps_token);
+    }
+}
+
 bool RpcRateLimiter::QpsToken::get_token(std::function<int()>& get_bvar_qps) {
     using namespace std::chrono;
     auto now = steady_clock::now();
@@ -110,4 +200,9 @@ bool RpcRateLimiter::QpsToken::get_token(std::function<int()>& get_bvar_qps) {
     return current_qps_ < max_qps_limit_;
 }
 
-} // namespace doris::cloud
+void RpcRateLimiter::QpsToken::set_max_qps_limit(int64_t max_qps_limit) {
+    std::lock_guard<bthread::Mutex> l(mutex_);
+    max_qps_limit_ = max_qps_limit;
+}
+
+} // namespace doris::cloud